"""Exercise background delegation through the real SDK runner and session history."""

import asyncio
import json
from copy import deepcopy

import pytest
import pytest_asyncio

from bumblehive import BumblehiveRuntime
from bumblehive.providers.base import ModelProvider, ModelResponse
from bumblehive.protocols import ToolCall
from bumblehive.session.manager import SessionManager
from bumblehive_server.runtime_service import RuntimeBusyError, RuntimeService
from bumblehive_server.schemas import ChatRequest
from bumblehive_server.session_reader import SessionReader


class DelegationProvider(ModelProvider):
    def __init__(self):
        self.plans = {}
        self.children = {}
        self.parent_release = asyncio.Event()
        self.parent_release.set()
        self.requests = []
        self.child_error = False
        self.parent_error = False
        self.approval_path = None

    async def generate(self, request):
        messages = request.messages
        self.requests.append(deepcopy(messages))
        users = [m["content"] for m in messages if m["role"] == "user"]
        first = users[0]
        if "\nTask:\nCHILD:" in first:
            name = first.rsplit("CHILD:", 1)[1].splitlines()[0]
            if messages[-1]["role"] == "tool":
                return ModelResponse(
                    content=f"REPORT:{name}: {messages[-1]['content']}"
                )
            gate = self.children.setdefault(name, asyncio.Event())
            await gate.wait()
            if self.child_error:
                raise RuntimeError("child execution failed")
            return ModelResponse(
                None,
                [ToolCall("read-child", "read_file", {"path": "data.txt"})],
                "tool_calls",
            )
        root = first.removeprefix("ROOT:").splitlines()[0]
        if len(messages) == len([m for m in messages if m["role"] == "system"]) + 1:
            count, approval = self.plans[root]
            calls = [
                ToolCall(
                    f"delegate-{i}",
                    "sub_agent",
                    {
                        "title": f"Review {root}-{i}",
                        "task": f"CHILD:{root}-{i}",
                    },
                )
                for i in range(count)
            ]
            if approval:
                calls.append(
                    ToolCall(
                        "write-parent",
                        "write_file",
                        {"path": str(self.approval_path), "content": "approved"},
                    )
                )
            return ModelResponse(None, calls, "tool_calls")
        await self.parent_release.wait()
        if self.parent_error:
            raise RuntimeError("parent execution failed")
        reports = [text for text in users if text.startswith("[Background tool result]")]
        return ModelResponse(
            content=f"PARENT:{root}:"
            + ("\n".join(reports) if reports else "initial answer")
        )

    async def generate_stream(self, request, *, callbacks=None):
        response = await self.generate(request)
        if response.content and callbacks and callbacks.on_content_delta:
            await callbacks.on_content_delta(response.content)
        return response


@pytest_asyncio.fixture
async def delegation(tmp_path):
    reader = SessionReader(tmp_path / "sessions")
    provider = DelegationProvider()
    provider.approval_path = tmp_path / "approval.txt"
    runtimes = []

    def factory(config):
        runtime = BumblehiveRuntime(config)
        runtime.sessions = SessionManager(reader.directory)

        async def get_provider(_config):
            return provider

        runtime._get_provider = get_provider
        runtimes.append(runtime)
        return runtime

    service = RuntimeService(
        tmp_path / "config.json", session_reader=reader, runtime_factory=factory
    )
    await service.startup()
    try:
        yield service, provider, reader, tmp_path
    finally:
        await service.shutdown()
    assert service._active_runs == 0
    assert service.subagents._tasks == {}


async def receive(queue, kind, session_id=None, status=None):
    async with asyncio.timeout(4):
        while True:
            value = await queue.get()
            if (
                value["type"] == kind
                and (session_id is None or value.get("session_id") == session_id)
                and (status is None or value.get("status") == status)
            ):
                return value


async def begin(delegation, root, count, *, approval=False, workspace=None):
    service, provider, reader, tmp_path = delegation
    workspace = workspace or tmp_path / root
    workspace.mkdir(exist_ok=True)
    (workspace / "data.txt").write_text(f"DATA:{root}")
    session_id = await reader.create(workspace)
    controller = service.get_session_controller(session_id)
    queue, _ = controller.subscribe()
    provider.plans[root] = (count, approval)
    await controller.submit(
        ChatRequest(
            content=f"ROOT:{root}",
            config={"runtime": {"workspace": str(workspace)}},
        )
    )
    children = [await receive(queue, "session_created") for _ in range(count)]
    return session_id, controller, queue, children


