"""Tests for streaming execution."""

from __future__ import annotations

import json

import httpx
import pytest
import respx

from baponi import Baponi, BaponiError
from baponi._events import KeepaliveEvent, OutputEvent, ResultEvent, StatusEvent

FAKE_KEY = "sk-test-key-for-unit-tests"
EXECUTE_URL = "https://api.baponi.ai/v1/sandbox/execute"

MOCK_RESULT = {
    "success": True,
    "stdout": "0\n1\n2\n",
    "stderr": "",
    "exit_code": 0,
}


def _ndjson_bytes(*lines: dict) -> bytes:
    """Build NDJSON bytes from dicts."""
    return b"\n".join(json.dumps(line).encode() for line in lines) + b"\n"


def _stream_response(*lines: dict) -> httpx.Response:
    """Create a mock streaming response with NDJSON content."""
    content = _ndjson_bytes(*lines)
    return httpx.Response(
        200,
        headers={"content-type": "application/x-ndjson"},
        stream=httpx.ByteStream(content),
    )


class TestExecuteStream:
    @respx.mock
    def test_yields_typed_events(self):
        ndjson_lines = [
            {"type": "status", "trace_id": "t-1", "status": "running", "seq": 1},
            {"type": "output", "stream": "stdout", "data": "0\n", "seq": 2},
            {"type": "output", "stream": "stdout", "data": "1\n", "seq": 3},
            {"type": "keepalive", "seq": 4},
            {
                "type": "result",
                "trace_id": "t-1",
                "status": "success",
                "result": MOCK_RESULT,
                "output_truncated": False,
                "seq": 5,
            },
        ]
        respx.post(EXECUTE_URL).mock(return_value=_stream_response(*ndjson_lines))

        with Baponi(api_key=FAKE_KEY) as client:
            with client.execute_stream("for i in range(3): print(i)") as stream:
                events = list(stream)

        assert len(events) == 5
        assert isinstance(events[0], StatusEvent)
        assert isinstance(events[1], OutputEvent)
        assert isinstance(events[2], OutputEvent)
        assert isinstance(events[3], KeepaliveEvent)
        assert isinstance(events[4], ResultEvent)

    @respx.mock
    def test_get_final_result(self):
        ndjson_lines = [
            {"type": "status", "trace_id": "t-1", "status": "running", "seq": 1},
            {
                "type": "result",
                "trace_id": "t-1",
                "status": "success",
                "result": MOCK_RESULT,
                "seq": 2,
            },
        ]
        respx.post(EXECUTE_URL).mock(return_value=_stream_response(*ndjson_lines))

        with Baponi(api_key=FAKE_KEY) as client:
            with client.execute_stream("print('hi')") as stream:
                for _ in stream:
                    pass
                result = stream.get_final_result()

        assert result.success is True
        assert result.stdout == "0\n1\n2\n"

    @respx.mock
    def test_until_done(self):
        ndjson_lines = [
            {"type": "status", "trace_id": "t-1", "status": "running", "seq": 1},
            {"type": "output", "stream": "stdout", "data": "hi\n", "seq": 2},
            {
                "type": "result",
                "trace_id": "t-1",
                "status": "success",
                "result": MOCK_RESULT,
                "seq": 3,
            },
        ]
        respx.post(EXECUTE_URL).mock(return_value=_stream_response(*ndjson_lines))

        with Baponi(api_key=FAKE_KEY) as client:
            with client.execute_stream("print('hi')") as stream:
                result = stream.until_done()

        assert result.success is True

    @respx.mock
    def test_get_final_result_raises_if_no_result(self):
        ndjson_lines = [
            {"type": "status", "trace_id": "t-1", "status": "running", "seq": 1},
        ]
        respx.post(EXECUTE_URL).mock(return_value=_stream_response(*ndjson_lines))

        with Baponi(api_key=FAKE_KEY) as client:
            with client.execute_stream("print('hi')") as stream:
                for _ in stream:
                    pass
                with pytest.raises(BaponiError, match="did not produce a result"):
                    stream.get_final_result()

    @respx.mock
    def test_trace_id_and_last_seq(self):
        ndjson_lines = [
            {"type": "status", "trace_id": "t-abc", "status": "running", "seq": 1},
            {"type": "keepalive", "seq": 2},
            {
                "type": "result",
                "trace_id": "t-abc",
                "status": "success",
                "result": MOCK_RESULT,
                "seq": 3,
            },
        ]
        respx.post(EXECUTE_URL).mock(return_value=_stream_response(*ndjson_lines))

        with Baponi(api_key=FAKE_KEY) as client:
            with client.execute_stream("x") as stream:
                for _ in stream:
                    pass
                assert stream.trace_id == "t-abc"
                assert stream.last_seq == 3

    @respx.mock
    def test_error_response_raises_before_streaming(self):
        respx.post(EXECUTE_URL).mock(
            return_value=httpx.Response(
                429, json={"error": "rate_limited", "message": "Too many requests"}
            )
        )

        with Baponi(api_key=FAKE_KEY) as client:
            from baponi.exceptions import RateLimitError

            with pytest.raises(RateLimitError):
                client.execute_stream("x")

    @respx.mock
    def test_sends_ndjson_accept_header(self):
        ndjson_lines = [
            {"type": "status", "trace_id": "t-1", "status": "running", "seq": 1},
            {
                "type": "result",
                "trace_id": "t-1",
                "status": "success",
                "result": MOCK_RESULT,
                "seq": 2,
            },
        ]
        route = respx.post(EXECUTE_URL).mock(
            return_value=_stream_response(*ndjson_lines)
        )

        with Baponi(api_key=FAKE_KEY) as client:
            with client.execute_stream("x") as stream:
                stream.until_done()

        assert route.calls[0].request.headers["accept"] == "application/x-ndjson"
