"""
Tests for Agentbyte LLM response types.

Run with: uv run pytest tests/llm/test_llm_types.py -v
"""
import pytest
from pydantic import ValidationError

from agentbyte.llm import (
    AssistantMessage,
    ChatCompletionChunk,
    ChatCompletionResult,
    Usage,
)


# ============================================================================
# Tests: ChatCompletionResult
# ============================================================================


class TestChatCompletionResult:
    """Test LLM chat completion results."""

    def test_completion_result_creation(self):
        """Test creating completion result."""
        result = ChatCompletionResult(
            message=AssistantMessage(content="Hello", source="gpt-4"),
            usage=Usage(tokens_input=10, tokens_output=5),
            model="gpt-4.1-mini",
        )
        assert result.message.content == "Hello"
        assert result.usage.total_tokens == 15
        assert result.model == "gpt-4.1-mini"

    def test_completion_result_with_metadata(self):
        """Test completion result with metadata."""
        result = ChatCompletionResult(
            message=AssistantMessage(content="Result", source="gpt-4"),
            usage=Usage(tokens_input=10, tokens_output=5),
            model="gpt-4.1-mini",
            metadata={"finish_reason": "stop", "log_probs": 0.95},
        )
        assert result.metadata["finish_reason"] == "stop"

    def test_completion_result_with_finish_reason(self):
        """Test completion result with finish_reason."""
        result = ChatCompletionResult(
            message=AssistantMessage(content="Result", source="gpt-4"),
            usage=Usage(tokens_input=10, tokens_output=5),
            model="gpt-4.1-mini",
            finish_reason="stop",
        )
        assert result.finish_reason == "stop"

    def test_completion_result_requires_message(self):
        """Test completion result requires message."""
        with pytest.raises(ValidationError):
            ChatCompletionResult(
                usage=Usage(tokens_input=1, tokens_output=1),
                model="test",
            )


# ============================================================================
# Tests: ChatCompletionChunk
# ============================================================================


class TestChatCompletionChunk:
    """Test streaming chunks."""

    def test_chunk_creation(self):
        """Test creating a chunk."""
        chunk = ChatCompletionChunk(
            content="Hello",
            model="gpt-4.1-mini",
            is_complete=False,
        )
        assert chunk.content == "Hello"
        assert chunk.model == "gpt-4.1-mini"
        assert chunk.is_complete is False

    def test_chunk_empty_content(self):
        """Test chunk can have empty content."""
        chunk = ChatCompletionChunk(content="", model="test", is_complete=True)
        assert chunk.content == ""
        assert chunk.is_complete is True

    def test_chunk_with_metadata(self):
        """Test chunk with metadata."""
        chunk = ChatCompletionChunk(
            content="text",
            model="test",
            is_complete=False,
            metadata={"finish_reason": None},
        )
        assert chunk.metadata["finish_reason"] is None


# ============================================================================
# Tests: Serialization
# ============================================================================


class TestSerialization:
    """Test JSON serialization/deserialization."""

    def test_completion_result_serialization(self):
        """Test serializing completion result."""
        result = ChatCompletionResult(
            message=AssistantMessage(content="Result", source="gpt-4"),
            usage=Usage(tokens_input=10, tokens_output=5),
            model="gpt-4.1-mini",
        )
        json_str = result.model_dump_json()
        restored = ChatCompletionResult.model_validate_json(json_str)
        assert restored.message.content == result.message.content
        assert restored.usage.total_tokens == result.usage.total_tokens
