"""Unit tests for NDJSON bridge transport and protocol state machine."""

import asyncio
import json
import pytest

from backend.transport import BridgeState, BridgeTransport, NDJSONProtocolError


class MockStream:
    def __init__(self):
        self._queue = asyncio.Queue()

    async def readline(self) -> bytes:
        data = await self._queue.get()
        return data

    def write(self, data: bytes) -> None:
        self._queue.put_nowait(data)

    async def drain(self) -> None:
        pass


@pytest.mark.asyncio
async def test_transport_lifecycle():
    to_backend = MockStream()
    to_agent = MockStream()

    backend_transport = BridgeTransport(reader=to_backend, writer=to_agent)

    # 1. Start from INIT
    assert backend_transport.state == BridgeState.INIT
    await backend_transport.send_start(
        prompt="Test prompt",
        system_instruction="System",
        remaining_ms=180000,
        model_descriptor={"id": "test-model", "context_window_tokens": 1000, "max_output_tokens": 500},
    )
    assert backend_transport.state == BridgeState.RUNNING

    # 2. Agent sends tool request
    to_backend.write(
        json.dumps({
            "v": 1,
            "type": "tool_request",
            "id": "a_1",
            "payload": {"tool": "confluence_view", "parameters": {"page_id": "123"}},
        }).encode("utf-8") + b"\n"
    )

    msg = await backend_transport.read_message_bounded()
    assert msg["type"] == "tool_request"
    assert msg["id"] == "a_1"
    assert "a_1" in backend_transport._inflight_calls

    # 3. Agent sends collection_start while tool call in flight -> should fail!
    to_backend.write(
        json.dumps({
            "v": 1,
            "type": "collection_start",
            "id": "a_2",
            "payload": {"markdown": "# Done", "warnings": []},
        }).encode("utf-8") + b"\n"
    )

    with pytest.raises(NDJSONProtocolError, match="still in flight"):
        await backend_transport.read_message_bounded()

    # 4. Backend answers tool request
    await backend_transport.send_tool_response(
        reply_to="a_1",
        result={"page_id": "123", "markdown": "content", "truncated": False},
        error=None,
    )
    assert "a_1" not in backend_transport._inflight_calls

    # 5. Now agent sends collection_start -> transition to COLLECTING
    to_backend.write(
        json.dumps({
            "v": 1,
            "type": "collection_start",
            "id": "a_3",
            "payload": {"markdown": "# Done", "warnings": []},
        }).encode("utf-8") + b"\n"
    )

    c_msg = await backend_transport.read_message_bounded()
    assert c_msg["type"] == "collection_start"
    assert backend_transport.state == BridgeState.COLLECTING

    # 6. Backend sends collection_ready
    await backend_transport.send_collection_ready(reply_to="a_3")

    # 7. Agent sends complete -> transition to COMPLETE
    to_backend.write(
        json.dumps({
            "v": 1,
            "type": "complete",
            "id": "a_4",
            "payload": {"accepted_transfer_count": 0},
        }).encode("utf-8") + b"\n"
    )

    comp_msg = await backend_transport.read_message_bounded()
    assert comp_msg["type"] == "complete"
    assert backend_transport.state == BridgeState.COMPLETE


@pytest.mark.asyncio
async def test_duplicate_message_ids():
    to_backend = MockStream()
    to_agent = MockStream()
    transport = BridgeTransport(reader=to_backend, writer=to_agent)
    await transport.send_start("p", "s", 1000, {"id": "m", "context_window_tokens": 1, "max_output_tokens": 1})

    # First message with ID a_1
    to_backend.write(
        json.dumps({"v": 1, "type": "tool_request", "id": "a_1", "payload": {"tool": "t", "parameters": {}}}).encode()
        + b"\n"
    )
    await transport.read_message_bounded()

    # Duplicate message with same ID a_1
    to_backend.write(
        json.dumps({"v": 1, "type": "tool_request", "id": "a_1", "payload": {"tool": "t", "parameters": {}}}).encode()
        + b"\n"
    )
    with pytest.raises(NDJSONProtocolError, match="Duplicate message ID"):
        await transport.read_message_bounded()


@pytest.mark.asyncio
async def test_max_inflight_remote_calls():
    to_backend = MockStream()
    to_agent = MockStream()
    transport = BridgeTransport(reader=to_backend, writer=to_agent)
    await transport.send_start("p", "s", 1000, {"id": "m", "context_window_tokens": 1, "max_output_tokens": 1})

    # Send 4 in-flight requests (allowed)
    for i in range(4):
        to_backend.write(
            json.dumps({"v": 1, "type": "tool_request", "id": f"a_{i}", "payload": {"tool": "t", "parameters": {}}}).encode()
            + b"\n"
        )
        await transport.read_message_bounded()

    # 5th in-flight request -> exceeds 4
    to_backend.write(
        json.dumps({"v": 1, "type": "tool_request", "id": "a_5", "payload": {"tool": "t", "parameters": {}}}).encode()
        + b"\n"
    )
    with pytest.raises(NDJSONProtocolError, match="in-flight"):
        await transport.read_message_bounded()