import importlib
import sys
from json import loads
from logging import getLogger
from pathlib import Path
from types import ModuleType
from unittest import IsolatedAsyncioTestCase

from avalan.agent.orchestrator import Orchestrator
from avalan.entities import MessageContentFile, MessageContentText
from avalan.model import TextGenerationResponse


class StreamingOrchestrator(Orchestrator):
    def __init__(self) -> None:  # type: ignore[no-untyped-def]
        pass

    async def __call__(self, messages, settings=None):
        async def gen():
            yield "a"
            yield "b"

        return TextGenerationResponse(
            gen, logger=getLogger(), use_async_generator=True
        )


class SimpleOrchestrator(Orchestrator):
    def __init__(self) -> None:  # type: ignore[no-untyped-def]
        self.synced = False
        self._model_ids = {"server-model"}
        self.last_messages = None
        self.last_settings = None

    async def __call__(self, messages, settings=None):
        self.last_messages = messages
        self.last_settings = settings

        def output_fn(**_):
            return "c"

        return TextGenerationResponse(
            output_fn,
            logger=getLogger(),
            use_async_generator=False,
            inputs={"input_ids": [[1, 2, 3]]},
        )

    async def sync_messages(self):  # type: ignore[override]
        self.synced = True


class ResponsesEndpointTestCase(IsolatedAsyncioTestCase):
    def setUp(self):
        from fastapi import FastAPI
        from fastapi.testclient import TestClient

        sys.modules.pop("avalan.server.routers.responses", None)
        sys.modules.pop("avalan.server.routers", None)
        sys.modules.pop("avalan.server", None)

        self.FastAPI = FastAPI
        self.TestClient = TestClient

        server_pkg = ModuleType("avalan.server")
        server_pkg.__path__ = [str(Path("src/avalan/server").resolve())]

        from fastapi import Request

        def _get_logger(request: Request):
            return getLogger()

        def _get_orchestrator(request: Request):
            return request.app.state.orchestrator

        server_pkg.di_get_logger = _get_logger
        server_pkg.di_get_orchestrator = _get_orchestrator
        sys.modules["avalan.server"] = server_pkg
        self.responses = importlib.import_module(
            "avalan.server.routers.responses"
        )

    def tearDown(self):
        sys.modules.pop("avalan.server.routers.responses", None)
        sys.modules.pop("avalan.server", None)

    async def test_streaming_responses(self):
        app = self.FastAPI()
        app.state.orchestrator = StreamingOrchestrator()
        app.include_router(self.responses.router)

        client = self.TestClient(app)
        payload = {
            "model": "m",
            "input": [{"role": "user", "content": "hi"}],
            "stream": True,
        }
        with client.stream("POST", "/responses", json=payload) as resp:
            self.assertEqual(resp.status_code, 200)
            lines = list(resp.iter_lines())

        events = [
            i
            for i, line in enumerate(lines)
            if line == "event: response.output_text.delta"
        ]
        first_data = loads(lines[events[0] + 1][6:])
        self.assertEqual(first_data["delta"], "a")
        second_data = loads(lines[events[1] + 1][6:])
        self.assertEqual(second_data["delta"], "b")

        self.assertIn("event: response.completed", lines)
        self.assertEqual(lines[-1], "")

    async def test_non_streaming_response(self):
        app = self.FastAPI()
        orchestrator = SimpleOrchestrator()
        app.state.orchestrator = orchestrator
        app.include_router(self.responses.router)

        client = self.TestClient(app)
        payload = {
            "input": [{"role": "user", "content": "hi"}],
            "stream": False,
        }
        resp = client.post("/responses", json=payload)
        self.assertEqual(resp.status_code, 200)
        body = resp.json()
        self.assertEqual(body["model"], "server-model")
        self.assertEqual(body["output"][0]["content"][0]["text"], "c")
        self.assertTrue(orchestrator.synced)
        self.assertEqual(
            body["usage"],
            {
                "input_text_tokens": 3,
                "output_text_tokens": 1,
                "total_tokens": 4,
            },
        )

    async def test_non_streaming_response_accepts_string_input(self):
        app = self.FastAPI()
        orchestrator = SimpleOrchestrator()
        app.state.orchestrator = orchestrator
        app.include_router(self.responses.router)

        client = self.TestClient(app)
        resp = client.post(
            "/responses",
            json={"input": "Find the claim", "stream": False},
        )

        self.assertEqual(resp.status_code, 200)
        assert orchestrator.last_messages is not None
        self.assertEqual(len(orchestrator.last_messages), 1)
        message = orchestrator.last_messages[0]
        self.assertEqual(message.role.value, "user")
        self.assertEqual(
            message.content,
            MessageContentText(type="text", text="Find the claim"),
        )

    async def test_non_streaming_response_accepts_input_text_blocks(self):
        app = self.FastAPI()
        orchestrator = SimpleOrchestrator()
        app.state.orchestrator = orchestrator
        app.include_router(self.responses.router)

        client = self.TestClient(app)
        payload = {
            "input": [
                {
                    "role": "user",
                    "content": [
                        {"type": "input_text", "text": "Find the claim"},
                        {
                            "type": "input_file",
                            "file_id": "file-1",
                            "filename": "report.pdf",
                        },
                    ],
                }
            ],
            "stream": False,
        }
        resp = client.post("/responses", json=payload)

        self.assertEqual(resp.status_code, 200)
        assert orchestrator.last_messages is not None
        message = orchestrator.last_messages[0]
        assert isinstance(message.content, list)
        self.assertEqual(
            message.content,
            [
                MessageContentText(type="text", text="Find the claim"),
                MessageContentFile(
                    type="file",
                    file={"file_id": "file-1", "filename": "report.pdf"},
                ),
            ],
        )

    async def test_non_streaming_response_accepts_text_format(self):
        app = self.FastAPI()
        orchestrator = SimpleOrchestrator()
        app.state.orchestrator = orchestrator
        app.include_router(self.responses.router)

        client = self.TestClient(app)
        payload = {
            "input": [{"role": "user", "content": "hi"}],
            "text": {
                "format": {"type": "json_object"},
                "stop": "DONE",
            },
            "stream": False,
        }
        resp = client.post("/responses", json=payload)

        self.assertEqual(resp.status_code, 200)
        assert orchestrator.last_settings is not None
        self.assertEqual(
            orchestrator.last_settings.response_format,
            {"type": "json_object"},
        )
        self.assertEqual(orchestrator.last_settings.stop_strings, "DONE")
