"""Agent families: spawn, forwarding, addressing, snapshots, and routing."""

import json
import tempfile
import threading
import time
import unittest
import urllib.error
import urllib.request
from pathlib import Path

from harness import Albedo, Provider, exclusive, text


def user_message(request):
    return next(
        (
            item.get("content", "")
            for item in reversed(request["messages"])
            if item.get("role") == "user"
        ),
        "",
    )


def wait_for(predicate, timeout=40):
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        value = predicate()
        if value:
            return value
        time.sleep(0.05)
    raise AssertionError("agent event never arrived")


class AgentsTests(unittest.TestCase):
    def setUp(self):
        def answer(request):
            user = user_message(request)
            line = user.splitlines()[-2 if "</mail>" in user else -1]
            return text("answered: " + line)

        self.provider = Provider(answer)
        self.alpha = f"alpha-{self.provider.route}"
        self.beta = f"beta-{self.provider.route}"
        providers = {
            profile: {
                "baseUrl": self.provider.url + f"/{name}/v1",
                "apiKey": "key",
                "model": f"fixture-{name}-{self.provider.route}",
                "protocol": "chat_completions",
            }
            for profile, name in ((self.alpha, "alpha"), (self.beta, "beta"))
        }

        # The cached catalog lists a model for the fixture endpoint that no profile names.
        # models.json is daemon-wide, so only the exclusive catalog test writes it.
        def prepare(app):
            (app.home / "models.json").write_text(
                json.dumps(
                    {
                        "fixture-gateway": {
                            "api": "http://127.0.0.1/v1",
                            "env": [],
                            "models": {"fixture-listed": {"id": "fixture-listed"}},
                        }
                    }
                )
            )

        catalog = (
            self._testMethodName
            == "test_generic_profile_lists_its_endpoints_catalog_models"
        )
        self.app = Albedo(
            self.provider, providers=providers, prepare=prepare if catalog else None
        )
        self.app.__enter__()
        self.addCleanup(self.provider.close)
        self.addCleanup(self.app.__exit__, None, None, None)

    def api(self, path, body=None, method=None):
        with self.app.api(path, body, method=method) as response:
            return json.load(response)

    def asked(self, fragment):
        return [
            user_message(record["request"])
            for record in self.provider.requests
            if fragment in user_message(record["request"])
        ]

    def spawn(self, parent, name="coder", task="map every wake path", model=None):
        body = {"name": name, "task": task}
        if model is not None:
            body["model"] = model
        return self.api(f"/sessions/{parent}/children", body)

    def listen(self):
        heard = []
        ready = threading.Event()

        def consume():
            request = urllib.request.Request(
                self.app.base + "/agents/stream",
                headers={"Authorization": "Bearer " + self.app.connection["token"]},
            )
            with urllib.request.urlopen(request, timeout=30) as response:
                for raw in response:
                    if raw.startswith(b"data: "):
                        heard.extend(json.loads(raw[6:])["events"])
                        ready.set()

        threading.Thread(target=consume, daemon=True).start()
        self.assertTrue(
            ready.wait(timeout=2), "agent stream did not send its initial frame"
        )
        return heard

    def test_spawn_task_and_unreviewed_forward(self):
        parent = self.app.session()
        made = self.spawn(parent)
        child = made["session"]["id"]
        self.assertEqual(
            made["member"],
            {
                "session": child,
                "parent": parent,
                "name": "coder",
                "depth": 1,
                "closed": False,
            },
        )
        self.assertEqual(made["session"]["title"], "coder")
        self.assertNotIn(child, [s["id"] for s in self.api("/sessions")])
        task = wait_for(lambda: self.asked('kind="task"'))[0]
        self.assertIn(f'session="{parent}"', task)
        self.assertIn("map every wake path", task)
        forwarded = wait_for(lambda: self.asked('kind="unreviewed"'))[0]
        self.assertIn('from="coder"', forwarded)
        self.assertIn("answered: map every wake path", forwarded)
        self.app.idle(parent)
        self.app.idle(child)
        self.assertEqual(
            [m["name"] for m in self.api(f"/sessions/{parent}/children")], ["coder"]
        )

    @exclusive
    def test_tree_snapshot_and_event_stream(self):
        parent = self.app.session()
        heard = self.listen()
        child = self.spawn(parent)["session"]["id"]
        wait_for(lambda: self.asked('kind="unreviewed"'))
        for asker in (parent, child):
            tree = self.api(f"/agents?session={asker}")
            self.assertEqual(tree["root"], parent)
            self.assertEqual(
                [(n["parent"], n["depth"]) for n in tree["nodes"]],
                [(None, 0), (parent, 1)],
            )
            self.assertEqual(tree["nodes"][1]["name"], "coder")
        expected = {
            ("spawn", child),
            ("mail", child),
            ("mail", parent),
            ("running", child),
            ("running", parent),
            ("text", child),
        }

        def kinds():
            return {(e["type"], e.get("session") or e.get("to")) for e in list(heard)}

        try:
            wait_for(lambda: expected <= kinds(), timeout=10)
        except AssertionError:
            self.fail(f"missing stream events: {expected - kinds()}; got {kinds()}")

    @exclusive
    def test_rename_keeps_family_address_and_reports_stream_event(self):
        parent = self.app.session()
        child = self.spawn(parent)["session"]["id"]
        self.app.idle(child)
        wait_for(lambda: self.asked('kind="unreviewed"'))
        self.app.idle(parent)
        heard = self.listen()
        renamed = self.api(
            f"/sessions/{child}", {"name": "  wake path\naudit "}, "PATCH"
        )
        self.assertEqual(renamed["title"], "wake path audit")
        self.assertEqual(
            self.api(f"/sessions/{parent}", {"name": "orchestration"}, "PATCH")[
                "title"
            ],
            "orchestration",
        )
        self.assertEqual(
            [s["title"] for s in self.api("/sessions") if s["id"] == parent],
            ["orchestration"],
        )
        nodes = {
            n["session"]["id"]: n
            for n in self.api(f"/agents?session={parent}")["nodes"]
        }
        self.assertEqual(
            (nodes[parent]["name"], nodes[parent]["address"]), ("orchestration", None)
        )
        self.assertEqual(
            (nodes[child]["name"], nodes[child]["address"]),
            ("wake path audit", "coder"),
        )
        try:
            wait_for(
                lambda: any(
                    e["type"] == "renamed"
                    and e["session"] == child
                    and e["name"] == "wake path audit"
                    for e in list(heard)
                ),
                timeout=10,
            )
        except AssertionError:
            self.fail(f"rename stream missing: {heard!r}")
        with self.assertRaises(urllib.error.HTTPError) as caught:
            self.api("/sessions/no-such-session", {"name": "x"}, "PATCH")
        self.assertEqual(caught.exception.code, 409)
        receipt = self.api(
            f"/sessions/{parent}/mail", {"to": "coder", "body": "also cover schedules"}
        )
        self.assertEqual((receipt["to"], receipt["name"]), (child, "coder"))
        self.assertIn(receipt["status"], ("delivered", "queued"))
        wait_for(lambda: self.asked("also cover schedules"))
        wait_for(lambda: len(self.asked('kind="unreviewed"')) == 2)
        self.app.idle(parent)

    def test_reply_and_external_id_mail(self):
        parent = self.app.session()
        radio = self.app.session()
        child = self.spawn(parent)["session"]["id"]
        wait_for(lambda: self.asked('kind="unreviewed"'))
        self.app.idle(parent)
        before = len(self.provider.requests)
        receipt = self.api(
            f"/sessions/{child}/mail", {"to": "parent", "body": "schedules covered"}
        )
        self.assertEqual(receipt["to"], parent)
        wait_for(lambda: self.asked("schedules covered"))
        self.app.idle(parent)
        time.sleep(1)
        self.assertEqual(len(self.provider.requests), before + 1)
        with self.assertRaises(urllib.error.HTTPError):
            self.api(f"/sessions/{child}/mail", {"to": "radio", "body": "by name"})
        receipt = self.api(
            f"/sessions/{child}/mail", {"to": radio, "body": "deploy when green"}
        )
        self.assertEqual(receipt["status"], "pending")
        wait_for(lambda: self.asked("deploy when green"), timeout=40)
        self.app.idle(radio)

    def test_tree_deletion_and_renamed_workspace(self):
        parent = self.app.session()
        child = self.spawn(parent)["session"]["id"]
        wait_for(lambda: self.asked('kind="unreviewed"'))
        self.app.idle(parent)
        with self.assertRaises(urllib.error.HTTPError) as caught:
            self.api(f"/sessions/{parent}", method="DELETE")
        self.assertEqual(caught.exception.code, 409)
        self.assertIn("child", caught.exception.read().decode())
        self.assertEqual(
            self.api(f"/sessions/{parent}?tree=1", method="DELETE")["deleted"], 2
        )
        listed = [s["id"] for s in self.api("/sessions")]
        self.assertNotIn(parent, listed)
        self.assertNotIn(child, listed)
        radio = self.app.session()
        self.api(f"/sessions/{radio}", {"name": "radio desk"}, "PATCH")
        moved = self.app.root / "moved"
        moved.mkdir()
        self.assertEqual(
            self.api(f"/sessions/{radio}/workspace", {"workspace": str(moved)})[
                "title"
            ],
            "radio desk",
        )
        self.assertEqual(
            [s["title"] for s in self.api("/sessions") if s["id"] == radio],
            ["radio desk"],
        )

    def test_children_follow_a_workspace_move(self):
        release = threading.Event()
        answer = self.provider.script

        def held(request):
            if "hold this turn" in user_message(request):
                release.wait(30)
            return answer(request)

        self.provider.script = held
        parent = self.app.session()
        still = self.spawn(parent, "still")["session"]["id"]
        away = self.spawn(parent, "away")["session"]["id"]

        # Both answers can reach the parent in one request.
        def heard():
            return " ".join(
                event.get("text", "") for event in self.app.history(parent)["events"]
            )

        wait_for(lambda: 'from="still"' in heard() and 'from="away"' in heard())
        for session in (still, away, parent):
            self.app.idle(session)
        elsewhere = Path(tempfile.mkdtemp(prefix="elsewhere-", dir=self.app.root))
        self.api(f"/sessions/{away}/workspace", {"workspace": str(elsewhere)})
        busy = self.spawn(parent, "busy", "hold this turn")["session"]["id"]
        wait_for(lambda: self.asked("hold this turn"))
        moved = Path(tempfile.mkdtemp(prefix="moved-", dir=self.app.root))
        self.assertEqual(
            self.api(f"/sessions/{parent}/workspace", {"workspace": str(moved)})[
                "workspace"
            ],
            str(moved),
        )

        def workspaces():
            return {
                node["session"]["id"]: node["session"]["workspace"]
                for node in self.api(f"/agents?session={parent}")["nodes"]
            }

        self.assertEqual(
            workspaces(),
            {
                parent: str(moved),
                still: str(moved),
                away: str(elsewhere),
                busy: str(self.app.workspace),
            },
        )
        cursor = self.app.stream_page(busy)
        release.set()
        self.app.idle(busy)
        self.assertEqual(workspaces()[busy], str(moved))
        self.assertIn(
            "workspace moved with parent; python variables were cleared, "
            "the transcript is intact",
            [
                event.get("text")
                for event in self.app.stream_page(busy, cursor)["events"]
            ],
        )

    def test_qualified_and_inferred_cross_provider_models(self):
        parent = self.app.session()
        self.api(f"/sessions/{parent}", {"name": "radio desk"}, "PATCH")
        with self.assertRaises(urllib.error.HTTPError) as caught:
            self.spawn(
                parent, "missing", "should not run", f"{self.beta}/no-such-model"
            )
        self.assertEqual(caught.exception.code, 409)
        for name, model in (
            ("explicit", f"{self.beta}/fixture-beta-{self.provider.route}"),
            ("inferred", f"fixture-beta-{self.provider.route}"),
        ):
            made = self.spawn(parent, name, "cross-provider task", model)
            self.assertEqual(made["session"]["provider"], self.beta)
            self.assertEqual(
                made["session"]["model"], f"fixture-beta-{self.provider.route}"
            )
            self.assertEqual(made["member"]["parent"], parent)
            self.app.idle(made["session"]["id"])
        tasks = wait_for(
            lambda: (
                self.asked('kind="task">\ncross-provider')
                if len(self.asked('kind="task">\ncross-provider')) == 2
                else None
            )
        )
        self.assertTrue(all('from="radio desk"' in task for task in tasks))
        self.assertGreaterEqual(
            sum("/beta/" in r["path"] for r in self.provider.requests), 2
        )

    @exclusive
    def test_generic_profile_lists_its_endpoints_catalog_models(self):
        parent = self.app.session()
        made = self.spawn(
            parent, "listed", "catalog task", f"{self.beta}/fixture-listed"
        )
        self.assertEqual(
            (made["session"]["provider"], made["session"]["model"]),
            (self.beta, "fixture-listed"),
        )
        self.app.idle(made["session"]["id"])