import asyncio
from contextlib import asynccontextmanager
from types import SimpleNamespace

import pytest

from bumblehive.observability.events import make_event
from bumblehive.agent.control import _RunControl
from bumblehive_server.chat.session import InputError, SessionController
from bumblehive_server.schemas import ChatRequest


class Stream:
    def __init__(self, content, session_id):
        self.content = content
        self.session_id = session_id
        self.release = asyncio.Event()
        self.inputs = []
        self.control = _RunControl()
        self.closed = False
        self.error = None
        self.failure = None

    def steer(self, text):
        if self.closed:
            raise RuntimeError("closed")
        self.inputs.append(text)
        self.control.steer(text)

    def consume_inputs(self):
        return self.control._drain_all()

    def take_pending(self):
        return self.control.take_pending()

    def __aiter__(self):
        return self.iterate()

    async def iterate(self):
        yield make_event(
            "model.stream.content_delta",
            run_id="run",
            session_id=self.session_id,
            delta=self.content,
        )
        await self.release.wait()
        if self.failure:
            raise self.failure
        if self.error is None:
            self.consume_inputs()
        self.closed = True

    async def result(self):
        return SimpleNamespace(
            final_content=self.content,
            tools_used=[],
            usage={},
            stop_reason="model_error" if self.error else "completed",
            error=self.error,
        )

    async def aclose(self):
        self.closed = True
        self.control._close()


class Service:
    def __init__(self):
        self.streams = []
        self.lease_gate = asyncio.Event()
        self.lease_gate.set()

    @asynccontextmanager
    async def lease(self):
        await self.lease_gate.wait()
        yield self

    def stream(self, content, *, session_id, **kwargs):
        stream = Stream(content, session_id)
        self.streams.append(stream)
        return stream


def message(text):
    return ChatRequest(message_id=text, content=text)


async def frame(queue, kind, *, status=None, message_id=None):
    async with asyncio.timeout(2):
        while True:
            item = await queue.get()
            if (
                item["type"] == kind
                and (status is None or item.get("status") == status)
                and (message_id is None or item.get("message_id") == message_id)
            ):
                return item


@pytest.mark.asyncio
async def test_followups_are_fifo_and_result_precedes_next_started():
    service = Service()
    controller = SessionController(service.lease, "one")
    queue, _ = controller.subscribe()
    try:
        await controller.submit(message("A"))
        await frame(queue, "message_status", status="started")
        await controller.submit(message("B"))
        await controller.submit(message("C"))
        assert [s.content for s in service.streams] == ["A"]
        for current, following in [("A", "B"), ("B", "C")]:
            service.streams[-1].release.set()
            assert (await frame(queue, "result"))["final_content"] == current
            await frame(queue, "message_status", status="started", message_id=following)
            assert service.streams[-1].content == following
        last_task = controller.active_task
        service.streams[-1].release.set()
        await frame(queue, "result")
        await last_task
        assert controller.active_task is None
        assert [s.content for s in service.streams] == ["A", "B", "C"]
    finally:
        await controller.close()


@pytest.mark.asyncio
async def test_multiple_steers_non_head_delete_and_direct_input():
    service = Service()
    controller = SessionController(service.lease, "one")
    queue, _ = controller.subscribe()
    try:
        await controller.submit(message("A"))
        await frame(queue, "message_status", status="started")
        for text in ("B", "C", "D", "E"):
            await controller.submit(message(text))
        await controller.steer("C")
        await controller.steer("B")
        await controller.submit(message("backend input"), steer=True)
        await controller.remove_pending("D")
        assert service.streams[0].inputs == ["C", "B", "backend input"]
        assert [item.request.content for item in controller.pending] == ["E"]
        for method in (controller.steer, controller.remove_pending):
            with pytest.raises(InputError, match="no longer pending"):
                await method("C")
        service.streams[0].release.set()
        await frame(queue, "message_status", status="started", message_id="E")
        assert [s.content for s in service.streams] == ["A", "E"]
    finally:
        await controller.close()


@pytest.mark.asyncio
async def test_steer_before_stream_and_after_input_closed():
    service = Service()
    service.lease_gate.clear()
    controller = SessionController(service.lease, "one")
    queue, _ = controller.subscribe()
    try:
        await controller.submit(message("A"))
        await controller.submit(message("B"))
        await controller.steer("B")
        service.lease_gate.set()
        await frame(queue, "message_status", status="steered", message_id="B")
        assert service.streams[0].inputs == ["B"]
        service.streams[0].closed = True
        await controller.submit(message("C"))
        await controller.steer("C")
        assert [item.request.content for item in controller.pending] == ["C"]
        service.streams[0].release.set()
        await frame(queue, "message_status", status="started", message_id="C")
    finally:
        await controller.close()


