from unittest.mock import AsyncMock

from haiku.rag.agents.research.models import Citation, ResearchReport
from haiku.rag.client import HaikuRAG
from haiku.rag.skills.rag import (
    STATE_NAMESPACE,
    STATE_TYPE,
    RAGState,
    instructions,
    skill_metadata,
    state_metadata,
)
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.tools.document import DocumentInfo
from haiku.rag.tools.qa import QAHistoryEntry
from haiku.skills.models import SkillMetadata, StateMetadata

from .conftest import _get_tool, _make_ctx


class TestRAGModuleAPI:
    def test_state_type_is_rag_state(self):
        assert STATE_TYPE is RAGState

    def test_state_namespace(self):
        assert STATE_NAMESPACE == "rag"

    def test_state_metadata_returns_state_metadata(self):
        result = state_metadata()
        assert isinstance(result, StateMetadata)
        assert result.namespace == "rag"
        assert result.type is RAGState
        assert result.schema == RAGState.model_json_schema()

    def test_skill_metadata_returns_skill_metadata(self):
        result = skill_metadata()
        assert isinstance(result, SkillMetadata)
        assert result.name == "rag"

    def test_instructions_returns_string(self):
        result = instructions()
        assert isinstance(result, str)
        assert len(result) > 0

    def test_constants_match_create_skill(self, temp_db_path):
        from haiku.rag.skills.rag import create_skill

        skill = create_skill(db_path=temp_db_path)
        assert skill.state_type is STATE_TYPE
        assert skill.state_namespace == STATE_NAMESPACE
        assert skill.metadata == skill_metadata()
        assert skill.instructions == instructions()


class TestRAGSkillCreation:
    def test_create_skill_returns_valid_skill(self, temp_db_path):
        from haiku.rag.skills.rag import create_skill

        skill = create_skill(db_path=temp_db_path)
        assert skill.metadata.name == "rag"
        assert skill.metadata.description
        assert skill.instructions

    def test_create_skill_has_expected_tools(self, temp_db_path):
        from haiku.rag.skills.rag import create_skill

        skill = create_skill(db_path=temp_db_path)
        tool_names = {getattr(t, "__name__") for t in skill.tools if callable(t)}
        assert tool_names == {
            "search",
            "list_documents",
            "get_document",
            "ask",
            "research",
        }

    def test_create_skill_has_state(self, temp_db_path):
        from haiku.rag.skills.rag import RAGState, create_skill

        skill = create_skill(db_path=temp_db_path)
        assert skill._state_type is RAGState
        assert skill._state_namespace == "rag"

    def test_create_skill_from_env(self, monkeypatch, temp_db_path):
        monkeypatch.setenv("HAIKU_RAG_DB", str(temp_db_path))
        from haiku.rag.skills.rag import create_skill

        skill = create_skill()
        assert skill.metadata.name == "rag"


class TestSearchTool:
    async def test_search_returns_formatted_string(self, rag_db):
        from haiku.rag.skills.rag import create_skill

        skill = create_skill(db_path=rag_db)
        search = _get_tool(skill, "search")
        ctx = _make_ctx()
        result = await search(ctx, query="artificial intelligence")
        assert isinstance(result, str)
        assert len(result) > 0

    async def test_search_updates_state(self, rag_db):
        from haiku.rag.skills.rag import RAGState, create_skill

        skill = create_skill(db_path=rag_db)
        search = _get_tool(skill, "search")
        state = RAGState()
        ctx = _make_ctx(state)
        await search(ctx, query="artificial intelligence")
        assert "artificial intelligence" in state.searches
        results = state.searches["artificial intelligence"]
        assert len(results) > 0
        assert isinstance(results[0], SearchResult)

    async def test_search_applies_document_filter_from_state(self, rag_db):
        from haiku.rag.skills.rag import RAGState, create_skill

        skill = create_skill(db_path=rag_db)
        search = _get_tool(skill, "search")
        state = RAGState(document_filter="title = 'AI Overview'")
        ctx = _make_ctx(state)
        result = await search(ctx, query="artificial intelligence")
        assert "AI Overview" in result
        assert "ML Basics" not in result

    async def test_search_without_state(self, rag_db):
        from haiku.rag.skills.rag import create_skill

        skill = create_skill(db_path=rag_db)
        search = _get_tool(skill, "search")
        ctx = _make_ctx(state=None)
        result = await search(ctx, query="artificial intelligence")
        assert isinstance(result, str)


