import uuid

import pytest

from haiku.rag.a2a import (
    extract_question_from_task,
    load_message_history,
    save_message_history,
)
from haiku.rag.client import HaikuRAG

pytest.importorskip("fasta2a")

from fasta2a.schema import Message, TextPart  # noqa: E402
from pydantic_ai.messages import (  # noqa: E402
    ModelMessage,
    ModelRequest,
    ModelResponse,
    ToolCallPart,
    ToolReturnPart,
)
from pydantic_ai.messages import (
    TextPart as AITextPart,
)


@pytest.mark.asyncio
async def test_save_and_load_message_history():
    """Test round-trip of saving and loading message history."""
    # Create sample message history with proper part_kind for ModelRequest
    from pydantic_ai.messages import UserPromptPart

    original_history: list[ModelMessage] = [
        ModelRequest(parts=[UserPromptPart(content="What is Python?")]),
        ModelResponse(parts=[AITextPart(content="Python is a programming language")]),
    ]

    # Save to A2A format
    saved_message = save_message_history(original_history)

    # Verify structure
    assert saved_message["role"] == "agent"
    assert saved_message["kind"] == "message"
    assert len(saved_message["parts"]) == 1
    assert saved_message["parts"][0]["kind"] == "data"
    metadata = saved_message["parts"][0].get("metadata")
    assert metadata is not None
    assert metadata.get("type") == "conversation_state"

    # Load it back
    loaded_history = load_message_history([saved_message])

    # Verify it matches
    assert len(loaded_history) == len(original_history)
    # First message is a request with UserPromptPart
    assert isinstance(loaded_history[0], ModelRequest)
    first_part = loaded_history[0].parts[0]
    assert hasattr(first_part, "content")
    assert first_part.content == "What is Python?"  # type: ignore
    # Second message is a response with TextPart
    assert isinstance(loaded_history[1], ModelResponse)
    second_part = loaded_history[1].parts[0]
    assert hasattr(second_part, "content")
    assert second_part.content == "Python is a programming language"  # type: ignore


@pytest.mark.asyncio
async def test_save_and_load_message_history_with_tool_calls():
    """Test saving and loading message history that includes tool calls."""
    from pydantic_ai.messages import UserPromptPart

    original_history: list[ModelMessage] = [
        ModelRequest(parts=[UserPromptPart(content="Search for Python")]),
        ModelResponse(
            parts=[
                ToolCallPart(
                    tool_name="search_documents",
                    args={"query": "Python", "limit": 3},
                    tool_call_id="call_1",
                )
            ]
        ),
        ModelRequest(
            parts=[
                ToolReturnPart(
                    tool_name="search_documents",
                    content="Python is a high-level programming language",
                    tool_call_id="call_1",
                )
            ]
        ),
        ModelResponse(
            parts=[AITextPart(content="Based on the search, Python is a language")]
        ),
    ]

    # Save and load
    saved_message = save_message_history(original_history)
    loaded_history = load_message_history([saved_message])

    # Verify tool calls are preserved
    assert len(loaded_history) == 4
    assert isinstance(loaded_history[1].parts[0], ToolCallPart)
    assert loaded_history[1].parts[0].tool_name == "search_documents"
    assert isinstance(loaded_history[2].parts[0], ToolReturnPart)
    assert loaded_history[2].parts[0].tool_name == "search_documents"


@pytest.mark.asyncio
async def test_extract_question_from_task():
    """Test extracting user question from task history."""
    task_history: list[Message] = [
        Message(
            role="user",
            parts=[TextPart(kind="text", text="What is Python?")],
            kind="message",
            message_id=str(uuid.uuid4()),
        )
    ]

    question = extract_question_from_task(task_history)
    assert question == "What is Python?"


@pytest.mark.asyncio
async def test_extract_question_from_task_no_text():
    """Test extracting question when no text part exists."""
    task_history: list[Message] = [
        Message(
            role="user",
            parts=[],
            kind="message",
            message_id=str(uuid.uuid4()),
        )
    ]

    question = extract_question_from_task(task_history)
    assert question is None


