"""Tests for llama_cpp_server_proxy.py — the Ollama-backed proxy on port 8080.

All httpx calls to Ollama are mocked so no external service is required.
"""

from __future__ import annotations

from typing import Any
from unittest.mock import MagicMock, patch

from fastapi.testclient import TestClient

# Import the FastAPI app directly (no side effects at import time beyond
# logging configuration and environment variable reads).
from llama_cpp_server_proxy import app


class MockResponse:
    """A minimal mock for httpx.Response."""

    def __init__(self, status_code: int, json_data: Any = None, text: str = "") -> None:
        self.status_code = status_code
        self._json = json_data or {}
        self.text = text

    def json(self) -> Any:
        return self._json


class MockAsyncClient:
    """Mock for httpx.AsyncClient used as an async context manager."""

    def __init__(self, response: MockResponse) -> None:
        self._response = response

    async def __aenter__(self) -> MockAsyncClient:
        return self

    async def __aexit__(self, *args: Any) -> None:
        return None

    async def get(self, *args: Any, **kwargs: Any) -> MockResponse:
        return self._response

    async def post(self, *args: Any, **kwargs: Any) -> MockResponse:
        return self._response


def _mock_client(response: MockResponse) -> MagicMock:
    """Return a MagicMock that replaces httpx.AsyncClient for one call."""
    mock = MagicMock()
    mock.return_value = MockAsyncClient(response)
    return mock


# ---------------------------------------------------------------------------
# GET / — health/info
# ---------------------------------------------------------------------------

def test_root_endpoint() -> None:
    client = TestClient(app)
    resp = client.get("/")
    assert resp.status_code == 200
    body = resp.json()
    assert body["arch"] == "llama.cpp proxy"
    assert body["port"] == 8080
    assert "os" in body
    assert "model" in body


# ---------------------------------------------------------------------------
# GET /api/tokenize — list_models
# ---------------------------------------------------------------------------

def test_list_models_success() -> None:
    models_payload = {
        "models": [
            {
                "name": "phi3.5:latest",
                "details": {"context_length": 131072},
            },
            {
                "name": "dolphin-cb26:latest",
                "details": {"context_length": 32768},
            },
        ]
    }
    with patch("llama_cpp_server_proxy.httpx.AsyncClient", _mock_client(MockResponse(200, models_payload))):
        client = TestClient(app)
        resp = client.get("/api/tokenize")
    assert resp.status_code == 200
    body = resp.json()
    assert len(body["models"]) == 2
    assert body["models"][0]["model"] == "phi3.5:latest"
    assert body["models"][0]["context_length"] == 131072


def test_list_models_non_200() -> None:
    with patch(
        "llama_cpp_server_proxy.httpx.AsyncClient",
        _mock_client(MockResponse(500, None, "internal error")),
    ):
        client = TestClient(app)
        resp = client.get("/api/tokenize")
    assert resp.status_code == 200  # endpoint returns 200 with error message
    body = resp.json()
    assert body["models"] == []
    assert "error" in body["message"].lower()


def test_list_models_httpx_error() -> None:
    def _raise(*args: Any, **kwargs: Any) -> MockAsyncClient:
        raise Exception("connection refused")

    with patch("llama_cpp_server_proxy.httpx.AsyncClient", _raise):
        client = TestClient(app)
        resp = client.get("/api/tokenize")
    assert resp.status_code == 200
    body = resp.json()
    assert body["models"] == []
    assert "connection refused" in body["message"]


# ---------------------------------------------------------------------------
# POST /api/chat
# ---------------------------------------------------------------------------

def test_chat_success() -> None:
    chat_payload = {
        "message": {"role": "assistant", "content": "Hello"},
        "done": True,
    }
    with patch(
        "llama_cpp_server_proxy.httpx.AsyncClient",
        _mock_client(MockResponse(200, chat_payload)),
    ):
        client = TestClient(app)
        resp = client.post(
            "/api/chat",
            json={"model": "phi3.5:latest", "messages": [{"role": "user", "content": "hi"}]},
        )
    assert resp.status_code == 200
    body = resp.json()
    assert body["object"] == "chat.completion"
    assert body["model"] == "phi3.5:latest"
    assert body["choices"][0]["message"]["role"] == "assistant"
    assert body["choices"][0]["message"]["content"] == "Hello"
    assert body["choices"][0]["finish_reason"] == "stop"
    assert "usage" in body


def test_chat_non_200() -> None:
    with patch(
        "llama_cpp_server_proxy.httpx.AsyncClient",
        _mock_client(MockResponse(404, None, "model not found")),
    ):
        client = TestClient(app)
        resp = client.post(
            "/api/chat",
            json={"model": "nonexistent", "messages": [{"role": "user", "content": "hi"}]},
        )
    assert resp.status_code == 500
    body = resp.json()
    assert "error" in body


def test_chat_streaming_endpoint_exists() -> None:
    """Streaming endpoint is registered (async generator can't be tested via TestClient).

    Starlette's TestClient does not support async_generator bodies for
    StreamingResponse, so we only verify the route exists and is wired.
    """
    from llama_cpp_server_proxy import StreamChatIterator

    assert StreamChatIterator is not None
    routes = [r.path for r in app.routes]
    assert "/api/chat" in routes


# ---------------------------------------------------------------------------
# POST /api/completion
# ---------------------------------------------------------------------------

