from __future__ import annotations

import json

import httpx
import pytest
import respx

from rungate import Rungate
from rungate._exceptions import BudgetExceededError
from rungate._sse import SSEDecoder
from rungate._streaming import ChatCompletionStream, RungateStream, _accumulate_chunks
from rungate.types.rungate_metadata import RungateMetadata
from tests.conftest import MOCK_BASE_URL


def _make_sse_body(*events: str) -> str:
    """Build a raw SSE body from event strings.

    Each event is separated by a blank line (SSE event boundary).
    Events may contain multiple lines (e.g., 'event: rungate\\ndata: {...}').
    """
    return "\n\n".join(events) + "\n\n"


def _chunk_event(content: str, index: int = 0, model: str = "gpt-4o", chunk_id: str = "chatcmpl-test") -> str:
    """Create an OpenAI streaming chunk SSE event."""
    chunk = {
        "id": chunk_id,
        "object": "chat.completion.chunk",
        "model": model,
        "choices": [{"index": index, "delta": {"content": content}, "finish_reason": None}],
    }
    return f"data: {json.dumps(chunk)}"


def _done_event() -> str:
    return "data: [DONE]"


def _rungate_event(
    run_id: str = "run-123",
    step_id: str = "step-456",
    provider: str = "openai",
    model: str = "gpt-4o",
    cost_usd: float = 0.001,
    tokens: dict | None = None,
) -> str:
    data = {
        "run_id": run_id,
        "step_id": step_id,
        "provider": provider,
        "model": model,
        "cost_usd": cost_usd,
        "tokens": tokens or {"input": 10, "output": 5},
    }
    return f"event: rungate\ndata: {json.dumps(data)}"


def _streaming_response(*events: str) -> httpx.Response:
    """Create a mock SSE streaming response."""
    body = _make_sse_body(*events)
    return httpx.Response(
        200,
        text=body,
        headers={"content-type": "text/event-stream", "x-rungate-run-id": "run-123"},
    )


# ─── SSE Decoder Tests ───────────────────────────────────


class TestSSEDecoder:
    def test_parses_data_events(self):
        body = 'data: {"hello": "world"}\n\ndata: {"foo": "bar"}\n\n'
        response = httpx.Response(200, text=body, headers={"content-type": "text/event-stream"})
        decoder = SSEDecoder()
        events = list(decoder.iter_events(response))
        assert len(events) == 2
        assert events[0].event == "message"
        assert events[0].data == '{"hello": "world"}'

    def test_parses_named_events(self):
        body = 'event: rungate\ndata: {"run_id": "abc"}\n\n'
        response = httpx.Response(200, text=body, headers={"content-type": "text/event-stream"})
        decoder = SSEDecoder()
        events = list(decoder.iter_events(response))
        assert len(events) == 1
        assert events[0].event == "rungate"
        assert events[0].data == '{"run_id": "abc"}'

    def test_handles_done_sentinel(self):
        body = "data: [DONE]\n\n"
        response = httpx.Response(200, text=body, headers={"content-type": "text/event-stream"})
        decoder = SSEDecoder()
        events = list(decoder.iter_events(response))
        assert len(events) == 1
        assert events[0].data == "[DONE]"

    def test_ignores_comment_lines(self):
        body = ": this is a comment\ndata: hello\n\n"
        response = httpx.Response(200, text=body, headers={"content-type": "text/event-stream"})
        decoder = SSEDecoder()
        events = list(decoder.iter_events(response))
        assert len(events) == 1
        assert events[0].data == "hello"

    def test_handles_empty_data(self):
        body = "data:\n\n"
        response = httpx.Response(200, text=body, headers={"content-type": "text/event-stream"})
        decoder = SSEDecoder()
        events = list(decoder.iter_events(response))
        assert len(events) == 1
        assert events[0].data == ""

    def test_handles_id_field(self):
        body = "id: evt-1\ndata: hello\n\n"
        response = httpx.Response(200, text=body, headers={"content-type": "text/event-stream"})
        decoder = SSEDecoder()
        events = list(decoder.iter_events(response))
        assert events[0].id == "evt-1"


# ─── RungateStream (Low-Level) Tests ─────────────────────