def assert_message_order(messages):
    waiting = set()
    for message in messages:
        role = message["role"]
        if role == "assistant":
            assert not waiting, "assistant interrupted a tool result batch"
            waiting = {call["id"] for call in message.get("tool_calls", [])}
        elif role == "tool":
            assert message["tool_call_id"] in waiting
            waiting.remove(message["tool_call_id"])
        elif role == "user":
            assert not waiting, "notification split an assistant/tool result batch"
    assert not waiting


@pytest.mark.asyncio
@pytest.mark.parametrize("count", [1, 3])
async def test_background_returns_then_wakes_idle_parent_and_preserves_order(
    delegation, count
):
    service, provider, reader, _ = delegation
    parent, controller, queue, children = await begin(
        delegation, "idle", count
    )
    assert all("background" not in child for child in children)
    await receive(queue, "result", parent)
    history = (await reader.get(parent)).messages
    assert [json.loads(m["content"]) for m in history if m["role"] == "tool"] == [
        {"status": "pending", "task_id": f"delegate-{i}"} for i in range(count)
    ]
    assert service._active_runs == count
    with pytest.raises(RuntimeBusyError):
        await service.update_config({"provider": {"model": "other"}})
    assert len(service.subagents.snapshot(parent)) == count
    for i in range(count):
        provider.children.setdefault(f"idle-{i}", asyncio.Event()).set()
    unfinished = {child["session_id"] for child in children}
    while unfinished:
        ended = await receive(queue, "result")
        unfinished.discard(ended["session_id"])
    async with asyncio.timeout(4):
        while service.subagents._tasks or controller.active_task:
            await asyncio.sleep(0.001)
    assert service.subagents.snapshot(parent) == []
    history = (await reader.get(parent)).messages
    reports = [
        m
        for m in history
        if str(m.get("content", "")).startswith("[Background tool result]")
    ]
    assert len(reports) == count
    assert {f"delegate-{i}" for i in range(count)} == {
        m["content"].split("Task ID: ", 1)[1].split("\n", 1)[0] for m in reports
    }
    assert all(m["role"] == "user" for m in reports)
    assert all(m["content"].split("\n\n", 1)[1].startswith("REPORT:idle-") for m in reports)
    for child in children:
        assert (await reader.get(child["session_id"])).messages
    assert_message_order(history)
    for request in provider.requests:
        assert_message_order(request)


@pytest.mark.asyncio
async def test_two_parent_sessions_keep_results_and_workspaces_separate(delegation):
    service, provider, reader, _ = delegation
    parents = await asyncio.gather(
        begin(delegation, "one", 2), begin(delegation, "two", 1)
    )
    for parent, _, queue, _ in parents:
        await receive(queue, "result", parent)
    for name in ["two-0", "one-1", "one-0"]:
        provider.children.setdefault(name, asyncio.Event()).set()
    async with asyncio.timeout(4):
        while service.subagents._tasks or any(c.active_task for _, c, _, _ in parents):
            await asyncio.sleep(0.001)
    for root, (parent, _, _, children) in zip(["one", "two"], parents):
        history = (await reader.get(parent)).messages
        other = "two" if root == "one" else "one"
        assert f"REPORT:{root}-" in str(history)
        assert f"REPORT:{other}-" not in str(history)
        assert f"DATA:{root}" in str(history)
        assert f"DATA:{other}" not in str(history)
        assert_message_order(history)


@pytest.mark.asyncio
async def test_completion_during_approval_is_steered_then_restored_on_stop(delegation):
    service, provider, reader, tmp_path = delegation
    parent, controller, queue, children = await begin(
        delegation, "stop", 1, approval=True
    )
    await receive(queue, "approval_request", parent)
    provider.children.setdefault("stop-0", asyncio.Event()).set()
    await receive(queue, "message_status", parent, "steered")
    await controller.cancel()
    assert len(controller.pending) == 1
    report_id = controller.pending[0].request.message_id
    assert report_id == "subagent:delegate-0"
    _, snapshot = controller.subscribe()
    assert snapshot["pending"][0]["message_id"] == report_id
    assert not snapshot["running"]
    assert not (tmp_path / "approval.txt").exists()
    await controller.steer(report_id)
    await receive(queue, "result", parent)
    history = (await reader.get(parent)).messages
    assert (
        sum(
            str(m.get("content", "")).startswith("[Background tool result]")
            for m in history
        )
        == 1
    )
    assert_message_order(history)


