"""
Unit tests for the LLM client interface.

Tests cover:
- LLMMessage dataclass
- LLMResponse dataclass with token tracking
- TokenUsage dataclass
- EmbeddingResponse dataclass
- LLMClient model validation
- Exception classes
- MockLLMClient functionality
"""

import json

import httpx
import pytest

from torivers_sdk.testing.mocks import MockLLMClient
from torivers_sdk.tools.llm import (
    APPROVED_CHAT_MODELS,
    APPROVED_EMBEDDING_MODELS,
    APPROVED_MODELS,
    DEFAULT_CHAT_MODEL,
    DEFAULT_EMBEDDING_MODEL,
    EmbeddingResponse,
    HttpLLMClient,
    LLMClient,
    LLMError,
    LLMMessage,
    LLMResponse,
    ModelNotApprovedError,
    RateLimitError,
    TokenLimitExceededError,
    TokenUsage,
)

# =============================================================================
# Constants Tests
# =============================================================================


class TestApprovedModels:
    """Test suite for approved model constants."""

    def test_approved_chat_models_contains_gpt4o(self) -> None:
        """Test that GPT-4o is in approved chat models."""
        assert "gpt-4o" in APPROVED_CHAT_MODELS
        assert "gpt-4o-mini" in APPROVED_CHAT_MODELS

    def test_approved_chat_models_contains_claude(self) -> None:
        """Test that Claude models are in approved chat models."""
        assert "claude-sonnet-4-20250514" in APPROVED_CHAT_MODELS
        assert "claude-3-5-haiku-20241022" in APPROVED_CHAT_MODELS

    def test_approved_embedding_models(self) -> None:
        """Test that embedding models are defined."""
        assert "text-embedding-3-small" in APPROVED_EMBEDDING_MODELS
        assert "text-embedding-3-large" in APPROVED_EMBEDDING_MODELS

    def test_approved_models_is_union(self) -> None:
        """Test that APPROVED_MODELS contains all models."""
        for model in APPROVED_CHAT_MODELS:
            assert model in APPROVED_MODELS
        for model in APPROVED_EMBEDDING_MODELS:
            assert model in APPROVED_MODELS

    def test_default_chat_model(self) -> None:
        """Test default chat model is approved."""
        assert DEFAULT_CHAT_MODEL in APPROVED_CHAT_MODELS
        assert DEFAULT_CHAT_MODEL == "gpt-4o-mini"

    def test_default_embedding_model(self) -> None:
        """Test default embedding model is approved."""
        assert DEFAULT_EMBEDDING_MODEL in APPROVED_EMBEDDING_MODELS
        assert DEFAULT_EMBEDDING_MODEL == "text-embedding-3-small"


# =============================================================================
# LLMMessage Tests
# =============================================================================


class TestLLMMessage:
    """Test suite for LLMMessage dataclass."""

    def test_create_system_message(self) -> None:
        """Test creating a system message."""
        msg = LLMMessage(role="system", content="You are a helpful assistant.")
        assert msg.role == "system"
        assert msg.content == "You are a helpful assistant."

    def test_create_user_message(self) -> None:
        """Test creating a user message."""
        msg = LLMMessage(role="user", content="Hello!")
        assert msg.role == "user"
        assert msg.content == "Hello!"

    def test_create_assistant_message(self) -> None:
        """Test creating an assistant message."""
        msg = LLMMessage(role="assistant", content="Hi there!")
        assert msg.role == "assistant"
        assert msg.content == "Hi there!"

    def test_to_dict(self) -> None:
        """Test converting message to dictionary."""
        msg = LLMMessage(role="user", content="Test message")
        d = msg.to_dict()
        assert d == {"role": "user", "content": "Test message"}


# =============================================================================
# TokenUsage Tests
# =============================================================================


