"""Tests for the central LLM client (app/ai/llm_client.py).

Covers llm_complete(), llm_embed(), helper functions, and the
LLMClient backward-compatibility class. All LiteLLM calls are mocked —
no real API requests are made.
"""

from __future__ import annotations

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

import pytest

from app.ai.llm_client import (
    BASE_BACKOFF_SECONDS,
    DEFAULT_MAX_RETRIES,
    LLMClient,
    LLMResponse,
    _classify_error,
    _extract_cost_usd,
    _extract_usage,
    build_model,
    get_llm_client,
    llm_complete,
    llm_embed,
    reset_llm_client,
)


# ──────────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────────


def _mock_completion_response(
    content: str = "Hello!",
    prompt_tokens: int = 10,
    completion_tokens: int = 5,
) -> MagicMock:
    """Build a fake LiteLLM completion response object."""
    resp = MagicMock()
    resp.choices = [MagicMock()]
    resp.choices[0].message.content = content
    resp.usage = MagicMock(
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        total_tokens=prompt_tokens + completion_tokens,
    )
    return resp


def _mock_embedding_response(count: int = 1, dims: int = 4) -> MagicMock:
    """Build a fake LiteLLM embedding response object."""
    resp = MagicMock()
    resp.data = [{"embedding": [0.1] * dims} for _ in range(count)]
    return resp


# ──────────────────────────────────────────────────────────────────────────
# TestLLMComplete
# ──────────────────────────────────────────────────────────────────────────


