"""Tests for SSE streaming endpoint on AgentMCPServer."""

from __future__ import annotations

import asyncio
import json
from pathlib import Path
from unittest.mock import MagicMock

import pytest
from starlette.requests import Request

from anima.shell.mcp_agent import AgentMCPServer


@pytest.fixture
def short_state_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
    """Create short-path state directory for Unix socket tests (macOS 104-char limit)."""
    state_dir = tmp_path / "state"
    state_dir.mkdir(parents=True, exist_ok=True)
    monkeypatch.setenv("ANIMA_STATE", str(state_dir))
    return state_dir


@pytest.fixture
def agent_mcp_server(short_state_dir: Path) -> AgentMCPServer:
    """Create AgentMCPServer instance for testing."""
    mock_runner = MagicMock()
    mock_runner.instance_id = "test-agent-123"
    return AgentMCPServer(
        instance_id="test-agent-123",
        agent_runner=mock_runner,
        interact_timeout_seconds=0.1,
    )


def _stream_request() -> Request:
    return Request(
        {
            "type": "http",
            "method": "GET",
            "path": "/stream",
            "query_string": b"",
            "headers": [],
        }
    )


@pytest.mark.asyncio
async def test_stream_endpoint_returns_409_for_ninth_consumer(
    agent_mcp_server: AgentMCPServer,
) -> None:
    responses = []
    for _ in range(8):
        responses.append(
            await agent_mcp_server._handle_stream_request(_stream_request())
        )

    response_ninth = await agent_mcp_server._handle_stream_request(_stream_request())
    assert response_ninth.status_code == 409
    body = json.loads(response_ninth.body)
    assert body["error"]["code"] == "STREAM_CONFLICT"

    for response in responses:
        await response.body_iterator.aclose()


@pytest.mark.asyncio
async def test_stream_endpoint_fans_out_to_concurrent_consumers(
    agent_mcp_server: AgentMCPServer,
) -> None:
    _response_one = await agent_mcp_server._handle_stream_request(_stream_request())
    _response_two = await agent_mcp_server._handle_stream_request(_stream_request())
    assert len(agent_mcp_server._stream_subscribers) == 2
    subscriber_queues = list(agent_mcp_server._stream_subscribers.values())

    await agent_mcp_server.start_stream_session(turn_number=1)
    await agent_mcp_server.enqueue_stream_chunk(turn_number=1, chunk="hello")
    await agent_mcp_server.enqueue_stream_done(
        turn_number=1,
        stop_reason="complete",
        final_response_chars=5,
        prompt_tokens=1,
        completion_tokens=2,
        total_tokens=3,
    )

    first_queue_events = [
        subscriber_queues[0].get_nowait(),
        subscriber_queues[0].get_nowait(),
    ]
    second_queue_events = [
        subscriber_queues[1].get_nowait(),
        subscriber_queues[1].get_nowait(),
    ]

    assert first_queue_events[0].__class__.__name__ == "StreamChunkEvent"
    assert second_queue_events[0].__class__.__name__ == "StreamChunkEvent"
    assert first_queue_events[0].chunk == "hello"
    assert second_queue_events[0].chunk == "hello"
    assert first_queue_events[1].__class__.__name__ == "StreamDoneEvent"
    assert second_queue_events[1].__class__.__name__ == "StreamDoneEvent"


@pytest.mark.asyncio
async def test_stream_endpoint_drops_non_terminal_events_on_queue_full(
    agent_mcp_server: AgentMCPServer,
) -> None:
    response = await agent_mcp_server._handle_stream_request(_stream_request())
    assert len(agent_mcp_server._stream_subscribers) == 1
    subscriber_queue = next(iter(agent_mcp_server._stream_subscribers.values()))

    for _ in range(32):
        subscriber_queue.put_nowait("sentinel")

    await agent_mcp_server.enqueue_stream_chunk(turn_number=1, chunk="dropped")

    queued_items = list(subscriber_queue._queue)
    assert "sentinel" in queued_items
    assert all(not hasattr(item, "chunk") for item in queued_items)

    await response.body_iterator.aclose()


@pytest.mark.asyncio
async def test_stream_endpoint_terminal_event_guaranteed_after_queue_pressure(
    agent_mcp_server: AgentMCPServer,
) -> None:
    response = await agent_mcp_server._handle_stream_request(_stream_request())
    subscriber_queue = next(iter(agent_mcp_server._stream_subscribers.values()))

    for _ in range(32):
        subscriber_queue.put_nowait("sentinel")

    enqueue_task = asyncio.create_task(
        agent_mcp_server.enqueue_stream_done(
            turn_number=1,
            stop_reason="complete",
            final_response_chars=0,
            prompt_tokens=0,
            completion_tokens=0,
            total_tokens=0,
        )
    )

    await asyncio.sleep(0.05)
    assert not enqueue_task.done()

    for _ in range(32):
        subscriber_queue.get_nowait()

    await enqueue_task
    terminal = subscriber_queue.get_nowait()
    assert terminal is not None
    assert terminal.__class__.__name__ == "StreamDoneEvent"

    await response.body_iterator.aclose()


@pytest.mark.asyncio
async def test_stream_endpoint_cleans_up_subscriber_on_disconnect(
    agent_mcp_server: AgentMCPServer,
) -> None:
    response = await agent_mcp_server._handle_stream_request(_stream_request())
    assert len(agent_mcp_server._stream_subscribers) == 1

    iterator = response.body_iterator.__aiter__()
    first_read = asyncio.create_task(iterator.__anext__())
    await asyncio.sleep(0)
    first_read.cancel()
    with pytest.raises(asyncio.CancelledError):
        await first_read
    await asyncio.sleep(0)

    assert len(agent_mcp_server._stream_subscribers) == 0
    assert agent_mcp_server._stream_consumer_active is False
