"""
Tests for KivotClient.

Uses aiohttp.test_utils to create a mock Go backend, so no real
infrastructure is required. All 18 test cases from the plan are covered.
"""

import asyncio
import json

import pytest
from aiohttp import web
from aiohttp.test_utils import AioHTTPTestCase, TestServer

from .client import EventStream, KivotClient, RunHandle
from .errors import (
    KivotAPIError,
    KivotConnectionError,
    RunNotFoundError,
    StreamError,
)
from .event_types import EventType, StreamEvent, TERMINAL_EVENT_TYPES


# ---------------------------------------------------------------------------
# Mock Go backend
# ---------------------------------------------------------------------------

_MOCK_RUN_ID = "test-run-550e8400"


def _build_demo_events(run_id: str) -> list[dict]:
    """Build a demo event sequence matching Go DemoSequenceTypes."""
    types = [
        EventType.RUN_CREATED,
        EventType.RUN_STARTED,
        EventType.MODEL_DELTA,
        EventType.MODEL_DELTA,
        EventType.MODEL_DELTA,
        EventType.MODEL_DELTA,
        EventType.MODEL_DELTA,
        EventType.MESSAGE_CREATED,
        EventType.RUN_COMPLETED,
    ]
    events = []
    for i, event_type in enumerate(types, start=1):
        payload: dict = {}
        if event_type == EventType.MODEL_DELTA:
            payload = {"content": "Hello "}
        elif event_type == EventType.MESSAGE_CREATED:
            payload = {"content": "Hello from durable agents MVP"}
        events.append(
            {
                "cursor": i,
                "event": {
                    "run_id": run_id,
                    "event_id": f"evt_{i:04d}",
                    "seq": i,
                    "type": int(event_type),
                    "step_id": "",
                    "author": "system",
                    "trace_id": "trace-abc",
                    "agent_id": "test-agent",
                    "payload": payload,
                    "ts": "2026-02-25T10:00:00Z",
                },
            }
        )
    return events


# Track received request body for assertions.
_last_start_body: dict | None = None
_last_approval_body: dict | None = None


async def handle_start(request: web.Request) -> web.Response:
    global _last_start_body
    try:
        _last_start_body = await request.json()
    except Exception:
        _last_start_body = None

    return web.json_response(
        {
            "run_id": _MOCK_RUN_ID,
            "ws_url": f"ws://localhost/v1/runs/stream?run_id={_MOCK_RUN_ID}",
        }
    )


async def handle_cancel(request: web.Request) -> web.Response:
    body = await request.json()
    run_id = body.get("run_id", "")
    if run_id == "unknown-run-id":
        return web.json_response(
            {"error": "run not found"},
            status=404,
        )
    return web.json_response({"ok": True})


async def handle_status(request: web.Request) -> web.Response:
    run_id = request.match_info["run_id"]
    if run_id == "unknown":
        return web.json_response({"error": "not found"}, status=404)
    return web.json_response(
        {
            "run_id": run_id,
            "state": "completed",
            "event_count": 9,
            "latest_cursor": 9,
            "created_at": "2026-02-25T10:00:00Z",
            "updated_at": "2026-02-25T10:00:01Z",
            "terminal_event": {"type": 3, "seq": 9},
        }
    )


async def handle_events(request: web.Request) -> web.Response:
    run_id = request.match_info["run_id"]
    cursor = int(request.query.get("cursor", "0"))
    limit = int(request.query.get("limit", "100"))

    all_events = _build_demo_events(run_id)
    filtered = [e for e in all_events if e["cursor"] > cursor][:limit]

    return web.json_response(
        {
            "run_id": run_id,
            "events": filtered,
            "count": len(filtered),
            "next_cursor": filtered[-1]["cursor"] if filtered else cursor,
        }
    )


async def handle_ws_stream(request: web.Request) -> web.WebSocketResponse:
    """Mock WebSocket stream: send demo sequence events."""
    ws = web.WebSocketResponse()
    await ws.prepare(request)

    run_id = request.query.get("run_id", "test-run")
    cursor = int(request.query.get("cursor", "0"))

    events = _build_demo_events(run_id)

    for event in events:
        if event["cursor"] > cursor:
            await ws.send_json(event)

    await ws.close()
    return ws


# Counter to support the reconnect test.
_ws_reconnect_call_count = 0


