from __future__ import annotations

import io
import json
import unittest

from backend.http.handler import HttpTransportMixin


class StreamError(RuntimeError):
    pass


class TransportStub(HttpTransportMixin):
    def __init__(self) -> None:
        self.statuses: list[int] = []
        self.response_headers: list[tuple[str, str]] = []
        self.wfile = io.BytesIO()
        self.close_connection = False

    def send_response(self, status: int) -> None:
        self.statuses.append(int(status))

    def send_header(self, name: str, value: str) -> None:
        self.response_headers.append((name, value))

    def end_headers(self) -> None:
        pass


class HttpStreamingTests(unittest.TestCase):
    def test_stream_transport_preserves_headers_events_and_completion(self) -> None:
        handler = TransportStub()
        handler.send_ndjson_stream(
            ({"type": "delta", "content": value} for value in ("甲", "乙")),
            (StreamError,),
        )

        self.assertEqual(handler.statuses, [200])
        self.assertEqual(
            dict(handler.response_headers),
            {
                "Content-Type": "application/x-ndjson; charset=utf-8",
                "Cache-Control": "no-cache, no-transform",
                "X-Accel-Buffering": "no",
                "Connection": "close",
            },
        )
        events = [
            json.loads(line)
            for line in handler.wfile.getvalue().decode("utf-8").splitlines()
        ]
        self.assertEqual(
            events,
            [
                {"type": "delta", "content": "甲"},
                {"type": "delta", "content": "乙"},
                {"type": "done"},
            ],
        )
        self.assertTrue(handler.close_connection)

    def test_stream_transport_preserves_feature_error_event(self) -> None:
        def events():
            yield {"type": "delta", "content": "partial"}
            raise StreamError("stream failed")

        handler = TransportStub()
        handler.send_ndjson_stream(events(), (StreamError,))
        payloads = [
            json.loads(line)
            for line in handler.wfile.getvalue().decode("utf-8").splitlines()
        ]
        self.assertEqual(
            payloads,
            [
                {"type": "delta", "content": "partial"},
                {"type": "error", "error": "stream failed"},
            ],
        )
        self.assertTrue(handler.close_connection)


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