class TestRungateStream:
    def test_yields_chunks_as_dicts(self):
        response = _streaming_response(
            _chunk_event("Hello"),
            _chunk_event(" world"),
            _done_event(),
            _rungate_event(),
        )
        stream = RungateStream(response)
        chunks = list(stream)
        assert len(chunks) == 2
        assert chunks[0]["choices"][0]["delta"]["content"] == "Hello"
        assert chunks[1]["choices"][0]["delta"]["content"] == " world"

    def test_captures_rungate_metadata(self):
        response = _streaming_response(
            _chunk_event("Hi"),
            _done_event(),
            _rungate_event(run_id="run-abc", cost_usd=0.005, tokens={"input": 100, "output": 50}),
        )
        stream = RungateStream(response)
        list(stream)  # consume

        assert stream.rungate.run_id == "run-abc"
        assert stream.rungate.cost_usd == 0.005
        assert stream.rungate.tokens_used == 150  # input + output

    def test_rungate_event_not_yielded_to_user(self):
        response = _streaming_response(
            _chunk_event("Hi"),
            _done_event(),
            _rungate_event(),
        )
        stream = RungateStream(response)
        chunks = list(stream)
        # Only the data chunk, not the rungate event
        assert len(chunks) == 1

    def test_done_stops_iteration(self):
        response = _streaming_response(
            _chunk_event("Hi"),
            _done_event(),
        )
        stream = RungateStream(response)
        chunks = list(stream)
        assert len(chunks) == 1

    def test_context_manager_closes_response(self):
        response = _streaming_response(_chunk_event("Hi"), _done_event(), _rungate_event())
        with RungateStream(response) as stream:
            list(stream)
        assert response.is_stream_consumed or response.is_closed

    def test_handles_float_tokens_in_rungate_event(self):
        """Float tokens should not crash — converted to int."""
        rungate_data = json.dumps(
            {
                "run_id": "r1",
                "step_id": "s1",
                "provider": "openai",
                "model": "gpt-4o",
                "cost_usd": 0.001,
                "tokens": {"input": 10.5, "output": 5.7},
            }
        )
        response = _streaming_response(
            _chunk_event("Hi"),
            _done_event(),
            f"event: rungate\ndata: {rungate_data}",
        )
        stream = RungateStream(response)
        list(stream)
        assert stream.rungate.tokens_used == 15  # int(10.5) + int(5.7)

    def test_handles_malformed_tokens_in_rungate_event(self):
        """String tokens should not crash — tokens_used becomes None."""
        rungate_data = json.dumps(
            {
                "run_id": "r1",
                "step_id": "s1",
                "provider": "openai",
                "model": "gpt-4o",
                "cost_usd": 0.001,
                "tokens": {"input": "ten", "output": "five"},
            }
        )
        response = _streaming_response(
            _chunk_event("Hi"),
            _done_event(),
            f"event: rungate\ndata: {rungate_data}",
        )
        stream = RungateStream(response)
        list(stream)
        assert stream.rungate.tokens_used is None
        assert stream.rungate.run_id == "r1"  # other fields still parsed

    def test_handles_null_tokens_in_rungate_event(self):
        """Null tokens field should not crash."""
        rungate_data = json.dumps(
            {
                "run_id": "r1",
                "step_id": "s1",
                "provider": "openai",
                "model": "gpt-4o",
                "cost_usd": 0.001,
                "tokens": None,
            }
        )
        response = _streaming_response(
            _chunk_event("Hi"),
            _done_event(),
            f"event: rungate\ndata: {rungate_data}",
        )
        stream = RungateStream(response)
        list(stream)
        assert stream.rungate.tokens_used is None

    def test_handles_invalid_json_in_rungate_event(self):
        """Invalid JSON in rungate event should not crash — silently ignored."""
        response = _streaming_response(
            _chunk_event("Hi"),
            _done_event(),
            "event: rungate\ndata: not-json",
        )
        stream = RungateStream(response)
        list(stream)
        assert stream.rungate.run_id is None  # fallback to empty metadata

    def test_rungate_property_returns_empty_metadata_if_no_event(self):
        response = _streaming_response(_chunk_event("Hi"), _done_event())
        stream = RungateStream(response)
        list(stream)
        assert stream.rungate.run_id is None
        assert stream.rungate.cost_usd is None


# ─── ChatCompletionStream (High-Level) Tests ─────────────