async def handle_ws_stream_reconnect(request: web.Request) -> web.WebSocketResponse:
    """
    Mock WebSocket that closes after 3 events on first connect,
    then sends remaining events on reconnect.
    """
    global _ws_reconnect_call_count
    _ws_reconnect_call_count += 1

    ws = web.WebSocketResponse()
    await ws.prepare(request)

    run_id = request.query.get("run_id", "test-run")
    cursor = int(request.query.get("cursor", "0"))

    events = _build_demo_events(run_id)
    filtered = [e for e in events if e["cursor"] > cursor]

    if _ws_reconnect_call_count == 1:
        # First connect: send only first 3 events, then close.
        for event in filtered[:3]:
            await ws.send_json(event)
        await ws.close()
    else:
        # Reconnect: send remaining events.
        for event in filtered:
            await ws.send_json(event)
        await ws.close()

    return ws


async def handle_approval(request: web.Request) -> web.Response:
    global _last_approval_body
    run_id = request.match_info["run_id"]
    _last_approval_body = await request.json()
    return web.json_response(
        {
            "ok": True,
            "run_id": run_id,
            "request_id": _last_approval_body.get("request_id", ""),
            "approved": _last_approval_body.get("approved", True),
        }
    )


_healthz_call_count = 0


async def handle_healthz(request: web.Request) -> web.Response:
    return web.Response(status=200, text="ok")


async def handle_readyz(request: web.Request) -> web.Response:
    return web.Response(status=200, text="ok")


_wait_healthy_call_count = 0


async def handle_healthz_flaky(request: web.Request) -> web.Response:
    """Return 503 twice, then 200."""
    global _wait_healthy_call_count
    _wait_healthy_call_count += 1
    if _wait_healthy_call_count <= 2:
        return web.Response(status=503, text="not ready")
    return web.Response(status=200, text="ok")


async def handle_bad_request(request: web.Request) -> web.Response:
    return web.json_response({"error": "invalid input"}, status=400)


async def handle_service_unavailable(request: web.Request) -> web.Response:
    return web.json_response({"error": "draining"}, status=503)


def create_mock_app() -> web.Application:
    """Create a mock Go backend application."""
    app = web.Application()
    app.router.add_post("/v1/runs/start", handle_start)
    app.router.add_post("/v1/runs/cancel", handle_cancel)
    app.router.add_get("/v1/runs/{run_id}/status", handle_status)
    app.router.add_get("/v1/runs/{run_id}/events", handle_events)
    app.router.add_get("/v1/runs/stream", handle_ws_stream)
    app.router.add_post("/v1/runs/{run_id}/approval", handle_approval)
    app.router.add_get("/healthz", handle_healthz)
    app.router.add_get("/readyz", handle_readyz)
    return app


def create_mock_app_reconnect() -> web.Application:
    """Mock app that closes WebSocket after 3 events on first connect."""
    app = web.Application()
    app.router.add_post("/v1/runs/start", handle_start)
    app.router.add_get("/v1/runs/stream", handle_ws_stream_reconnect)
    app.router.add_get("/healthz", handle_healthz)
    return app


def create_mock_app_flaky_health() -> web.Application:
    """Mock app with flaky healthz endpoint."""
    app = web.Application()
    app.router.add_get("/healthz", handle_healthz_flaky)
    return app


def create_mock_app_errors() -> web.Application:
    """Mock app that returns errors for testing error handling."""
    app = web.Application()
    app.router.add_post("/v1/runs/start", handle_bad_request)
    app.router.add_post("/v1/runs/cancel", handle_service_unavailable)
    return app


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture
async def mock_server():
    """Start mock Go backend, yield (server, base_url), then cleanup."""
    app = create_mock_app()
    server = TestServer(app)
    await server.start_server()
    yield server
    await server.close()


@pytest.fixture
async def client(mock_server):
    """Create KivotClient pointing at the mock server."""
    base_url = f"http://localhost:{mock_server.port}"
    async with KivotClient(base_url, reconnect_delay=0.01) as c:
        yield c


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


class TestCreateRun:
    """Tests for client.runs.create()."""

    @pytest.mark.asyncio
    async def test_create_run(self, client: KivotClient):
        """client.runs.create() returns RunHandle with valid run_id and ws_url."""
        handle = await client.runs.create()
        assert handle.run_id == _MOCK_RUN_ID
        assert "ws://" in handle.ws_url
        assert _MOCK_RUN_ID in handle.ws_url
        assert isinstance(handle, RunHandle)

    @pytest.mark.asyncio
    async def test_create_run_with_params(self, client: KivotClient):
        """create() sends agent_id, input, model_config in the request body."""
        global _last_start_body
        _last_start_body = None

        handle = await client.runs.create(
            agent_id="my-agent",
            input={"message": "hello"},
            model_config={"provider": "openai", "model_name": "gpt-4o"},
        )

        assert handle.run_id == _MOCK_RUN_ID
        assert _last_start_body is not None
        assert _last_start_body["agent_id"] == "my-agent"
        assert _last_start_body["input"] == {"message": "hello"}
        assert _last_start_body["model_config"]["provider"] == "openai"