@pytest.mark.asyncio
async def test_late_report_after_stop_stays_queued_until_manual_resume(delegation):
    service, provider, reader, _ = delegation
    provider.parent_release.clear()
    parent, controller, queue, children = await begin(delegation, "late", 1)
    while (await receive(queue, "event"))["kind"] != "tool.call.finished":
        pass
    await controller.cancel()
    provider.children.setdefault("late-0", asyncio.Event()).set()
    await receive(queue, "message_status", parent, "queued")
    async with asyncio.timeout(4):
        while service.subagents._tasks:
            await asyncio.sleep(0.001)
    assert controller.active_task is None
    assert len(controller.pending) == 1
    assert controller.pending[0].request.config["runtime"]["workspace"].endswith("late")
    provider.parent_release.set()
    await controller.steer(controller.pending[0].request.message_id)
    await receive(queue, "result", parent)
    assert "REPORT:late-0" in str((await reader.get(parent)).messages)


@pytest.mark.asyncio
async def test_disconnect_does_not_cancel_background_and_snapshot_is_authoritative(
    delegation,
):
    service, provider, reader, _ = delegation
    parent, controller, queue, children = await begin(delegation, "reconnect", 1)
    await receive(queue, "result", parent)
    controller.unsubscribe(queue)
    again, snapshot = controller.subscribe()
    assert not snapshot["running"]
    assert (
        service.subagents.snapshot(parent)[0]["session_id"] == children[0]["session_id"]
    )
    controller.unsubscribe(again)
    provider.children.setdefault("reconnect-0", asyncio.Event()).set()
    async with asyncio.timeout(4):
        while service.subagents._tasks or controller.active_task:
            await asyncio.sleep(0.001)
    _, final = controller.subscribe()
    assert not final["running"]
    assert service.subagents.snapshot(parent) == []
    assert "REPORT:reconnect-0" in str((await reader.get(parent)).messages)


@pytest.mark.asyncio
async def test_failure_reports_once_and_delete_cancels_without_waking_parent(
    delegation,
):
    service, provider, reader, _ = delegation
    provider.child_error = True
    parent, controller, queue, children = await begin(delegation, "error", 1)
    await receive(queue, "result", parent)
    provider.children.setdefault("error-0", asyncio.Event()).set()
    await receive(queue, "message_status", parent, "started")
    await receive(queue, "result", parent)
    reports = [
        m
        for m in (await reader.get(parent)).messages
        if str(m.get("content", "")).startswith("[Background tool result]")
    ]
    assert len(reports) == 1
    assert reports[0]["content"] == (
        "[Background tool result]\nTask ID: delegate-0\n\n"
        "Error:\nchild execution failed"
    )
    provider.child_error = False
    parent2, controller2, queue2, children2 = await begin(delegation, "delete", 1)
    await receive(queue2, "result", parent2)
    await service.delete_session(children2[0]["session_id"])
    assert controller2.active_task is None and not controller2.pending
    assert not service.subagents.snapshot(parent2)


@pytest.mark.asyncio
async def test_report_arriving_during_final_model_response_follows_assistant(
    delegation,
):
    service, provider, reader, _ = delegation
    provider.parent_release.clear()
    parent, controller, queue, _ = await begin(delegation, "final", 1)
    # Wait until the parent has requested its answer with the delegation result.
    async with asyncio.timeout(4):
        while not any(
            messages[-1]["role"] == "tool"
            and str(messages[1].get("content", "")).startswith("ROOT:final")
            for messages in provider.requests
        ):
            await asyncio.sleep(0.001)
    provider.children.setdefault("final-0", asyncio.Event()).set()
    await receive(queue, "message_status", parent, "steered")
    provider.parent_release.set()
    await receive(queue, "result", parent)
    history = (await reader.get(parent)).messages
    index = next(
        i
        for i, m in enumerate(history)
        if str(m.get("content", "")).startswith("[Background tool result]")
    )
    assert history[index - 1]["role"] == "assistant"
    assert not history[index - 1].get("tool_calls")
    assert_message_order(history)
    await controller.cancel()
    assert not controller.pending, "a consumed report must not be restored"


