#!/usr/bin/env python3
import json
import unittest
from pathlib import Path
from tempfile import TemporaryDirectory

from model_stream import SSEDecoder, StreamProtocolError, consume_sse


def event(kind, payload):
    return f"event: {kind}\ndata: {json.dumps(payload, separators=(',', ':'))}\n\n".encode()


def completed(*, output_tokens=2, status="completed", incomplete_reason=None):
    response = {
        "id": "resp-1",
        "status": status,
        "usage": {
            "input_tokens": 20,
            "output_tokens": output_tokens,
            "total_tokens": 20 + output_tokens,
        },
    }
    if incomplete_reason is not None:
        response["incomplete_details"] = {"reason": incomplete_reason}
    kind = "response.incomplete" if status == "incomplete" else "response.completed"
    return event(kind, {"type": kind, "response": response})


class ModelStreamTests(unittest.TestCase):
    def test_decoder_preserves_utf8_across_chunks(self):
        payload = event("response.output_text.delta", {
            "type": "response.output_text.delta",
            "delta": "text",
        })
        decoder = SSEDecoder()
        split = payload.index(b"text") + 2
        self.assertEqual(decoder.feed(payload[:split]), [])
        parsed = decoder.feed(payload[split:])
        self.assertEqual(len(parsed), 1)
        self.assertIn("text", parsed[0].data)

    def test_completed_stream_uses_provider_usage(self):
        delta = event("response.output_text.delta", {
            "type": "response.output_text.delta",
            "delta": "hello",
        })
        result = consume_sse(
            [delta[:7], delta[7:], completed(output_tokens=2)],
            context_id="c1",
            packet_id="p1",
            max_output_tokens=8,
        )
        self.assertEqual(result["status"], "completed")
        self.assertTrue(result["reusable"])
        self.assertEqual(result["usage"]["output_tokens"], 2)

    def test_disconnect_writes_metadata_only_checkpoint(self):
        def disconnected():
            yield event("response.output_text.delta", {
                "type": "response.output_text.delta",
                "delta": "secret-partial-text",
            })
            raise ConnectionError("body ended early")

        with TemporaryDirectory() as directory:
            path = Path(directory) / "checkpoint.json"
            result = consume_sse(
                disconnected(),
                context_id="c1",
                packet_id="p1",
                checkpoint_path=path,
                max_output_tokens=64,
            )
            checkpoint_text = path.read_text(encoding="utf-8")
        self.assertEqual(result["status"], "interrupted")
        self.assertFalse(result["reusable"])
        self.assertEqual(result["reason_code"], "stream_disconnected")
        self.assertNotIn("secret-partial-text", checkpoint_text)
        self.assertFalse(json.loads(checkpoint_text)["reusable"])

    def test_stream_ending_without_terminal_event_is_not_reusable(self):
        result = consume_sse(
            [event("response.created", {"type": "response.created"})],
            context_id="c1",
            packet_id="p1",
        )
        self.assertEqual(result["reason_code"], "stream_ended_without_completion")
        self.assertFalse(result["reusable"])

    def test_incomplete_utf8_is_a_protocol_error(self):
        result = consume_sse(
            [b"data: {\"delta\":\"\xe4\xb8"],
            context_id="c1",
            packet_id="p1",
        )
        self.assertEqual(result["status"], "protocol_error")
        self.assertEqual(result["reason_code"], "stream_decode_error")

    def test_context_failure_requires_smaller_fresh_packet(self):
        failed = event("response.failed", {
            "type": "response.failed",
            "response": {
                "id": "resp-2",
                "status": "failed",
                "error": {
                    "code": "context_length_exceeded",
                    "message": "context window exceeded",
                },
            },
        })
        result = consume_sse(
            [failed],
            context_id="c1",
            packet_id="p1",
            profile="feature",
        )
        self.assertEqual(result["status"], "context_budget_exceeded")
        self.assertEqual(result["retry"]["profile"], "doc")
        self.assertFalse(result["retry"]["reuse_packet"])

    def test_provider_output_overrun_is_rejected(self):
        result = consume_sse(
            [completed(output_tokens=9)],
            context_id="c1",
            packet_id="p1",
            max_output_tokens=8,
        )
        self.assertEqual(result["status"], "output_budget_exceeded")
        self.assertFalse(result["reusable"])

    def test_incomplete_response_is_checkpointable_not_reusable(self):
        result = consume_sse(
            [completed(
                output_tokens=8,
                status="incomplete",
                incomplete_reason="max_output_tokens",
            )],
            context_id="c1",
            packet_id="p1",
            max_output_tokens=8,
        )
        self.assertEqual(result["status"], "truncated")
        self.assertFalse(result["reusable"])
        self.assertEqual(result["finish_reason"], "max_output_tokens")

    def test_decoder_enforces_event_budget(self):
        decoder = SSEDecoder(max_event_bytes=8, max_response_bytes=32)
        with self.assertRaisesRegex(StreamProtocolError, "event exceeds"):
            decoder.feed(b"data: 123456789\n\n")


if __name__ == "__main__":
    unittest.main()