class TestTokenUsage:
    """Test suite for TokenUsage dataclass."""

    def test_create_token_usage(self) -> None:
        """Test creating token usage."""
        usage = TokenUsage(prompt_tokens=10, completion_tokens=20, total_tokens=30)
        assert usage.prompt_tokens == 10
        assert usage.completion_tokens == 20
        assert usage.total_tokens == 30

    def test_default_values(self) -> None:
        """Test default values are zero."""
        usage = TokenUsage()
        assert usage.prompt_tokens == 0
        assert usage.completion_tokens == 0
        assert usage.total_tokens == 0

    def test_from_dict(self) -> None:
        """Test creating from dictionary."""
        data = {"prompt_tokens": 15, "completion_tokens": 25, "total_tokens": 40}
        usage = TokenUsage.from_dict(data)
        assert usage.prompt_tokens == 15
        assert usage.completion_tokens == 25
        assert usage.total_tokens == 40

    def test_from_dict_missing_keys(self) -> None:
        """Test creating from partial dictionary."""
        usage = TokenUsage.from_dict({"total_tokens": 50})
        assert usage.prompt_tokens == 0
        assert usage.completion_tokens == 0
        assert usage.total_tokens == 50

    def test_to_dict(self) -> None:
        """Test converting to dictionary."""
        usage = TokenUsage(prompt_tokens=10, completion_tokens=20, total_tokens=30)
        d = usage.to_dict()
        assert d == {
            "prompt_tokens": 10,
            "completion_tokens": 20,
            "total_tokens": 30,
        }


# =============================================================================
# LLMResponse Tests
# =============================================================================


class TestLLMResponse:
    """Test suite for LLMResponse dataclass."""

    def test_create_response(self) -> None:
        """Test creating an LLM response."""
        response = LLMResponse(
            content="Hello!",
            model="gpt-4o-mini",
            usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
            finish_reason="stop",
        )
        assert response.content == "Hello!"
        assert response.model == "gpt-4o-mini"
        assert response.finish_reason == "stop"

    def test_tokens_used_property(self) -> None:
        """Test tokens_used property returns total tokens."""
        response = LLMResponse(
            content="Test",
            model="gpt-4o",
            usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
        )
        assert response.tokens_used == 30

    def test_prompt_tokens_property(self) -> None:
        """Test prompt_tokens property."""
        response = LLMResponse(
            content="Test",
            model="gpt-4o",
            usage={"prompt_tokens": 15, "completion_tokens": 25, "total_tokens": 40},
        )
        assert response.prompt_tokens == 15

    def test_completion_tokens_property(self) -> None:
        """Test completion_tokens property."""
        response = LLMResponse(
            content="Test",
            model="gpt-4o",
            usage={"prompt_tokens": 15, "completion_tokens": 25, "total_tokens": 40},
        )
        assert response.completion_tokens == 25

    def test_token_usage_property(self) -> None:
        """Test token_usage returns TokenUsage object."""
        response = LLMResponse(
            content="Test",
            model="gpt-4o",
            usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
        )
        token_usage = response.token_usage
        assert isinstance(token_usage, TokenUsage)
        assert token_usage.prompt_tokens == 10
        assert token_usage.completion_tokens == 20
        assert token_usage.total_tokens == 30

    def test_default_usage(self) -> None:
        """Test default usage is empty dict."""
        response = LLMResponse(content="Test", model="gpt-4o")
        assert response.usage == {}
        assert response.tokens_used == 0

    def test_default_finish_reason(self) -> None:
        """Test default finish reason is stop."""
        response = LLMResponse(content="Test", model="gpt-4o")
        assert response.finish_reason == "stop"


# =============================================================================
# EmbeddingResponse Tests
# =============================================================================


class TestEmbeddingResponse:
    """Test suite for EmbeddingResponse dataclass."""

    def test_create_embedding_response(self) -> None:
        """Test creating an embedding response."""
        embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
        response = EmbeddingResponse(
            embeddings=embeddings,
            model="text-embedding-3-small",
            usage={"total_tokens": 10},
        )
        assert response.embeddings == embeddings
        assert response.model == "text-embedding-3-small"

    def test_tokens_used_property(self) -> None:
        """Test tokens_used property."""
        response = EmbeddingResponse(
            embeddings=[[0.1]],
            model="text-embedding-3-small",
            usage={"total_tokens": 25},
        )
        assert response.tokens_used == 25

    def test_dimensions_property(self) -> None:
        """Test dimensions property."""
        embeddings = [[0.1] * 1536, [0.2] * 1536]
        response = EmbeddingResponse(
            embeddings=embeddings,
            model="text-embedding-3-small",
        )
        assert response.dimensions == 1536

    def test_dimensions_empty_embeddings(self) -> None:
        """Test dimensions with empty embeddings."""
        response = EmbeddingResponse(
            embeddings=[],
            model="text-embedding-3-small",
        )
        assert response.dimensions == 0

    def test_default_usage(self) -> None:
        """Test default usage is empty dict."""
        response = EmbeddingResponse(embeddings=[[0.1]], model="test")
        assert response.usage == {}
        assert response.tokens_used == 0