@pytest.mark.asyncio
@pytest.mark.parametrize("before_start", [True, False])
async def test_cancel_retains_queue_and_manual_send_resumes(before_start):
    service = Service()
    controller = SessionController(service.lease, "one")
    queue, _ = controller.subscribe()
    try:
        await controller.submit(message("A"))
        if not before_start:
            await frame(queue, "message_status", status="started")
        await controller.submit(message("B"))
        await controller.cancel()
        assert controller.active_task is None
        assert not service.streams or service.streams[0].closed
        await controller.steer("B")
        await frame(queue, "message_status", status="started", message_id="B")
        assert service.streams[-1].content == "B"
    finally:
        await controller.close()


@pytest.mark.asyncio
async def test_two_sessions_reconnect_multiple_subscribers_and_steers_are_isolated():
    service = Service()
    one = SessionController(service.lease, "one")
    two = SessionController(service.lease, "two")
    first_queue, _ = one.subscribe()
    second_queue, _ = two.subscribe()
    try:
        await asyncio.gather(one.submit(message("A")), two.submit(message("X")))
        await frame(first_queue, "message_status", status="started")
        await frame(second_queue, "message_status", status="started")
        await asyncio.gather(one.submit(message("B")), two.submit(message("Y")))
        one.unsubscribe(first_queue)
        assert not service.streams[0].closed
        reconnected, snapshot = one.subscribe()
        assert snapshot["running"] is True
        assert [m["message_id"] for m in snapshot["pending"]] == ["B"]
        extra, _ = one.subscribe()
        await one.steer("B")
        await two.steer("Y")
        assert [(s.session_id, s.inputs) for s in service.streams] == [
            ("one", ["B"]),
            ("two", ["Y"]),
        ]
        await frame(reconnected, "message_status", status="steered", message_id="B")
        await frame(extra, "message_status", status="steered", message_id="B")
        with pytest.raises(InputError):
            await one.steer("Y")
        one.unsubscribe(reconnected)
        one.unsubscribe(extra)
        for stream in service.streams:
            stream.release.set()
        await asyncio.gather(one.active_task, two.active_task)
        assert one.active_task is two.active_task is None
        _, completed = one.subscribe()
        assert completed["running"] is False and completed["pending"] == []
    finally:
        await asyncio.gather(one.close(), two.close())


@pytest.mark.asyncio
async def test_model_error_pauses_pending_and_close_does_not_restart():
    from bumblehive.protocols.errors import AgentError

    service = Service()
    controller = SessionController(service.lease, "one")
    queue, _ = controller.subscribe()
    await controller.submit(message("A"))
    await frame(queue, "message_status", status="started")
    await controller.submit(message("B"))
    service.streams[0].error = AgentError(code="model_error", message="failed")
    service.streams[0].release.set()
    await asyncio.gather(controller.active_task)
    assert [s.content for s in service.streams] == ["A"]
    await controller.submit(message("C"))
    assert controller.active_task is None
    assert [i.request.content for i in controller.pending] == ["B", "C"]
    await controller.close()
    assert controller.pending == []
    with pytest.raises(InputError):
        await controller.submit(message("D"))


@pytest.mark.asyncio
async def test_service_delete_and_shutdown_close_only_the_owned_runs(tmp_path):
    from bumblehive_server.runtime_service import RuntimeService

    streams = Service()
    deleted = []
    closed = []

    async def delete(session_id):
        deleted.append(session_id)
        return True

    async def close():
        closed.append(True)

    service = RuntimeService(tmp_path / "unused.json")
    service._runtime = SimpleNamespace(
        stream=streams.stream, delete_session=delete, close=close
    )
    one = service.get_session_controller("one")
    two = service.get_session_controller("two")
    first, _ = one.subscribe()
    second, _ = two.subscribe()
    try:
        await asyncio.gather(one.submit(message("A")), two.submit(message("B")))
        await frame(first, "message_status", status="started")
        await frame(second, "message_status", status="started")
        await two.submit(message("B follow-up"))
        assert service._active_runs == 2
        assert await service.delete_session("one") is True
        assert deleted == ["one"]
        assert one.active_task is None
        assert streams.streams[0].closed
        assert not streams.streams[1].closed
        assert two.active_task is not None
        assert service._active_runs == 1
    finally:
        await service.shutdown()
    assert two.active_task is None and two.pending == []
    assert streams.streams[1].closed and closed == [True]
    assert service._active_runs == 0
    assert service._session_controllers == {}