class TestListDocumentsTool:
    async def test_list_documents_returns_results(self, rag_db):
        from haiku.rag.skills.rag import create_skill

        skill = create_skill(db_path=rag_db)
        list_docs = _get_tool(skill, "list_documents")
        ctx = _make_ctx()
        results = await list_docs(ctx)
        assert isinstance(results, list)
        assert len(results) == 2

    async def test_list_documents_updates_state(self, rag_db):
        from haiku.rag.skills.rag import RAGState, create_skill

        skill = create_skill(db_path=rag_db)
        list_docs = _get_tool(skill, "list_documents")
        state = RAGState()
        ctx = _make_ctx(state)
        await list_docs(ctx)
        assert len(state.documents) == 2
        assert isinstance(state.documents[0], DocumentInfo)
        assert state.documents[0].id is not None

    async def test_list_documents_no_duplicates_in_state(self, rag_db):
        from haiku.rag.skills.rag import RAGState, create_skill

        skill = create_skill(db_path=rag_db)
        list_docs = _get_tool(skill, "list_documents")
        state = RAGState()
        ctx = _make_ctx(state)
        await list_docs(ctx)
        await list_docs(ctx)
        assert len(state.documents) == 2


class TestGetDocumentTool:
    async def test_get_document_by_title(self, rag_db):
        from haiku.rag.skills.rag import create_skill

        skill = create_skill(db_path=rag_db)
        get_doc = _get_tool(skill, "get_document")
        ctx = _make_ctx()
        result = await get_doc(ctx, query="AI Overview")
        assert result is not None
        assert result["title"] == "AI Overview"

    async def test_get_document_updates_state(self, rag_db):
        from haiku.rag.skills.rag import RAGState, create_skill

        skill = create_skill(db_path=rag_db)
        get_doc = _get_tool(skill, "get_document")
        state = RAGState()
        ctx = _make_ctx(state)
        await get_doc(ctx, query="AI Overview")
        assert len(state.documents) == 1
        assert isinstance(state.documents[0], DocumentInfo)
        assert state.documents[0].title == "AI Overview"

    async def test_get_document_not_found(self, rag_db):
        from haiku.rag.skills.rag import create_skill

        skill = create_skill(db_path=rag_db)
        get_doc = _get_tool(skill, "get_document")
        ctx = _make_ctx()
        result = await get_doc(ctx, query="nonexistent document xyz")
        assert result is None


