import pytest
from pydantic_ai.models.test import TestModel

from haiku.rag.client import HaikuRAG
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState


@pytest.mark.asyncio
async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path):
    """Test deep Q&A graph with mocked LLM using TestModel."""

    # Mock get_model to return TestModel which generates valid schema-compliant data
    def test_model_factory(provider, model, config=None):
        return TestModel()

    # Patch all locations where get_model is imported
    monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
    monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
    monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
    monkeypatch.setattr("haiku.rag.graph.deep_qa.graph.get_model", test_model_factory)

    graph = build_deep_qa_graph()

    state = DeepQAState(
        context=DeepQAContext(original_question="What is haiku.rag?"),
        max_sub_questions=3,
    )

    # Use real client but with TestModel for LLM calls
    client = HaikuRAG(temp_db_path, create=True)
    deps = DeepQADeps(client=client)

    result = await graph.run(state=state, deps=deps)

    # TestModel will generate valid structured output based on schemas
    assert result.answer is not None
    assert isinstance(result.answer, str)

    client.close()