@pytest.mark.asyncio
@pytest.mark.parametrize("consume_first", [False, True])
async def test_cancel_restores_only_unconsumed_steers_with_original_ids_config_and_order(
    consume_first,
):
    service = Service()
    controller = SessionController(service.lease, "one")
    queue, _ = controller.subscribe()
    left, tail = message("left"), message("tail")
    first = ChatRequest(
        message_id="B",
        content="same text",
        config={"generation": {"temperature": 0.2}},
        approval_mode="full_access",
    )
    second = ChatRequest(
        message_id="C", content="same text", config={"generation": {"temperature": 0.8}}
    )
    try:
        await controller.submit(message("A"))
        await frame(queue, "message_status", status="started")
        for request in (left, first, second, tail):
            await controller.submit(request)
        await controller.steer("C")
        if consume_first:
            assert service.streams[0].consume_inputs() == [
                {"role": "user", "content": "same text"}
            ]
        await controller.steer("B")
        while not queue.empty():
            queue.get_nowait()
        await controller.cancel()
        expected = [left, first, tail] if consume_first else [left, first, second, tail]
        assert [item.request for item in controller.pending] == expected
        assert controller.pending[1].request is first
        assert controller.active_task is None
        assert len(service.streams) == 1  # Cancellation must leave the queue paused.
        receipts = []
        while not queue.empty():
            receipts.append(queue.get_nowait())
        assert [
            (item["message_id"], item["position"])
            for item in receipts
            if item["type"] == "message_status"
        ] == ([("B", 1)] if consume_first else [("B", 1), ("C", 2)])
        assert receipts[-1]["type"] == "cancelled"
        _, snapshot = controller.subscribe()
        assert [item["message_id"] for item in snapshot["pending"]] == [
            item.message_id for item in expected
        ]
        await controller.steer("B")
        await frame(queue, "message_status", status="started", message_id="B")
        assert controller.active_request is first
        assert controller.approvals.mode == "full_access"
        assert service.streams[-1].inputs == []
    finally:
        await controller.close()


@pytest.mark.asyncio
@pytest.mark.parametrize("failure_kind", ["exception", "result_error"])
async def test_failed_run_restores_unconsumed_steers_and_keeps_followups_paused(
    failure_kind,
):
    from bumblehive.protocols.errors import AgentError

    service = Service()
    controller = SessionController(service.lease, "one")
    queue, _ = controller.subscribe()
    try:
        await controller.submit(message("A"))
        await frame(queue, "message_status", status="started")
        await controller.submit(message("B"), steer=True)
        await controller.submit(message("C"))
        stream = service.streams[0]
        if failure_kind == "exception":
            stream.failure = RuntimeError("request failed")
        else:
            stream.error = AgentError(code="model_error", message="request failed")
        task = controller.active_task
        stream.release.set()
        await task
        assert controller.active_task is None
        assert [item.request.message_id for item in controller.pending] == ["B", "C"]
        assert len(service.streams) == 1
        assert controller._steered == []
    finally:
        await controller.close()


@pytest.mark.asyncio
async def test_cancel_one_session_does_not_restore_or_consume_another_sessions_steer():
    service = Service()
    one = SessionController(service.lease, "one")
    two = SessionController(service.lease, "two")
    first, _ = one.subscribe()
    second, _ = two.subscribe()
    try:
        await asyncio.gather(one.submit(message("A")), two.submit(message("X")))
        await frame(first, "message_status", status="started")
        await frame(second, "message_status", status="started")
        await one.submit(message("same"), steer=True)
        await two.submit(message("same"), steer=True)
        await one.cancel()
        assert [item.request.message_id for item in one.pending] == ["same"]
        assert two.pending == [] and two.active_task is not None
        assert two.active_stream.inputs == ["same"]
        task = two.active_task
        two.active_stream.release.set()
        await task
        assert two.pending == [] and two._steered == []
    finally:
        await asyncio.gather(one.close(), two.close())


@pytest.mark.asyncio
async def test_cancel_during_startup_clears_deferred_steering_intent():
    service = Service()
    service.lease_gate.clear()
    controller = SessionController(service.lease, "one")
    queue, _ = controller.subscribe()
    try:
        await controller.submit(message("A"))
        await controller.submit(message("B"), steer=True)
        await controller.submit(message("C"))
        await controller.cancel()
        assert all(not item.steer_requested for item in controller.pending)
        service.lease_gate.set()
        await controller.steer("C")
        await frame(queue, "message_status", status="started", message_id="C")
        assert service.streams[-1].inputs == []
    finally:
        await controller.close()