class TestAskTool:
    async def test_ask_returns_answer_with_citations(self, rag_db, monkeypatch):
        from haiku.rag.skills.rag import create_skill

        citations = [
            Citation(
                document_id="d1",
                chunk_id="c1",
                document_uri="test://ai-overview",
                document_title="AI Overview",
                content="AI is transforming industries.",
            )
        ]
        monkeypatch.setattr(
            HaikuRAG,
            "ask",
            AsyncMock(return_value=("AI transforms industries worldwide.", citations)),
        )

        skill = create_skill(db_path=rag_db)
        ask = _get_tool(skill, "ask")
        ctx = _make_ctx()
        result = await ask(ctx, question="What is AI?")
        assert isinstance(result, str)
        assert "AI transforms industries" in result

    async def test_ask_updates_state(self, rag_db, monkeypatch):
        from haiku.rag.skills.rag import RAGState, create_skill

        citations = [
            Citation(
                document_id="d1",
                chunk_id="c1",
                document_uri="test://ai-overview",
                content="AI content",
            )
        ]
        monkeypatch.setattr(
            HaikuRAG,
            "ask",
            AsyncMock(return_value=("AI transforms industries.", citations)),
        )

        skill = create_skill(db_path=rag_db)
        ask = _get_tool(skill, "ask")
        state = RAGState()
        ctx = _make_ctx(state)
        await ask(ctx, question="What is AI?")
        assert len(state.citations) == 1
        assert len(state.qa_history) == 1
        assert isinstance(state.qa_history[0], QAHistoryEntry)
        assert state.qa_history[0].question == "What is AI?"

    async def test_ask_assigns_citation_indices(self, rag_db, monkeypatch):
        from haiku.rag.skills.rag import RAGState, create_skill

        first_citations = [
            Citation(
                document_id="d1",
                chunk_id="c1",
                document_uri="test://doc1",
                content="First.",
            ),
            Citation(
                document_id="d2",
                chunk_id="c2",
                document_uri="test://doc2",
                content="Second.",
            ),
        ]
        second_citations = [
            Citation(
                document_id="d3",
                chunk_id="c3",
                document_uri="test://doc3",
                content="Third.",
            ),
        ]

        call_count = 0

        async def mock_ask(self, question, **kwargs):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                return ("Answer 1", first_citations)
            return ("Answer 2", second_citations)

        monkeypatch.setattr(HaikuRAG, "ask", mock_ask)

        skill = create_skill(db_path=rag_db)
        ask = _get_tool(skill, "ask")
        state = RAGState()
        ctx = _make_ctx(state)

        await ask(ctx, question="First question")
        assert state.citations[0].index == 1
        assert state.citations[1].index == 2

        await ask(ctx, question="Second question")
        assert state.citations[2].index == 3

    async def test_ask_applies_document_filter_from_state(self, rag_db, monkeypatch):
        from haiku.rag.skills.rag import RAGState, create_skill

        captured_kwargs = {}

        async def mock_ask(self, question, **kwargs):
            captured_kwargs.update(kwargs)
            return ("Answer.", [])

        monkeypatch.setattr(HaikuRAG, "ask", mock_ask)

        skill = create_skill(db_path=rag_db)
        ask = _get_tool(skill, "ask")
        state = RAGState(document_filter="title = 'AI Overview'")
        ctx = _make_ctx(state)
        await ask(ctx, question="What is AI?")
        assert captured_kwargs.get("filter") == "title = 'AI Overview'"

    async def test_ask_includes_prior_qa_context(self, rag_db, monkeypatch):
        import random

        from haiku.rag.skills.rag import RAGState, create_skill
        from tests.skills.conftest import VECTOR_DIM

        captured_questions = []

        async def mock_ask(self, question, **kwargs):
            captured_questions.append(question)
            return ("Answer about AI.", [])

        monkeypatch.setattr(HaikuRAG, "ask", mock_ask)

        skill = create_skill(db_path=rag_db)
        ask = _get_tool(skill, "ask")

        # Pre-compute the embedding the fake embedder will produce for "Tell me about AI"
        query_text = "Tell me about AI"
        random.seed(hash(query_text) % (2**32))
        query_embedding = [random.random() for _ in range(VECTOR_DIM)]

        prior_citations = [
            Citation(
                document_id="d1",
                chunk_id="c1",
                document_uri="test://ai-overview",
                document_title="AI Overview",
                content="AI content from source.",
            )
        ]
        state = RAGState(
            qa_history=[
                QAHistoryEntry(
                    question="What is artificial intelligence?",
                    answer="AI is the simulation of human intelligence by machines.",
                    question_embedding=query_embedding,
                    citations=prior_citations,
                ),
            ]
        )
        ctx = _make_ctx(state)
        await ask(ctx, question=query_text)

        # rag.ask() should receive augmented question with prior context
        assert len(captured_questions) == 1
        augmented = captured_questions[0]
        assert "Context from prior questions" in augmented
        assert "What is artificial intelligence?" in augmented
        assert "AI is the simulation" in augmented
        assert "AI Overview" in augmented
        assert query_text in augmented

        # State should store the original question, not the augmented one
        assert state.qa_history[-1].question == query_text

    async def test_ask_embeds_prior_qa_on_demand(self, rag_db, monkeypatch):
        from haiku.rag.skills.rag import RAGState, create_skill
        from tests.skills.conftest import VECTOR_DIM

        captured_questions = []

        async def mock_ask(self, question, **kwargs):
            captured_questions.append(question)
            return ("Answer about AI.", [])

        monkeypatch.setattr(HaikuRAG, "ask", mock_ask)

        skill = create_skill(db_path=rag_db)
        ask = _get_tool(skill, "ask")

        # Use the same question text for the prior QA entry and query so
        # their fake embeddings are identical (cosine similarity = 1.0).
        prior_question = "Tell me about AI"
        query_text = prior_question

        # Leave question_embedding=None to exercise the lazy embedding path
        state = RAGState(
            qa_history=[
                QAHistoryEntry(
                    question=prior_question,
                    answer="AI is the simulation of human intelligence by machines.",
                    question_embedding=None,
                ),
            ]
        )
        ctx = _make_ctx(state)
        await ask(ctx, question=query_text)

        # The lazy embedding should have populated question_embedding
        assert state.qa_history[0].question_embedding is not None
        assert len(state.qa_history[0].question_embedding) == VECTOR_DIM

        # The augmented question should include prior context
        assert len(captured_questions) == 1
        assert "Context from prior questions" in captured_questions[0]
        assert prior_question in captured_questions[0]

    async def test_ask_no_prior_qa_context_when_irrelevant(self, rag_db, monkeypatch):
        from haiku.rag.skills.rag import RAGState, create_skill
        from tests.skills.conftest import VECTOR_DIM

        captured_questions = []

        async def mock_ask(self, question, **kwargs):
            captured_questions.append(question)
            return ("Answer.", [])

        monkeypatch.setattr(HaikuRAG, "ask", mock_ask)

        skill = create_skill(db_path=rag_db)
        ask = _get_tool(skill, "ask")

        # Use orthogonal embedding — won't match the fake embedder's output
        orthogonal = [1.0 if i % 2 == 0 else -1.0 for i in range(VECTOR_DIM)]
        state = RAGState(
            qa_history=[
                QAHistoryEntry(
                    question="What is the weather?",
                    answer="It is sunny today.",
                    question_embedding=orthogonal,
                ),
            ]
        )
        ctx = _make_ctx(state)
        await ask(ctx, question="Explain quantum computing")

        # rag.ask() should receive the original question unchanged
        assert len(captured_questions) == 1
        assert captured_questions[0] == "Explain quantum computing"