class TestChatCompletionStream:
    def test_text_stream_yields_content(self):
        response = _streaming_response(
            _chunk_event("Hello"),
            _chunk_event(" world"),
            _chunk_event("!"),
            _done_event(),
            _rungate_event(),
        )
        stream = ChatCompletionStream(response)
        texts = list(stream.text_stream)
        assert texts == ["Hello", " world", "!"]

    def test_text_stream_skips_non_content_chunks(self):
        # Chunk with no content in delta
        no_content = json.dumps(
            {
                "id": "chatcmpl-test",
                "object": "chat.completion.chunk",
                "model": "gpt-4o",
                "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
            }
        )
        response = _streaming_response(
            f"data: {no_content}",
            _chunk_event("Hello"),
            _done_event(),
            _rungate_event(),
        )
        stream = ChatCompletionStream(response)
        texts = list(stream.text_stream)
        assert texts == ["Hello"]

    def test_get_final_completion(self):
        response = _streaming_response(
            _chunk_event("Hello"),
            _chunk_event(" world"),
            _done_event(),
            _rungate_event(cost_usd=0.002),
        )
        stream = ChatCompletionStream(response)
        for _ in stream.text_stream:
            pass

        completion = stream.get_final_completion()
        assert completion.choices[0]["message"]["content"] == "Hello world"
        assert completion.rungate.cost_usd == 0.002

    def test_rungate_metadata_after_text_stream(self):
        response = _streaming_response(
            _chunk_event("Hi"),
            _done_event(),
            _rungate_event(run_id="run-xyz", provider="anthropic"),
        )
        stream = ChatCompletionStream(response)
        list(stream.text_stream)

        assert stream.rungate.run_id == "run-xyz"
        assert stream.rungate.provider == "anthropic"

    def test_context_manager_cleanup(self):
        response = _streaming_response(_chunk_event("Hi"), _done_event(), _rungate_event())
        with ChatCompletionStream(response) as stream:
            list(stream.text_stream)

    def test_text_stream_raises_on_reaccess_after_full_consumption(self):
        response = _streaming_response(_chunk_event("Hi"), _done_event(), _rungate_event())
        stream = ChatCompletionStream(response)
        list(stream.text_stream)  # fully consumed

        with pytest.raises(RuntimeError, match="already been consumed"):
            list(stream.text_stream)

    def test_text_stream_raises_on_reaccess_after_early_break(self):
        response = _streaming_response(
            _chunk_event("Hello"),
            _chunk_event(" world"),
            _done_event(),
            _rungate_event(),
        )
        stream = ChatCompletionStream(response)
        for _text in stream.text_stream:
            break  # early break — _started is True

        with pytest.raises(RuntimeError, match="already been consumed"):
            list(stream.text_stream)  # second access blocked

    def test_get_final_completion_auto_consumes(self):
        response = _streaming_response(
            _chunk_event("Auto"),
            _done_event(),
            _rungate_event(),
        )
        stream = ChatCompletionStream(response)
        # Don't iterate text_stream — call get_final_completion directly
        completion = stream.get_final_completion()
        assert completion.choices[0]["message"]["content"] == "Auto"

    def test_get_final_completion_after_early_break(self):
        """Breaking out of text_stream early, then calling get_final_completion."""
        response = _streaming_response(
            _chunk_event("Hello"),
            _chunk_event(" world"),
            _chunk_event("!"),
            _done_event(),
            _rungate_event(),
        )
        stream = ChatCompletionStream(response)
        for _text in stream.text_stream:
            break  # Early exit after first chunk

        # Should not crash — exhausts the suspended generator to get remaining chunks
        completion = stream.get_final_completion()
        content = completion.choices[0]["message"]["content"]
        # Must contain all chunks exactly once (no duplication)
        assert content == "Hello world!", f"Expected 'Hello world!', got {content!r}"


# ─── Integration Tests (Mocked HTTP) ─────────────────────