@pytest.mark.asyncio
async def test_lru_memory_storage_lru_eviction():
    """Test that LRUMemoryStorage evicts least recently used contexts."""
    from fasta2a.storage import InMemoryStorage

    from haiku.rag.a2a import LRUMemoryStorage

    base_storage = InMemoryStorage()
    storage = LRUMemoryStorage(storage=base_storage, max_contexts=3)

    # Add 3 contexts (at limit)
    await storage.update_context("ctx1", [])
    await storage.update_context("ctx2", [])
    await storage.update_context("ctx3", [])

    # All 3 should be tracked
    assert len(storage.context_order) == 3
    assert "ctx1" in storage.context_order
    assert "ctx2" in storage.context_order
    assert "ctx3" in storage.context_order

    # Add 4th context - should evict ctx1 (oldest)
    await storage.update_context("ctx4", [])
    assert len(storage.context_order) == 3
    assert "ctx1" not in storage.context_order
    assert "ctx2" in storage.context_order
    assert "ctx3" in storage.context_order
    assert "ctx4" in storage.context_order

    # Access ctx2 (moves it to end)
    await storage.load_context("ctx2")

    # Add 5th context - should evict ctx3 (now oldest since ctx2 was accessed)
    await storage.update_context("ctx5", [])
    assert len(storage.context_order) == 3
    assert "ctx3" not in storage.context_order
    assert "ctx2" in storage.context_order  # Still present (was accessed)
    assert "ctx4" in storage.context_order
    assert "ctx5" in storage.context_order


@pytest.mark.asyncio
async def test_lru_memory_storage_access_order():
    """Test that accessing contexts updates their order."""
    from fasta2a.storage import InMemoryStorage

    from haiku.rag.a2a import LRUMemoryStorage

    base_storage = InMemoryStorage()
    storage = LRUMemoryStorage(storage=base_storage, max_contexts=2)

    # Add 2 contexts
    await storage.update_context("ctx1", [])
    await storage.update_context("ctx2", [])

    # Order should be: ctx1, ctx2
    assert list(storage.context_order.keys()) == ["ctx1", "ctx2"]

    # Load ctx1 (moves to end)
    await storage.load_context("ctx1")
    # Order should be: ctx2, ctx1
    assert list(storage.context_order.keys()) == ["ctx2", "ctx1"]

    # Add ctx3 - should evict ctx2 (oldest)
    await storage.update_context("ctx3", [])
    assert "ctx2" not in storage.context_order
    assert "ctx1" in storage.context_order
    assert "ctx3" in storage.context_order


@pytest.mark.asyncio
async def test_a2a_app_creation(temp_db_path):
    """Test that A2A app can be created successfully."""
    from haiku.rag.a2a import create_a2a_app

    # Create a test database
    async with HaikuRAG(temp_db_path) as client:
        await client.create_document(
            content="Python is a high-level programming language known for its simplicity.",
            uri="python_doc",
        )

    # Create A2A app
    app = create_a2a_app(temp_db_path)

    # Verify app properties
    assert app.name == "haiku-rag"
    assert app.description is not None
    assert "conversational" in app.description.lower()


@pytest.mark.asyncio
async def test_a2a_app_has_skills(temp_db_path):
    """Test that A2A app exposes skills describing its capabilities."""
    from haiku.rag.a2a import create_a2a_app

    # Create a test database
    async with HaikuRAG(temp_db_path) as client:
        await client.create_document(content="Test document", uri="test_doc")

    # Create A2A app
    app = create_a2a_app(temp_db_path)

    # Verify app has skills
    assert app.skills is not None
    assert len(app.skills) > 0

    # Check that at least one skill exists
    skill = app.skills[0]
    assert "id" in skill
    assert "name" in skill
    assert "description" in skill
    assert "tags" in skill
    assert "input_modes" in skill
    assert "output_modes" in skill

    # Verify the skill describes document search/QA capabilities
    skill_text = f"{skill['name']} {skill['description']}".lower()
    assert any(
        keyword in skill_text
        for keyword in ["search", "question", "answer", "document", "knowledge"]
    )