class TestLLMComplete:
    """Tests for llm_complete() — mock mode, parameter pass-through, errors."""

    @pytest.mark.asyncio
    async def test_basic_completion_no_api_key(self) -> None:
        """llm_complete() without api_key should still call litellm.acompletion."""
        mock_resp = _mock_completion_response(content="Hi there")
        with patch("app.ai.llm_client.litellm.acompletion", new_callable=AsyncMock) as mock_acompletion:
            mock_acompletion.return_value = mock_resp
            result = await llm_complete(
                model="gpt-4o",
                messages=[{"role": "user", "content": "hello"}],
                api_key=None,
            )
        assert result["content"] == "Hi there"
        assert result["model"] == "gpt-4o"
        assert result["usage"]["prompt_tokens"] == 10
        assert result["usage"]["completion_tokens"] == 5
        assert result["usage"]["total_tokens"] == 15
        # api_key should not be in kwargs
        call_kwargs = mock_acompletion.call_args.kwargs
        assert "api_key" not in call_kwargs

    @pytest.mark.asyncio
    async def test_response_format_passed_through(self) -> None:
        """response_format parameter is passed to litellm.acompletion."""
        mock_resp = _mock_completion_response(content='{"key": "value"}')
        rf = {"type": "json_object"}
        with patch("app.ai.llm_client.litellm.acompletion", new_callable=AsyncMock) as mock_acompletion:
            mock_acompletion.return_value = mock_resp
            await llm_complete(
                model="gpt-4o",
                messages=[{"role": "user", "content": "return json"}],
                response_format=rf,
                api_key="test-key",
            )
        call_kwargs = mock_acompletion.call_args.kwargs
        assert call_kwargs["response_format"] == rf

    @pytest.mark.asyncio
    async def test_tools_parameter_passed_through(self) -> None:
        """tools parameter is passed to litellm.acompletion."""
        mock_resp = _mock_completion_response(content="ok")
        tools = [{"type": "function", "function": {"name": "get_weather"}}]
        with patch("app.ai.llm_client.litellm.acompletion", new_callable=AsyncMock) as mock_acompletion:
            mock_acompletion.return_value = mock_resp
            await llm_complete(
                model="gpt-4o",
                messages=[{"role": "user", "content": "weather?"}],
                tools=tools,
                api_key="test-key",
            )
        call_kwargs = mock_acompletion.call_args.kwargs
        assert call_kwargs["tools"] == tools

    @pytest.mark.asyncio
    async def test_timeout_transient_error_retries(self) -> None:
        """Timeout error is classified as transient and retried."""
        mock_resp = _mock_completion_response(content="success")
        with (
            patch("app.ai.llm_client.litellm.acompletion", new_callable=AsyncMock) as mock_acompletion,
            patch("app.ai.llm_client.asyncio.sleep", new_callable=AsyncMock) as mock_sleep,
        ):
            mock_acompletion.side_effect = [asyncio.TimeoutError("Request timed out"), mock_resp]
            result = await llm_complete(
                model="gpt-4o",
                messages=[{"role": "user", "content": "hi"}],
                api_key="test-key",
                max_retries=2,
            )
        assert result["content"] == "success"
        assert mock_acompletion.call_count == 2
        assert mock_sleep.call_count == 1  # one backoff before retry

    @pytest.mark.asyncio
    async def test_rate_limit_429_transient_retries_with_backoff(self) -> None:
        """429 rate-limit error is transient and retried with exponential backoff."""
        mock_resp = _mock_completion_response(content="ok")
        with (
            patch("app.ai.llm_client.litellm.acompletion", new_callable=AsyncMock) as mock_acompletion,
            patch("app.ai.llm_client.asyncio.sleep", new_callable=AsyncMock) as mock_sleep,
        ):
            mock_acompletion.side_effect = [Exception("Rate limit exceeded: 429"), mock_resp]
            result = await llm_complete(
                model="gpt-4o",
                messages=[{"role": "user", "content": "hi"}],
                api_key="test-key",
                max_retries=2,
            )
        assert result["content"] == "ok"
        assert mock_acompletion.call_count == 2
        # First backoff = BASE_BACKOFF_SECONDS * 2^0 = 1.0
        mock_sleep.assert_called_once_with(BASE_BACKOFF_SECONDS * 1)

    @pytest.mark.asyncio
    async def test_auth_error_permanent_no_retry(self) -> None:
        """401 auth error is permanent — no retry, immediate raise."""
        auth_exc = Exception("Authentication error: 401 Unauthorized")
        with (
            patch("app.ai.llm_client.litellm.acompletion", new_callable=AsyncMock) as mock_acompletion,
            patch("app.ai.llm_client.asyncio.sleep", new_callable=AsyncMock) as mock_sleep,
        ):
            mock_acompletion.side_effect = auth_exc
            with pytest.raises(Exception, match="401"):
                await llm_complete(
                    model="gpt-4o",
                    messages=[{"role": "user", "content": "hi"}],
                    api_key="bad-key",
                    max_retries=3,
                )
        assert mock_acompletion.call_count == 1  # no retry
        assert mock_sleep.call_count == 0

    @pytest.mark.asyncio
    async def test_max_retries_zero_no_retry(self) -> None:
        """max_retries=0 means no retry on transient error."""
        timeout_exc = asyncio.TimeoutError("timed out")
        with (
            patch("app.ai.llm_client.litellm.acompletion", new_callable=AsyncMock) as mock_acompletion,
            patch("app.ai.llm_client.asyncio.sleep", new_callable=AsyncMock) as mock_sleep,
        ):
            mock_acompletion.side_effect = timeout_exc
            with pytest.raises(asyncio.TimeoutError):
                await llm_complete(
                    model="gpt-4o",
                    messages=[{"role": "user", "content": "hi"}],
                    api_key="test-key",
                    max_retries=0,
                )
        assert mock_acompletion.call_count == 1
        assert mock_sleep.call_count == 0

    @pytest.mark.asyncio
    async def test_max_retries_2_then_final_error(self) -> None:
        """max_retries=2 → 2 retries (3 total attempts) then final error."""
        timeout_exc = asyncio.TimeoutError("timed out")
        with (
            patch("app.ai.llm_client.litellm.acompletion", new_callable=AsyncMock) as mock_acompletion,
            patch("app.ai.llm_client.asyncio.sleep", new_callable=AsyncMock) as mock_sleep,
        ):
            mock_acompletion.side_effect = timeout_exc
            with pytest.raises(asyncio.TimeoutError):
                await llm_complete(
                    model="gpt-4o",
                    messages=[{"role": "user", "content": "hi"}],
                    api_key="test-key",
                    max_retries=2,
                )
        # 1 initial + 2 retries = 3 total calls
        assert mock_acompletion.call_count == 3
        assert mock_sleep.call_count == 2

    @pytest.mark.asyncio
    async def test_provider_prefix_applied(self) -> None:
        """provider parameter causes build_model prefix to be applied."""
        mock_resp = _mock_completion_response(content="ok")
        with patch("app.ai.llm_client.litellm.acompletion", new_callable=AsyncMock) as mock_acompletion:
            mock_acompletion.return_value = mock_resp
            await llm_complete(
                model="gpt-4o",
                messages=[{"role": "user", "content": "hi"}],
                provider="anthropic",
                api_key="test-key",
            )
        call_kwargs = mock_acompletion.call_args.kwargs
        assert call_kwargs["model"] == "anthropic/gpt-4o"