class TestCancelRun:
    """Tests for client.runs.cancel()."""

    @pytest.mark.asyncio
    async def test_cancel_run(self, client: KivotClient):
        """client.runs.cancel() returns ok response."""
        result = await client.runs.cancel(_MOCK_RUN_ID)
        assert result["ok"] is True

    @pytest.mark.asyncio
    async def test_cancel_unknown_run(self, client: KivotClient):
        """Cancelling an unknown run raises RunNotFoundError."""
        with pytest.raises(RunNotFoundError) as exc_info:
            await client.runs.cancel("unknown-run-id")
        assert exc_info.value.status_code == 404


class TestGetStatus:
    """Tests for client.runs.status()."""

    @pytest.mark.asyncio
    async def test_get_status(self, client: KivotClient):
        """client.runs.status() returns full status dict."""
        status = await client.runs.status(_MOCK_RUN_ID)
        assert status["run_id"] == _MOCK_RUN_ID
        assert status["state"] == "completed"
        assert status["event_count"] == 9
        assert status["latest_cursor"] == 9
        assert "terminal_event" in status
        assert status["terminal_event"]["type"] == 3


class TestGetEvents:
    """Tests for client.runs.events()."""

    @pytest.mark.asyncio
    async def test_get_events(self, client: KivotClient):
        """client.runs.events() returns events list and next_cursor."""
        result = await client.runs.events(_MOCK_RUN_ID)
        assert result["run_id"] == _MOCK_RUN_ID
        assert result["count"] == 9
        assert len(result["events"]) == 9
        assert result["next_cursor"] == 9


class TestStreamFullSequence:
    """Tests for WebSocket streaming."""

    @pytest.mark.asyncio
    async def test_stream_full_sequence(self, client: KivotClient):
        """Streaming all demo events yields 9 events ending with terminal."""
        stream = client.runs.stream(_MOCK_RUN_ID)
        events: list[StreamEvent] = []
        async for event in stream:
            events.append(event)

        assert len(events) == 9
        assert events[0].type == EventType.RUN_CREATED
        assert events[0].type_name == "RUN_CREATED"
        assert events[-1].type == EventType.RUN_COMPLETED
        assert events[-1].is_terminal is True

    @pytest.mark.asyncio
    async def test_stream_with_cursor_resume(self, client: KivotClient):
        """Streaming with cursor=5 returns only events after seq 5."""
        stream = client.runs.stream(_MOCK_RUN_ID, cursor=5)
        events: list[StreamEvent] = []
        async for event in stream:
            events.append(event)

        # Events 6..9 (4 events: MODEL_DELTA, MODEL_DELTA, MESSAGE_CREATED, RUN_COMPLETED)
        assert len(events) == 4
        assert all(e.cursor > 5 for e in events)
        assert events[-1].is_terminal is True


class TestStreamReconnect:
    """Tests for WebSocket reconnect with cursor resume."""

    @pytest.mark.asyncio
    async def test_stream_reconnect(self):
        """Client reconnects after server closes and resumes from cursor."""
        global _ws_reconnect_call_count
        _ws_reconnect_call_count = 0

        app = create_mock_app_reconnect()
        server = TestServer(app)
        await server.start_server()

        try:
            base_url = f"http://localhost:{server.port}"
            async with KivotClient(
                base_url,
                reconnect_delay=0.01,
                max_reconnect_attempts=3,
            ) as client:
                stream = client.runs.stream(_MOCK_RUN_ID)
                events: list[StreamEvent] = []
                async for event in stream:
                    events.append(event)

                # Should get all 9 events across both connections.
                assert len(events) == 9
                assert events[0].type == EventType.RUN_CREATED
                assert events[-1].type == EventType.RUN_COMPLETED
                # Verify reconnect happened.
                assert _ws_reconnect_call_count >= 2
        finally:
            await server.close()


class TestCreateAndStream:
    """Tests for create_and_stream convenience method."""

    @pytest.mark.asyncio
    async def test_create_and_stream(self, client: KivotClient):
        """create_and_stream() creates run and returns iterable stream."""
        stream = await client.runs.create_and_stream()
        assert isinstance(stream, EventStream)

        events: list[StreamEvent] = []
        async for event in stream:
            events.append(event)

        assert len(events) == 9
        assert events[-1].is_terminal is True