def test_completion_success() -> None:
    completion_payload = {"response": "The quick brown fox"}
    with patch(
        "llama_cpp_server_proxy.httpx.AsyncClient",
        _mock_client(MockResponse(200, completion_payload)),
    ):
        client = TestClient(app)
        resp = client.post(
            "/api/completion",
            json={"model": "phi3.5:latest", "prompt": "The quick"},
        )
    assert resp.status_code == 200
    body = resp.json()
    assert body["object"] == "text_completion"
    assert body["choices"][0]["text"] == "The quick brown fox"


def test_completion_non_200() -> None:
    with patch(
        "llama_cpp_server_proxy.httpx.AsyncClient",
        _mock_client(MockResponse(500, None, "server error")),
    ):
        client = TestClient(app)
        resp = client.post(
            "/api/completion",
            json={"model": "phi3.5:latest", "prompt": "test"},
        )
    assert resp.status_code == 500


# ---------------------------------------------------------------------------
# POST /api/embeddings
# ---------------------------------------------------------------------------

def test_embeddings_success() -> None:
    embeddings_payload = {"embeddings": [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]}
    with patch(
        "llama_cpp_server_proxy.httpx.AsyncClient",
        _mock_client(MockResponse(200, embeddings_payload)),
    ):
        client = TestClient(app)
        resp = client.post(
            "/api/embeddings",
            json={"model": "phi3.5:latest", "input": ["hello", "world"]},
        )
    assert resp.status_code == 200
    body = resp.json()
    assert len(body["embeddings"]) == 2
    assert body["embeddings"][0] == [0.1, 0.2, 0.3]


def test_embeddings_string_input() -> None:
    embeddings_payload = {"embeddings": [[0.1, 0.2, 0.3]]}
    with patch(
        "llama_cpp_server_proxy.httpx.AsyncClient",
        _mock_client(MockResponse(200, embeddings_payload)),
    ):
        client = TestClient(app)
        resp = client.post(
            "/api/embeddings",
            json={"model": "phi3.5:latest", "input": "hello"},
        )
    assert resp.status_code == 200


def test_embeddings_non_200() -> None:
    with patch(
        "llama_cpp_server_proxy.httpx.AsyncClient",
        _mock_client(MockResponse(422, None, "invalid input")),
    ):
        client = TestClient(app)
        resp = client.post(
            "/api/embeddings",
            json={"model": "phi3.5:latest", "input": "test"},
        )
    assert resp.status_code == 500


# ---------------------------------------------------------------------------
# POST /api/tokenize — returns 501
# ---------------------------------------------------------------------------

def test_tokenize_not_implemented() -> None:
    client = TestClient(app)
    resp = client.post(
        "/api/tokenize",
        json={"model": "phi3.5:latest", "content": "hello"},
    )
    assert resp.status_code == 501
    body = resp.json()
    assert "error" in body
    assert "not supported" in body["error"].lower()


# ---------------------------------------------------------------------------
# POST /api/detokenize
# ---------------------------------------------------------------------------

def test_detokenize_success() -> None:
    detokenize_payload = {"response": "hello world"}
    with patch(
        "llama_cpp_server_proxy.httpx.AsyncClient",
        _mock_client(MockResponse(200, detokenize_payload)),
    ):
        client = TestClient(app)
        resp = client.post(
            "/api/detokenize",
            json={"model": "phi3.5:latest", "tokens": [100, 200, 300]},
        )
    assert resp.status_code == 200
    body = resp.json()
    assert "text" in body


def test_detokenize_tokens_as_list_converted() -> None:
    """Verify tokens list is converted to space-separated string before sending."""
    captured_body = {}

    class CapturingAsyncClient:
        async def __aenter__(self):
            return self

        async def __aexit__(self, *args):
            return None

        async def post(self, url, json=None, **kwargs):
            captured_body.update(json or {})
            return MockResponse(200, {"response": "ok"})

    with patch("llama_cpp_server_proxy.httpx.AsyncClient", lambda *a, **k: CapturingAsyncClient()):
        client = TestClient(app)
        resp = client.post(
            "/api/detokenize",
            json={"model": "phi3.5:latest", "tokens": [7, 8, 9]},
        )
    assert resp.status_code == 200
    assert captured_body.get("prompt") == "7 8 9"


def test_detokenize_non_200() -> None:
    with patch(
        "llama_cpp_server_proxy.httpx.AsyncClient",
        _mock_client(MockResponse(400, None, "bad request")),
    ):
        client = TestClient(app)
        resp = client.post(
            "/api/detokenize",
            json={"model": "phi3.5:latest", "tokens": [1, 2, 3]},
        )
    assert resp.status_code == 500


# ---------------------------------------------------------------------------
# Fallback model logic
# ---------------------------------------------------------------------------

def test_fallback_when_primary_unavailable() -> None:
    """When /api/show returns non-200 for primary model, fallback is used."""
    # First call (api/show for model check) returns 404, second call (api/chat) succeeds
    responses = iter(
        [
            MockResponse(404, None, "not found"),  # /api/show
            MockResponse(200, {"message": {"content": "fallback response"}}),  # /api/chat
        ]
    )

    def _client_factory(*args: Any, **kwargs: Any) -> MockAsyncClient:
        return MockAsyncClient(next(responses))

    with patch("llama_cpp_server_proxy.httpx.AsyncClient", _client_factory):
        client = TestClient(app)
        resp = client.post(
            "/api/chat",
            json={"model": "nonexistent-model", "messages": [{"role": "user", "content": "hi"}]},
        )
    assert resp.status_code == 200
    body = resp.json()
    assert body["choices"][0]["message"]["content"] == "fallback response"