class TestResearchTool:
    async def test_research_returns_report(self, rag_db, monkeypatch):
        from haiku.rag.skills.rag import create_skill

        report = ResearchReport(
            title="AI Research",
            executive_summary="AI is transforming industries.",
            main_findings=["Finding 1"],
            conclusions=["Conclusion 1"],
            sources_summary="Multiple sources consulted.",
        )
        monkeypatch.setattr(HaikuRAG, "research", AsyncMock(return_value=report))

        skill = create_skill(db_path=rag_db)
        research = _get_tool(skill, "research")
        ctx = _make_ctx()
        result = await research(ctx, question="What is AI?")
        assert isinstance(result, str)
        assert "AI Research" in result

    async def test_research_updates_state(self, rag_db, monkeypatch):
        from haiku.rag.skills.rag import RAGState, create_skill

        report = ResearchReport(
            title="AI Research",
            executive_summary="AI is transforming industries.",
            main_findings=["Finding 1"],
            conclusions=["Conclusion 1"],
            sources_summary="Multiple sources consulted.",
        )
        monkeypatch.setattr(HaikuRAG, "research", AsyncMock(return_value=report))

        skill = create_skill(db_path=rag_db)
        research = _get_tool(skill, "research")
        state = RAGState()
        ctx = _make_ctx(state)
        await research(ctx, question="What is AI?")
        assert len(state.reports) == 1
        assert state.reports[0].question == "What is AI?"
        assert len(state.qa_history) == 1
        assert state.qa_history[0].question == "What is AI?"
        assert state.qa_history[0].answer == "AI is transforming industries."

    async def test_research_applies_document_filter_from_state(
        self, rag_db, monkeypatch
    ):
        from haiku.rag.skills.rag import RAGState, create_skill

        captured_kwargs = {}

        report = ResearchReport(
            title="AI Research",
            executive_summary="Summary.",
            main_findings=["Finding"],
            conclusions=["Conclusion"],
            sources_summary="Sources.",
        )

        async def mock_research(self, question, **kwargs):
            captured_kwargs.update(kwargs)
            return report

        monkeypatch.setattr(HaikuRAG, "research", mock_research)

        skill = create_skill(db_path=rag_db)
        research = _get_tool(skill, "research")
        state = RAGState(document_filter="title = 'AI Overview'")
        ctx = _make_ctx(state)
        await research(ctx, question="What is AI?")
        assert captured_kwargs.get("filter") == "title = 'AI Overview'"

    async def test_research_without_state(self, rag_db, monkeypatch):
        from haiku.rag.skills.rag import create_skill

        report = ResearchReport(
            title="AI Research",
            executive_summary="Summary.",
            main_findings=["Finding"],
            conclusions=["Conclusion"],
            sources_summary="Sources.",
        )
        monkeypatch.setattr(HaikuRAG, "research", AsyncMock(return_value=report))

        skill = create_skill(db_path=rag_db)
        research = _get_tool(skill, "research")
        ctx = _make_ctx(state=None)
        result = await research(ctx, question="What is AI?")
        assert isinstance(result, str)