class TestStreamingIntegration:
    def test_create_stream_true_returns_rungate_stream(self, client: Rungate, mock_router: respx.Router):
        mock_router.post(f"{MOCK_BASE_URL}/v1/chat/completions").mock(
            return_value=_streaming_response(
                _chunk_event("Hi"),
                _done_event(),
                _rungate_event(),
            )
        )

        with client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": "Hello"}],
            stream=True,
        ) as stream:
            chunks = list(stream)
            assert len(chunks) == 1
            assert stream.rungate.run_id == "run-123"

    def test_completions_stream_returns_high_level(self, client: Rungate, mock_router: respx.Router):
        mock_router.post(f"{MOCK_BASE_URL}/v1/chat/completions").mock(
            return_value=_streaming_response(
                _chunk_event("Hello"),
                _chunk_event(" there"),
                _done_event(),
                _rungate_event(cost_usd=0.003),
            )
        )

        with client.chat.completions.stream(
            model="gpt-4o",
            messages=[{"role": "user", "content": "Hi"}],
        ) as stream:
            texts = list(stream.text_stream)
            assert texts == ["Hello", " there"]
            assert stream.rungate.cost_usd == 0.003

    def test_non_sse_200_raises_error(self, client: Rungate, mock_router: respx.Router):
        """200 with application/json (not SSE) raises APIError, not silent empty stream."""
        from rungate._exceptions import APIError

        mock_router.post(f"{MOCK_BASE_URL}/v1/chat/completions").mock(
            return_value=httpx.Response(
                200,
                json={"some": "response"},
                headers={"content-type": "application/json"},
            )
        )

        with pytest.raises(APIError, match="Expected streaming response"):
            client.chat.completions.create(
                model="gpt-4o",
                messages=[{"role": "user", "content": "Hi"}],
                stream=True,
            )

    def test_governance_error_on_streaming_request(self, client: Rungate, mock_router: respx.Router):
        """Governance errors are JSON, not SSE — raised at stream open time."""
        mock_router.post(f"{MOCK_BASE_URL}/v1/chat/completions").mock(
            return_value=httpx.Response(
                402,
                json={
                    "error": {
                        "type": "governance_block",
                        "code": "budget_exceeded",
                        "message": "Over budget",
                        "run_id": "run-123",
                        "details": {},
                    }
                },
                headers={"content-type": "application/json"},
            )
        )

        with pytest.raises(BudgetExceededError):
            client.chat.completions.create(
                model="gpt-4o",
                messages=[{"role": "user", "content": "Hi"}],
                stream=True,
            )


# ─── Async Streaming Tests ────────────────────────────────


class TestAsyncRungateStream:
    @pytest.mark.asyncio
    async def test_async_yields_chunks(self):
        from rungate._streaming import AsyncRungateStream

        response = _streaming_response(
            _chunk_event("Hello"),
            _chunk_event(" async"),
            _done_event(),
            _rungate_event(),
        )
        stream = AsyncRungateStream(response)
        chunks = [chunk async for chunk in stream]
        assert len(chunks) == 2
        assert chunks[0]["choices"][0]["delta"]["content"] == "Hello"
        assert chunks[1]["choices"][0]["delta"]["content"] == " async"

    @pytest.mark.asyncio
    async def test_async_captures_metadata(self):
        from rungate._streaming import AsyncRungateStream

        response = _streaming_response(
            _chunk_event("Hi"),
            _done_event(),
            _rungate_event(run_id="async-run", cost_usd=0.01),
        )
        stream = AsyncRungateStream(response)
        async for _ in stream:
            pass
        assert stream.rungate.run_id == "async-run"
        assert stream.rungate.cost_usd == 0.01


class TestAsyncChatCompletionStream:
    @pytest.mark.asyncio
    async def test_async_text_stream(self):
        from rungate._streaming import AsyncChatCompletionStream

        response = _streaming_response(
            _chunk_event("Async"),
            _chunk_event(" text"),
            _done_event(),
            _rungate_event(),
        )
        stream = AsyncChatCompletionStream(response)
        texts = [text async for text in stream.text_stream]
        assert texts == ["Async", " text"]

    @pytest.mark.asyncio
    async def test_async_get_final_completion(self):
        from rungate._streaming import AsyncChatCompletionStream

        response = _streaming_response(
            _chunk_event("Final"),
            _done_event(),
            _rungate_event(cost_usd=0.005),
        )
        stream = AsyncChatCompletionStream(response)
        async for _ in stream.text_stream:
            pass
        completion = await stream.get_final_completion()
        assert completion.choices[0]["message"]["content"] == "Final"
        assert completion.rungate.cost_usd == 0.005


# ─── Accumulator Tests ───────────────────────────────────


class TestAccumulateChunks:
    def test_accumulates_content(self):
        chunks = [
            {"id": "c1", "model": "gpt-4o", "choices": [{"delta": {"content": "Hello"}, "finish_reason": None}]},
            {"id": "c1", "model": "gpt-4o", "choices": [{"delta": {"content": " world"}, "finish_reason": "stop"}]},
        ]
        result = _accumulate_chunks(chunks, RungateMetadata())
        assert result.choices[0]["message"]["content"] == "Hello world"
        assert result.choices[0]["finish_reason"] == "stop"

    def test_empty_chunks(self):
        result = _accumulate_chunks([], RungateMetadata())
        assert result.id == ""
        assert result.choices == []

    def test_preserves_metadata(self):
        metadata = RungateMetadata(run_id="run-1", cost_usd=0.005)
        result = _accumulate_chunks(
            [{"id": "c1", "model": "gpt-4o", "choices": [{"delta": {"content": "Hi"}, "finish_reason": "stop"}]}],
            metadata,
        )
        assert result.rungate.run_id == "run-1"
        assert result.rungate.cost_usd == 0.005