# ──────────────────────────────────────────────────────────────────────────
# TestLLMEmbed
# ──────────────────────────────────────────────────────────────────────────


class TestLLMEmbed:
    """Tests for llm_embed() — mock mode, single/batch, dimensions."""

    @pytest.mark.asyncio
    async def test_embed_single_text(self) -> None:
        """llm_embed() with a single text returns list[list[float]]."""
        mock_resp = _mock_embedding_response(count=1, dims=4)
        with patch("app.ai.llm_client.litellm.aembedding", new_callable=AsyncMock) as mock_aembedding:
            mock_aembedding.return_value = mock_resp
            result = await llm_embed(
                texts="hello world",
                api_key="test-key",
                model="openai/text-embedding-3-small",
            )
        assert isinstance(result, list)
        assert len(result) == 1
        assert isinstance(result[0], list)
        assert all(isinstance(v, float) for v in result[0])

    @pytest.mark.asyncio
    async def test_embed_batch_texts(self) -> None:
        """llm_embed() with a list of texts returns batch embeddings."""
        mock_resp = _mock_embedding_response(count=3, dims=4)
        with patch("app.ai.llm_client.litellm.aembedding", new_callable=AsyncMock) as mock_aembedding:
            mock_aembedding.return_value = mock_resp
            result = await llm_embed(
                texts=["text one", "text two", "text three"],
                api_key="test-key",
                model="openai/text-embedding-3-small",
            )
        assert len(result) == 3
        assert all(len(emb) == 4 for emb in result)

    @pytest.mark.asyncio
    async def test_embed_dimensions_passed_through(self) -> None:
        """dimensions parameter is passed to litellm.aembedding for text-embedding-3 models."""
        mock_resp = _mock_embedding_response(count=1, dims=768)
        with patch("app.ai.llm_client.litellm.aembedding", new_callable=AsyncMock) as mock_aembedding:
            mock_aembedding.return_value = mock_resp
            await llm_embed(
                texts="hello",
                api_key="test-key",
                model="openai/text-embedding-3-small",
                dimensions=768,
            )
        call_kwargs = mock_aembedding.call_args.kwargs
        assert call_kwargs["dimensions"] == 768

    @pytest.mark.asyncio
    async def test_embed_empty_list_returns_empty(self) -> None:
        """llm_embed() with empty list returns empty list without calling API."""
        with patch("app.ai.llm_client.litellm.aembedding", new_callable=AsyncMock) as mock_aembedding:
            result = await llm_embed(texts=[], api_key="test-key")
        assert result == []
        assert mock_aembedding.call_count == 0

    @pytest.mark.asyncio
    async def test_embed_failure_returns_empty_vectors(self) -> None:
        """On API failure, llm_embed() returns empty vectors for each input text."""
        with patch("app.ai.llm_client.litellm.aembedding", new_callable=AsyncMock) as mock_aembedding:
            mock_aembedding.side_effect = Exception("connection refused")
            result = await llm_embed(
                texts=["a", "b"],
                api_key="test-key",
                model="openai/text-embedding-3-small",
            )
        assert result == [[], []]


# ──────────────────────────────────────────────────────────────────────────
# TestHelpers
# ──────────────────────────────────────────────────────────────────────────