# =============================================================================
# Exception Tests
# =============================================================================


class TestLLMError:
    """Test suite for LLMError exception."""

    def test_basic_error(self) -> None:
        """Test creating a basic LLM error."""
        error = LLMError("Something went wrong")
        assert str(error) == "Something went wrong"
        assert error.code is None
        assert error.details == {}

    def test_error_with_code(self) -> None:
        """Test creating error with code."""
        error = LLMError("API Error", code="api_error")
        assert error.code == "api_error"

    def test_error_with_details(self) -> None:
        """Test creating error with details."""
        error = LLMError("Error", details={"key": "value"})
        assert error.details == {"key": "value"}


class TestModelNotApprovedError:
    """Test suite for ModelNotApprovedError exception."""

    def test_chat_model_error(self) -> None:
        """Test error for unapproved chat model."""
        error = ModelNotApprovedError("invalid-model", "chat")
        assert error.model == "invalid-model"
        assert error.model_type == "chat"
        assert error.code == "model_not_approved"
        assert "invalid-model" in str(error)
        assert "not approved" in str(error)

    def test_embedding_model_error(self) -> None:
        """Test error for unapproved embedding model."""
        error = ModelNotApprovedError("invalid-embed", "embedding")
        assert error.model == "invalid-embed"
        assert error.model_type == "embedding"

    def test_error_details(self) -> None:
        """Test error includes model in details."""
        error = ModelNotApprovedError("bad-model", "chat")
        assert error.details == {"model": "bad-model"}


class TestRateLimitError:
    """Test suite for RateLimitError exception."""

    def test_basic_rate_limit_error(self) -> None:
        """Test creating basic rate limit error."""
        error = RateLimitError()
        assert "Rate limit" in str(error)
        assert error.code == "rate_limit_exceeded"
        assert error.retry_after is None

    def test_rate_limit_with_retry_after(self) -> None:
        """Test rate limit error with retry_after."""
        error = RateLimitError(retry_after=30.0)
        assert error.retry_after == 30.0


class TestTokenLimitExceededError:
    """Test suite for TokenLimitExceededError exception."""

    def test_basic_token_limit_error(self) -> None:
        """Test creating basic token limit error."""
        error = TokenLimitExceededError()
        assert "Token limit" in str(error)
        assert error.code == "token_limit_exceeded"

    def test_token_limit_with_details(self) -> None:
        """Test token limit error with token counts."""
        error = TokenLimitExceededError(
            tokens_requested=5000,
            tokens_available=4096,
        )
        assert error.tokens_requested == 5000
        assert error.tokens_available == 4096
        assert error.details["tokens_requested"] == 5000
        assert error.details["tokens_available"] == 4096


# =============================================================================
# LLMClient Class Method Tests
# =============================================================================


