import os
import pathlib
import unittest
from unittest import mock

from fastapi.testclient import TestClient

from mrjunejune.inference.litellm_proxy import (
    create_app,
    resolve_token_directory,
)


class Chunk:
    def __init__(self, content):
        self.content = content

    def model_dump(self, exclude_none=True):
        del exclude_none
        return {
            "choices": [
                {"index": 0, "delta": {"content": self.content}}
            ]
        }


class Stream:
    def __init__(self):
        self.chunks = iter([Chunk("hello "), Chunk("traveler")])

    def __aiter__(self):
        return self

    async def __anext__(self):
        try:
            return next(self.chunks)
        except StopIteration as error:
            raise StopAsyncIteration from error


class LiteLlmGatewayTest(unittest.TestCase):
    def setUp(self):
        os.environ["LITELLM_MASTER_KEY"] = "test-key"
        self.calls = []

        async def completion(**payload):
            self.calls.append(payload)
            return Stream()

        self.client = TestClient(create_app(completion))

    def test_health_does_not_require_provider_auth(self):
        response = self.client.get("/health/liveliness")
        self.assertEqual(response.status_code, 200)
        self.assertEqual(response.json(), {"status": "ready"})

    def test_stream_requires_key_and_maps_alias(self):
        denied = self.client.post(
            "/v1/chat/completions",
            json={"model": "jrpg-copilot", "messages": [], "stream": True},
        )
        self.assertEqual(denied.status_code, 401)

        with self.client.stream(
            "POST",
            "/v1/chat/completions",
            headers={"Authorization": "Bearer test-key"},
            json={
                "model": "jrpg-copilot",
                "messages": [{"role": "user", "content": "hello"}],
                "stream": True,
                "max_tokens": 99999,
            },
        ) as response:
            body = "".join(response.iter_text())
        self.assertEqual(response.status_code, 200)
        self.assertIn("hello ", body)
        self.assertIn("traveler", body)
        self.assertIn("data: [DONE]", body)
        self.assertEqual(self.calls[0]["model"], "github_copilot/gpt-4")
        self.assertEqual(self.calls[0]["max_tokens"], 1024)

    def test_token_directory_precedence(self):
        with mock.patch.dict(
            os.environ,
            {
                "GITHUB_COPILOT_TOKEN_DIR": "/tmp/copilot-env",
                "MRJUNEJUNE_INFERENCE_STATE": "/tmp/inference-state",
            },
            clear=True,
        ):
            self.assertEqual(
                resolve_token_directory("/tmp/copilot-argument"),
                pathlib.Path("/tmp/copilot-argument"),
            )
            self.assertEqual(
                resolve_token_directory(),
                pathlib.Path("/tmp/copilot-env"),
            )

    def test_token_directory_uses_persistent_state_root(self):
        with mock.patch.dict(
            os.environ,
            {"MRJUNEJUNE_INFERENCE_STATE": "/tmp/inference-state"},
            clear=True,
        ):
            self.assertEqual(
                resolve_token_directory(),
                pathlib.Path("/tmp/inference-state/litellm-copilot"),
            )


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