class TestHelpers:
    """Tests for build_model, _classify_error, _extract_cost_usd, _extract_usage."""

    def test_build_model_with_provider(self) -> None:
        """build_model() prepends provider prefix, stripping any existing prefix."""
        assert build_model("gpt-4o", "openai") == "openai/gpt-4o"
        assert build_model("openai/gpt-4o", "anthropic") == "anthropic/gpt-4o"
        assert build_model("claude-3-sonnet", "anthropic") == "anthropic/claude-3-sonnet"

    def test_build_model_without_provider(self) -> None:
        """build_model() returns model unchanged when provider is None."""
        assert build_model("gpt-4o", None) == "gpt-4o"
        assert build_model("openai/gpt-4o", None) == "openai/gpt-4o"

    def test_build_model_empty_provider(self) -> None:
        """build_model() with empty string provider returns model unchanged."""
        assert build_model("gpt-4o", "") == "gpt-4o"

    def test_classify_error_transient_timeout(self) -> None:
        """TimeoutError is classified as transient."""
        assert _classify_error(asyncio.TimeoutError("timed out")) == "transient"
        assert _classify_error(TimeoutError("operation timed out")) == "transient"

    def test_classify_error_transient_rate_limit(self) -> None:
        """Rate limit / 429 / 503 errors are transient."""
        assert _classify_error(Exception("rate limit exceeded")) == "transient"
        assert _classify_error(Exception("429 Too Many Requests")) == "transient"
        assert _classify_error(Exception("503 service unavailable")) == "transient"
        assert _classify_error(Exception("502 bad gateway")) == "transient"
        assert _classify_error(Exception("504 gateway timeout")) == "transient"

    def test_classify_error_permanent_auth(self) -> None:
        """Auth / 401 / 403 errors are permanent."""
        assert _classify_error(Exception("authentication failed")) == "permanent"
        assert _classify_error(Exception("401 Unauthorized")) == "permanent"
        assert _classify_error(Exception("403 Forbidden")) == "permanent"
        assert _classify_error(Exception("invalid api key")) == "permanent"
        assert _classify_error(Exception("invalid_api_key")) == "permanent"

    def test_classify_error_permanent_validation(self) -> None:
        """Validation / 400 / model_not_found errors are permanent."""
        assert _classify_error(Exception("invalid_request")) == "permanent"
        assert _classify_error(Exception("400 bad request")) == "permanent"
        assert _classify_error(Exception("model_not_found")) == "permanent"

    def test_classify_error_unknown_defaults_transient(self) -> None:
        """Unknown errors default to transient (safe to retry)."""
        assert _classify_error(Exception("something weird happened")) == "transient"
        assert _classify_error(ValueError("unexpected value")) == "transient"

    def test_classify_error_permanent_takes_priority(self) -> None:
        """If both permanent and transient keywords match, permanent wins."""
        # Contains both 'timeout' (transient) and '401' (permanent)
        exc = Exception("timeout during authentication: 401")
        assert _classify_error(exc) == "permanent"

    def test_extract_cost_usd_success(self) -> None:
        """_extract_cost_usd() returns cost from litellm.completion_cost."""
        resp = MagicMock()
        with patch("app.ai.llm_client.litellm.completion_cost", return_value=0.0025) as mock_cost:
            cost = _extract_cost_usd(resp, "openai/gpt-4o")
        assert cost == pytest.approx(0.0025)
        mock_cost.assert_called_once_with(resp)

    def test_extract_cost_usd_failure_returns_zero(self) -> None:
        """_extract_cost_usd() returns 0.0 when litellm.completion_cost fails."""
        resp = MagicMock()
        with patch("app.ai.llm_client.litellm.completion_cost", side_effect=Exception("no cost data")):
            cost = _extract_cost_usd(resp, "openai/gpt-4o")
        assert cost == 0.0

    def test_extract_cost_usd_none_returns_zero(self) -> None:
        """_extract_cost_usd() returns 0.0 when completion_cost returns None."""
        resp = MagicMock()
        with patch("app.ai.llm_client.litellm.completion_cost", return_value=None):
            cost = _extract_cost_usd(resp, "openai/gpt-4o")
        assert cost == 0.0

    def test_extract_usage_with_tokens(self) -> None:
        """_extract_usage() returns prompt, completion, total tokens from response."""
        resp = MagicMock()
        resp.usage = MagicMock(prompt_tokens=50, completion_tokens=30, total_tokens=80)
        usage = _extract_usage(resp)
        assert usage == {"prompt_tokens": 50, "completion_tokens": 30, "total_tokens": 80}

    def test_extract_usage_no_usage_attr(self) -> None:
        """_extract_usage() returns zeros when response has no usage attribute."""
        resp = MagicMock()
        resp.usage = None
        usage = _extract_usage(resp)
        assert usage == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}

    def test_extract_usage_missing_total(self) -> None:
        """_extract_usage() computes total_tokens when not present."""
        resp = MagicMock()
        resp.usage = MagicMock(prompt_tokens=20, completion_tokens=10, total_tokens=0)
        # total_tokens is 0 (falsy) → should fall back to prompt + completion
        usage = _extract_usage(resp)
        assert usage["prompt_tokens"] == 20
        assert usage["completion_tokens"] == 10
        assert usage["total_tokens"] == 30  # 20 + 10


# ──────────────────────────────────────────────────────────────────────────
# TestLLMClientCompat
# ──────────────────────────────────────────────────────────────────────────