class TestLLMClientClassMethods:
    """Test suite for LLMClient class methods."""

    def test_validate_model_approved_chat(self) -> None:
        """Test validating an approved chat model."""
        LLMClient.validate_model("gpt-4o", "chat")  # Should not raise

    def test_validate_model_approved_embedding(self) -> None:
        """Test validating an approved embedding model."""
        LLMClient.validate_model("text-embedding-3-small", "embedding")

    def test_validate_model_not_approved_raises(self) -> None:
        """Test that unapproved model raises error."""
        with pytest.raises(ModelNotApprovedError) as exc_info:
            LLMClient.validate_model("invalid-model", "chat")
        assert exc_info.value.model == "invalid-model"

    def test_is_model_approved_true(self) -> None:
        """Test is_model_approved returns True for approved model."""
        assert LLMClient.is_model_approved("gpt-4o", "chat") is True
        assert (
            LLMClient.is_model_approved("text-embedding-3-small", "embedding") is True
        )

    def test_is_model_approved_false(self) -> None:
        """Test is_model_approved returns False for unapproved model."""
        assert LLMClient.is_model_approved("invalid", "chat") is False
        assert LLMClient.is_model_approved("invalid", "embedding") is False

    def test_get_default_model_chat(self) -> None:
        """Test getting default chat model."""
        model = LLMClient.get_default_model("chat")
        assert model == DEFAULT_CHAT_MODEL

    def test_get_default_model_embedding(self) -> None:
        """Test getting default embedding model."""
        model = LLMClient.get_default_model("embedding")
        assert model == DEFAULT_EMBEDDING_MODEL

    def test_get_default_raises_outside_runtime(self) -> None:
        """Test that get_default raises outside runtime."""
        with pytest.raises(RuntimeError) as exc_info:
            LLMClient.get_default()
        assert "runtime environment" in str(exc_info.value)


def _build_llm(handler) -> HttpLLMClient:  # type: ignore[type-arg]
    """Build an HttpLLMClient with a mock transport."""
    client = HttpLLMClient(
        base_url="http://proxy:8001",
        token="test-jwt-token",
    )
    client._client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
    return client