@pytest.mark.asyncio
async def test_report_during_approval_is_consumed_after_complete_tool_batch(delegation):
    _, provider, reader, _ = delegation
    parent, controller, queue, _ = await begin(
        delegation, "approve", 1, approval=True
    )
    approval = await receive(queue, "approval_request", parent)
    provider.children.setdefault("approve-0", asyncio.Event()).set()
    await receive(queue, "message_status", parent, "steered")
    assert not controller.pending
    await controller.decide_approval(approval["approval_id"], False)
    await receive(queue, "result", parent)
    history = (await reader.get(parent)).messages
    index = next(
        i
        for i, m in enumerate(history)
        if str(m.get("content", "")).startswith("[Background tool result]")
    )
    assert history[index - 1]["role"] == "tool"
    assert_message_order(history)


@pytest.mark.asyncio
async def test_extremely_fast_children_still_report_after_the_tool_batch(delegation):
    service, provider, reader, _ = delegation
    for i in range(3):
        provider.children.setdefault(f"fast-{i}", asyncio.Event()).set()
    parent, controller, queue, _ = await begin(delegation, "fast", 3)
    async with asyncio.timeout(4):
        while service.subagents._tasks or controller.active_task:
            await asyncio.sleep(0.001)
    history = (await reader.get(parent)).messages
    assert (
        len(
            notices := [
                m
                for m in history
                if str(m.get("content", "")).startswith("[Background tool result]")
            ]
        )
        == 3
    )
    assert all("REPORT:fast-" in m["content"] for m in notices)
    assert_message_order(history)


@pytest.mark.asyncio
async def test_background_cancel_before_lease_does_not_leave_startup_waiting():
    from contextlib import asynccontextmanager
    from bumblehive_server.subagents import SubagentTasks
    from test_runtime_service import RecordingSubagentObserver

    @asynccontextmanager
    async def blocked_lease():
        await asyncio.Event().wait()
        yield None

    reports = []

    async def notify(*args):
        reports.append(args)

    observer = RecordingSubagentObserver()
    tasks = SubagentTasks(blocked_lease, notify)
    startup = asyncio.create_task(
        tasks.start(
            task_id="delegate-blocked",
            session_id="child",
            parent_session_id="parent",
            title="Waiting",
            workspace="/tmp",
            content="Inspect",
            observer=observer,
        )
    )
    await asyncio.sleep(0)  # Cancel before the worker can acquire its lease.
    await tasks.cancel("child")
    with pytest.raises(RuntimeError, match="stopped before startup"):
        await asyncio.wait_for(startup, 1)
    assert observer.cancelled == ["child"]
    assert reports == [] and tasks._tasks == {}


@pytest.mark.asyncio
async def test_lease_startup_failure_is_one_tool_error_without_completion_notice():
    from contextlib import asynccontextmanager
    from bumblehive_server.subagents import SubagentTasks
    from test_runtime_service import RecordingSubagentObserver

    @asynccontextmanager
    async def failed_lease():
        raise RuntimeError("runtime unavailable")
        yield None

    reports = []

    async def notify(*args):
        reports.append(args)

    observer = RecordingSubagentObserver()
    tasks = SubagentTasks(failed_lease, notify)
    with pytest.raises(RuntimeError, match="runtime unavailable"):
        await tasks.start(
            task_id="delegate-failed",
            session_id="child",
            parent_session_id="parent",
            title="Start",
            workspace="/tmp",
            content="Inspect",
            observer=observer,
        )
    assert observer.errors == [
        {"session_id": "child", "message": "runtime unavailable"}
    ]
    assert observer.cancelled == [] and reports == [] and tasks._tasks == {}


@pytest.mark.asyncio
async def test_parent_failure_does_not_auto_resume_for_late_report(delegation):
    service, provider, reader, _ = delegation
    provider.parent_error = True
    parent, controller, queue, _ = await begin(delegation, "parent-error", 1)
    await receive(queue, "error", parent)
    provider.children.setdefault("parent-error-0", asyncio.Event()).set()
    await receive(queue, "message_status", parent, "queued")
    async with asyncio.timeout(4):
        while service.subagents._tasks:
            await asyncio.sleep(0.001)
    assert controller.active_task is None and controller._paused
    assert len(controller.pending) == 1
    provider.parent_error = False
    await controller.steer(controller.pending[0].request.message_id)
    await receive(queue, "result", parent)
    assert "REPORT:parent-error-0" in str((await reader.get(parent)).messages)