class TestLLMClientCompat:
    """Tests for LLMClient class, get_llm_client(), reset_llm_client()."""

    @pytest.mark.asyncio
    async def test_mock_mode_calls_mock_generate(self) -> None:
        """LLMClient in mock mode calls _mock_generate (no API call)."""
        client = LLMClient(model=None, api_key=None)
        assert client.is_mock is True

        with patch.object(client, "_mock_generate", new_callable=AsyncMock) as mock_mock_gen:
            mock_mock_gen.return_value = LLMResponse(
                message="mocked", proposed_actions=[], confidence=0.5
            )
            result = await client.generate("create a contact")
        mock_mock_gen.assert_called_once()
        assert result.message == "mocked"

    @pytest.mark.asyncio
    async def test_mock_mode_keyword_matching(self) -> None:
        """LLLMClient mock mode maps keywords to actions via action_mapper."""
        client = LLMClient(model=None, api_key=None)
        result = await client.generate("create a new contact named John")
        assert isinstance(result, LLMResponse)
        assert len(result.proposed_actions) > 0
        assert result.proposed_actions[0]["method"] == "POST"

    @pytest.mark.asyncio
    async def test_mock_mode_no_match(self) -> None:
        """LLMClient mock mode returns empty actions for unrecognized query."""
        client = LLMClient(model=None, api_key=None)
        result = await client.generate("xyzzy nonsense")
        assert result.proposed_actions == []
        assert result.confidence < 0.5

    @pytest.mark.asyncio
    async def test_api_mode_calls_llm_complete(self) -> None:
        """LLMClient in API mode calls llm_complete (mocked)."""
        client = LLMClient(model="gpt-4o", api_key="test-key")
        assert client.is_mock is False

        llm_result = {
            "content": '{"message": "ok", "proposed_actions": [], "confidence": 0.9}',
            "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
            "cost_usd": 0.001,
            "model": "openai/gpt-4o",
            "raw_response": MagicMock(),
        }
        with patch("app.ai.llm_client.llm_complete", new_callable=AsyncMock) as mock_llm_complete:
            mock_llm_complete.return_value = llm_result
            result = await client.generate("list all contacts")
        mock_llm_complete.assert_called_once()
        assert result.message == "ok"
        assert result.confidence == 0.9

    @pytest.mark.asyncio
    async def test_api_mode_fallback_on_error(self) -> None:
        """LLMClient API mode falls back to empty actions on API error."""
        client = LLMClient(model="gpt-4o", api_key="test-key")
        with patch("app.ai.llm_client.llm_complete", new_callable=AsyncMock) as mock_llm_complete:
            mock_llm_complete.side_effect = Exception("API down")
            result = await client.generate("list contacts")
        assert result.proposed_actions == []
        assert result.confidence == 0.1
        assert "API call failed" in result.message

    def test_get_llm_client_singleton(self) -> None:
        """get_llm_client() returns the same instance on repeated calls."""
        reset_llm_client()
        client1 = get_llm_client()
        client2 = get_llm_client()
        assert client1 is client2
        assert isinstance(client1, LLMClient)

    def test_reset_llm_client_clears_instance(self) -> None:
        """reset_llm_client() clears the singleton, next get_llm_client() returns new instance."""
        reset_llm_client()
        client1 = get_llm_client()
        reset_llm_client()
        client2 = get_llm_client()
        assert client1 is not client2

    def test_llm_client_default_mock_mode(self) -> None:
        """LLMClient() with no args and no env vars defaults to mock mode."""
        with patch.dict("os.environ", {}, clear=False):
            # Ensure AI_MODEL and AI_API_KEY are not set
            import os
            env_copy = dict(os.environ)
            env_copy.pop("AI_MODEL", None)
            env_copy.pop("AI_API_KEY", None)
            with patch.dict(os.environ, env_copy, clear=True):
                client = LLMClient()
                assert client.is_mock is True

    def test_llm_client_api_mode_with_model_and_key(self) -> None:
        """LLMClient() with model and api_key is not in mock mode."""
        client = LLMClient(model="gpt-4o", api_key="sk-test")
        assert client.is_mock is False

    def test_llm_response_to_dict(self) -> None:
        """LLMResponse.to_dict() returns correct structure."""
        resp = LLMResponse(message="hello", proposed_actions=[{"method": "GET"}], confidence=0.9)
        d = resp.to_dict()
        assert d == {"message": "hello", "proposed_actions": [{"method": "GET"}], "confidence": 0.9}