class TestGetDefault:
    """Tests for LLMClient.get_default() auto-wiring."""

    def test_returns_http_client_when_env_set(
        self, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        monkeypatch.setenv("TORIVERS_LLM_PROXY_URL", "http://proxy:8001")
        monkeypatch.setenv("TORIVERS_LLM_TOKEN", "my-jwt")

        client = LLMClient.get_default()
        assert isinstance(client, HttpLLMClient)
        assert client._base_url == "http://proxy:8001"
        assert client._token == "my-jwt"


class TestHttpLLMClientChat:
    @pytest.mark.asyncio
    async def test_chat_returns_response(self) -> None:
        def handler(request: httpx.Request) -> httpx.Response:
            assert request.url.path == "/api/llm/chat"
            assert request.method == "POST"
            assert "Bearer test-jwt-token" in request.headers["authorization"]
            payload = json.loads(request.content.decode("utf-8"))
            assert payload["model"] == "gpt-4o-mini"
            assert payload["messages"][0]["role"] == "user"
            return httpx.Response(
                200,
                json={
                    "content": "ok",
                    "model": "gpt-4o-mini",
                    "finish_reason": "stop",
                    "usage": {
                        "prompt_tokens": 1,
                        "completion_tokens": 2,
                        "total_tokens": 3,
                    },
                },
            )

        llm = _build_llm(handler)
        resp = await llm.chat([LLMMessage(role="user", content="hi")])
        assert resp.content == "ok"
        assert resp.tokens_used == 3

    @pytest.mark.asyncio
    async def test_chat_raises_token_limit_on_429_without_retry_after(self) -> None:
        def handler(_: httpx.Request) -> httpx.Response:
            return httpx.Response(429, json={"detail": "Token limit exceeded"})

        llm = _build_llm(handler)
        with pytest.raises(TokenLimitExceededError):
            await llm.chat([LLMMessage(role="user", content="hi")])

    @pytest.mark.asyncio
    async def test_chat_raises_rate_limit_on_429_with_retry_after(self) -> None:
        def handler(_: httpx.Request) -> httpx.Response:
            return httpx.Response(
                429,
                json={"detail": "Too many requests"},
                headers={"Retry-After": "60"},
            )

        llm = _build_llm(handler)
        with pytest.raises(RateLimitError) as exc:
            await llm.chat([LLMMessage(role="user", content="hi")])
        assert exc.value.retry_after == 60.0


class TestHttpLLMClientEmbed:
    @pytest.mark.asyncio
    async def test_embed_with_metadata_returns_usage(self) -> None:
        def handler(request: httpx.Request) -> httpx.Response:
            assert request.url.path == "/api/llm/embed"
            payload = json.loads(request.content.decode("utf-8"))
            assert payload["model"] == "text-embedding-3-small"
            assert payload["texts"] == ["a", "b"]
            return httpx.Response(
                200,
                json={
                    "embeddings": [[0.1, 0.2], [0.3, 0.4]],
                    "model": "text-embedding-3-small",
                    "usage": {"prompt_tokens": 2, "total_tokens": 2},
                },
            )

        llm = _build_llm(handler)
        resp = await llm.embed_with_metadata(["a", "b"])
        assert resp.dimensions == 2
        assert resp.tokens_used == 2


# =============================================================================
# MockLLMClient Tests
# =============================================================================


class TestMockLLMClient:
    """Test suite for MockLLMClient."""

    @pytest.mark.asyncio
    async def test_chat_returns_default_response(self) -> None:
        """Test chat returns default response."""
        llm = MockLLMClient()
        response = await llm.chat([LLMMessage(role="user", content="Hello")])
        assert response.content == "Mock LLM response"

    @pytest.mark.asyncio
    async def test_chat_returns_queued_response(self) -> None:
        """Test chat returns queued response."""
        llm = MockLLMClient()
        llm.add_response("Custom response")
        response = await llm.chat([LLMMessage(role="user", content="Hello")])
        assert response.content == "Custom response"

    @pytest.mark.asyncio
    async def test_chat_consumes_queue(self) -> None:
        """Test chat consumes responses from queue."""
        llm = MockLLMClient()
        llm.add_responses(["First", "Second", "Third"])

        r1 = await llm.chat([LLMMessage(role="user", content="1")])
        r2 = await llm.chat([LLMMessage(role="user", content="2")])
        r3 = await llm.chat([LLMMessage(role="user", content="3")])
        r4 = await llm.chat([LLMMessage(role="user", content="4")])

        assert r1.content == "First"
        assert r2.content == "Second"
        assert r3.content == "Third"
        assert r4.content == "Mock LLM response"  # Falls back to default

    @pytest.mark.asyncio
    async def test_chat_records_calls(self) -> None:
        """Test chat records calls."""
        llm = MockLLMClient()
        await llm.chat(
            [LLMMessage(role="user", content="Test")],
            model="gpt-4o",
            temperature=0.5,
            max_tokens=100,
        )

        assert len(llm.calls) == 1
        call = llm.calls[0]
        assert call["method"] == "chat"
        assert call["model"] == "gpt-4o"
        assert call["temperature"] == 0.5
        assert call["max_tokens"] == 100

    @pytest.mark.asyncio
    async def test_chat_uses_default_model(self) -> None:
        """Test chat uses default model when none specified."""
        llm = MockLLMClient()
        response = await llm.chat([LLMMessage(role="user", content="Test")])
        assert response.model == DEFAULT_CHAT_MODEL

    @pytest.mark.asyncio
    async def test_chat_tracks_tokens(self) -> None:
        """Test chat tracks token usage."""
        llm = MockLLMClient(default_prompt_tokens=15, default_completion_tokens=25)
        response = await llm.chat([LLMMessage(role="user", content="Test")])

        assert response.tokens_used == 40
        assert response.prompt_tokens == 15
        assert response.completion_tokens == 25
        assert llm.total_tokens_used == 40

    @pytest.mark.asyncio
    async def test_complete_returns_response(self) -> None:
        """Test complete returns response."""
        llm = MockLLMClient()
        llm.add_response("Completed text")
        response = await llm.complete("Start of text")
        assert response.content == "Completed text"

    @pytest.mark.asyncio
    async def test_embed_returns_vectors(self) -> None:
        """Test embed returns embedding vectors."""
        llm = MockLLMClient(embedding_dimensions=768)
        embeddings = await llm.embed(["text1", "text2"])

        assert len(embeddings) == 2
        assert len(embeddings[0]) == 768
        assert len(embeddings[1]) == 768

    @pytest.mark.asyncio
    async def test_embed_custom_response(self) -> None:
        """Test embed with custom response."""
        llm = MockLLMClient()
        custom_embeddings = [[0.5, 0.5], [0.3, 0.7]]
        llm.add_embedding_response(custom_embeddings)

        result = await llm.embed(["a", "b"])
        assert result == custom_embeddings

    @pytest.mark.asyncio
    async def test_embed_with_metadata(self) -> None:
        """Test embed_with_metadata returns EmbeddingResponse."""
        llm = MockLLMClient()
        response = await llm.embed_with_metadata(["test text"])

        assert isinstance(response, EmbeddingResponse)
        assert len(response.embeddings) == 1
        assert response.model == DEFAULT_EMBEDDING_MODEL
        assert response.tokens_used > 0


class TestMockLLMClientModelValidation:
    """Test suite for MockLLMClient model validation."""

    @pytest.mark.asyncio
    async def test_validation_disabled_by_default(self) -> None:
        """Test model validation is disabled by default."""
        llm = MockLLMClient()
        # Should not raise even with invalid model
        response = await llm.chat(
            [LLMMessage(role="user", content="Test")],
            model="invalid-model",
        )
        assert response.model == "invalid-model"

    @pytest.mark.asyncio
    async def test_validation_enabled_rejects_invalid_chat_model(self) -> None:
        """Test validation rejects invalid chat model."""
        llm = MockLLMClient(validate_models=True)
        with pytest.raises(ModelNotApprovedError) as exc_info:
            await llm.chat(
                [LLMMessage(role="user", content="Test")],
                model="invalid-model",
            )
        assert exc_info.value.model == "invalid-model"
        assert exc_info.value.model_type == "chat"

    @pytest.mark.asyncio
    async def test_validation_enabled_accepts_valid_chat_model(self) -> None:
        """Test validation accepts valid chat model."""
        llm = MockLLMClient(validate_models=True)
        response = await llm.chat(
            [LLMMessage(role="user", content="Test")],
            model="gpt-4o",
        )
        assert response.model == "gpt-4o"

    @pytest.mark.asyncio
    async def test_validation_enabled_allows_default_model(self) -> None:
        """Test validation allows default model (None)."""
        llm = MockLLMClient(validate_models=True)
        # Should not raise with model=None
        response = await llm.chat([LLMMessage(role="user", content="Test")])
        assert response.model == DEFAULT_CHAT_MODEL

    @pytest.mark.asyncio
    async def test_validation_rejects_invalid_embedding_model(self) -> None:
        """Test validation rejects invalid embedding model."""
        llm = MockLLMClient(validate_models=True)
        with pytest.raises(ModelNotApprovedError) as exc_info:
            await llm.embed(["test"], model="invalid-embed")
        assert exc_info.value.model_type == "embedding"

    @pytest.mark.asyncio
    async def test_validation_accepts_valid_embedding_model(self) -> None:
        """Test validation accepts valid embedding model."""
        llm = MockLLMClient(validate_models=True)
        embeddings = await llm.embed(["test"], model="text-embedding-3-small")
        assert len(embeddings) == 1


class TestMockLLMClientAssertions:
    """Test suite for MockLLMClient assertion helpers."""

    @pytest.mark.asyncio
    async def test_assert_called(self) -> None:
        """Test assert_called passes when called."""
        llm = MockLLMClient()
        await llm.chat([LLMMessage(role="user", content="Test")])
        llm.assert_called()  # Should not raise

    def test_assert_called_fails_when_not_called(self) -> None:
        """Test assert_called fails when not called."""
        llm = MockLLMClient()
        with pytest.raises(AssertionError):
            llm.assert_called()

    @pytest.mark.asyncio
    async def test_assert_called_with_method(self) -> None:
        """Test assert_called with specific method."""
        llm = MockLLMClient()
        await llm.chat([LLMMessage(role="user", content="Test")])
        llm.assert_called("chat")  # Should not raise
        with pytest.raises(AssertionError):
            llm.assert_called("embed")

    def test_assert_not_called(self) -> None:
        """Test assert_not_called passes when not called."""
        llm = MockLLMClient()
        llm.assert_not_called()  # Should not raise

    @pytest.mark.asyncio
    async def test_assert_not_called_fails_when_called(self) -> None:
        """Test assert_not_called fails when called."""
        llm = MockLLMClient()
        await llm.chat([LLMMessage(role="user", content="Test")])
        with pytest.raises(AssertionError):
            llm.assert_not_called()

    @pytest.mark.asyncio
    async def test_assert_call_count(self) -> None:
        """Test assert_call_count."""
        llm = MockLLMClient()
        await llm.chat([LLMMessage(role="user", content="1")])
        await llm.chat([LLMMessage(role="user", content="2")])
        await llm.embed(["test"])

        llm.assert_call_count(3)
        llm.assert_call_count(2, method="chat")
        llm.assert_call_count(1, method="embed")

    @pytest.mark.asyncio
    async def test_get_last_call(self) -> None:
        """Test get_last_call returns last call."""
        llm = MockLLMClient()
        await llm.chat([LLMMessage(role="user", content="First")])
        await llm.chat([LLMMessage(role="user", content="Second")])

        last = llm.get_last_call()
        assert last is not None
        assert last["messages"][0].content == "Second"

    @pytest.mark.asyncio
    async def test_get_last_call_by_method(self) -> None:
        """Test get_last_call with method filter."""
        llm = MockLLMClient()
        await llm.chat([LLMMessage(role="user", content="Chat")])
        await llm.embed(["embed"])

        chat_call = llm.get_last_call("chat")
        assert chat_call is not None
        assert chat_call["method"] == "chat"

        embed_call = llm.get_last_call("embed")
        assert embed_call is not None
        assert embed_call["method"] == "embed"

    def test_get_last_call_returns_none_when_empty(self) -> None:
        """Test get_last_call returns None when no calls."""
        llm = MockLLMClient()
        assert llm.get_last_call() is None

    @pytest.mark.asyncio
    async def test_chat_calls_property(self) -> None:
        """Test chat_calls returns only chat calls."""
        llm = MockLLMClient()
        await llm.chat([LLMMessage(role="user", content="Test")])
        await llm.embed(["test"])
        await llm.complete("prompt")

        assert len(llm.chat_calls) == 1
        assert len(llm.embed_calls) == 1
        assert len(llm.complete_calls) == 1

    @pytest.mark.asyncio
    async def test_clear_calls(self) -> None:
        """Test clear_calls resets state."""
        llm = MockLLMClient()
        await llm.chat([LLMMessage(role="user", content="Test")])
        assert llm.total_tokens_used > 0
        assert len(llm.calls) > 0

        llm.clear_calls()
        assert llm.total_tokens_used == 0
        assert len(llm.calls) == 0

    @pytest.mark.asyncio
    async def test_clear_responses(self) -> None:
        """Test clear_responses empties queue."""
        llm = MockLLMClient()
        llm.add_responses(["A", "B", "C"])
        llm.clear_responses()

        response = await llm.chat([LLMMessage(role="user", content="Test")])
        assert response.content == "Mock LLM response"  # Falls back to default

    def test_set_default_response(self) -> None:
        """Test set_default_response changes default."""
        llm = MockLLMClient()
        llm.set_default_response("New default")
        assert llm._default_response == "New default"


class TestMockLLMClientTokenTracking:
    """Test suite for MockLLMClient token tracking."""

    @pytest.mark.asyncio
    async def test_total_tokens_accumulate(self) -> None:
        """Test total tokens accumulate across calls."""
        llm = MockLLMClient(default_prompt_tokens=10, default_completion_tokens=20)

        await llm.chat([LLMMessage(role="user", content="1")])
        assert llm.total_tokens_used == 30

        await llm.chat([LLMMessage(role="user", content="2")])
        assert llm.total_tokens_used == 60

    @pytest.mark.asyncio
    async def test_embed_tokens_tracked(self) -> None:
        """Test embedding tokens are tracked."""
        llm = MockLLMClient()
        # "hello" = 5 chars = ~2 tokens (5//4 + 1)
        await llm.embed(["hello"])
        assert llm.total_tokens_used > 0

    @pytest.mark.asyncio
    async def test_embed_call_records_tokens(self) -> None:
        """Test embed call records token usage."""
        llm = MockLLMClient()
        await llm.embed(["test string"])

        call = llm.get_last_call("embed")
        assert call is not None
        assert "tokens_used" in call
        assert call["tokens_used"] > 0