class TestStreamEventTypes:
    """Tests for StreamEvent type properties."""

    @pytest.mark.asyncio
    async def test_stream_event_types(self, client: KivotClient):
        """Verify type_name, is_terminal, is_model_delta on stream events."""
        stream = client.runs.stream(_MOCK_RUN_ID)
        events: list[StreamEvent] = []
        async for event in stream:
            events.append(event)

        # RUN_CREATED
        assert events[0].type_name == "RUN_CREATED"
        assert events[0].is_terminal is False
        assert events[0].is_model_delta is False

        # MODEL_DELTA (events 2-6, indices 2-6)
        for i in range(2, 7):
            assert events[i].type_name == "MODEL_DELTA"
            assert events[i].is_model_delta is True
            assert events[i].is_terminal is False
            assert events[i].content == "Hello "

        # MESSAGE_CREATED
        assert events[7].type_name == "MESSAGE_CREATED"
        assert events[7].content == "Hello from durable agents MVP"

        # RUN_COMPLETED (terminal)
        assert events[8].type_name == "RUN_COMPLETED"
        assert events[8].is_terminal is True
        assert events[8].is_model_delta is False


class TestHealthEndpoints:
    """Tests for health check endpoints."""

    @pytest.mark.asyncio
    async def test_healthz(self, client: KivotClient):
        """healthz() returns True when server is healthy."""
        assert await client.healthz() is True

    @pytest.mark.asyncio
    async def test_readyz(self, client: KivotClient):
        """readyz() returns True when server is ready."""
        assert await client.readyz() is True

    @pytest.mark.asyncio
    async def test_wait_for_healthy(self):
        """wait_for_healthy() succeeds after initially unhealthy server."""
        global _wait_healthy_call_count
        _wait_healthy_call_count = 0

        app = create_mock_app_flaky_health()
        server = TestServer(app)
        await server.start_server()

        try:
            base_url = f"http://localhost:{server.port}"
            async with KivotClient(base_url) as client:
                # Should succeed after 2 failures + 1 success.
                await client.wait_for_healthy(timeout=5.0, interval=0.05)
                assert _wait_healthy_call_count >= 3
        finally:
            await server.close()


class TestErrorHandling:
    """Tests for HTTP error mapping."""

    @pytest.mark.asyncio
    async def test_api_error_400(self):
        """Bad request raises KivotAPIError with status 400."""
        app = create_mock_app_errors()
        server = TestServer(app)
        await server.start_server()

        try:
            base_url = f"http://localhost:{server.port}"
            async with KivotClient(base_url) as client:
                with pytest.raises(KivotAPIError) as exc_info:
                    await client.runs.create()
                assert exc_info.value.status_code == 400
                assert "invalid input" in exc_info.value.response_body
        finally:
            await server.close()

    @pytest.mark.asyncio
    async def test_api_error_503(self):
        """Service unavailable raises KivotAPIError with status 503."""
        app = create_mock_app_errors()
        server = TestServer(app)
        await server.start_server()

        try:
            base_url = f"http://localhost:{server.port}"
            async with KivotClient(base_url) as client:
                with pytest.raises(KivotAPIError) as exc_info:
                    await client.runs.cancel("some-run")
                assert exc_info.value.status_code == 503
                assert "draining" in exc_info.value.response_body
        finally:
            await server.close()


class TestContextManager:
    """Tests for async context manager lifecycle."""

    @pytest.mark.asyncio
    async def test_context_manager(self, mock_server):
        """Session is created on enter and closed on exit."""
        base_url = f"http://localhost:{mock_server.port}"
        client = KivotClient(base_url)

        assert client._session is None

        async with client:
            assert client._session is not None
            assert not client._session.closed
            session_ref = client._session

        assert session_ref.closed


class TestApprove:
    """Tests for client.runs.approve()."""

    @pytest.mark.asyncio
    async def test_approve(self, client: KivotClient):
        """approve() sends correct body and returns ok response."""
        global _last_approval_body
        _last_approval_body = None

        result = await client.runs.approve(
            _MOCK_RUN_ID,
            "req_123",
            approved=True,
            decided_by="user@example.com",
            reason="Looks good",
        )

        assert result["ok"] is True
        assert result["run_id"] == _MOCK_RUN_ID
        assert result["request_id"] == "req_123"
        assert result["approved"] is True

        assert _last_approval_body is not None
        assert _last_approval_body["request_id"] == "req_123"
        assert _last_approval_body["approved"] is True
        assert _last_approval_body["decided_by"] == "user@example.com"
        assert _last_approval_body["reason"] == "Looks good"


class TestStreamClose:
    """Tests for stream.close()."""

    @pytest.mark.asyncio
    async def test_stream_close(self, client: KivotClient):
        """close() stops iteration and closes WebSocket."""
        stream = client.runs.stream(_MOCK_RUN_ID)

        events: list[StreamEvent] = []
        count = 0
        async for event in stream:
            events.append(event)
            count += 1
            if count >= 3:
                await stream.close()
                break

        assert len(events) == 3
        assert stream._closed is True

        # Verify further iteration raises StopAsyncIteration.
        more: list[StreamEvent] = []
        async for event in stream:
            more.append(event)
        assert more == []


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
