import random

import pytest
from docling_core.types.doc.document import ContentLayer, DoclingDocument
from docling_core.types.doc.labels import DocItemLabel

from haiku.rag.client import HaikuRAG
from haiku.rag.client.titles import extract_structural_title, resolve_title
from haiku.rag.config import AppConfig
from haiku.rag.config.models import ProcessingConfig
from haiku.rag.embeddings import EmbedderWrapper


@pytest.fixture(autouse=True)
def mock_embedder(monkeypatch):
    """Monkeypatch the embedder to return deterministic vectors."""

    async def fake_embed_query(self, text):
        random.seed(hash(text) % (2**32))
        return [random.random() for _ in range(2560)]

    async def fake_embed_documents(self, texts):
        result = []
        for t in texts:
            random.seed(hash(t) % (2**32))
            result.append([random.random() for _ in range(2560)])
        return result

    monkeypatch.setattr(EmbedderWrapper, "embed_query", fake_embed_query)
    monkeypatch.setattr(EmbedderWrapper, "embed_documents", fake_embed_documents)


# =========================================================================
# Structural title extraction
# =========================================================================


class TestExtractStructuralTitle:
    def test_furniture_title(self):
        """TITLE on FURNITURE layer (HTML <title>) is extracted."""
        doc = DoclingDocument(name="test")
        doc.add_text(
            label=DocItemLabel.TITLE,
            text="Website Page Title",
            content_layer=ContentLayer.FURNITURE,
        )
        doc.add_text(label=DocItemLabel.PARAGRAPH, text="Body text")

        assert extract_structural_title(doc) == "Website Page Title"

    def test_body_title(self):
        """TITLE on BODY layer (h1, PDF title) is extracted."""
        doc = DoclingDocument(name="test")
        doc.add_text(
            label=DocItemLabel.TITLE,
            text="Document Heading",
            content_layer=ContentLayer.BODY,
        )
        doc.add_text(label=DocItemLabel.PARAGRAPH, text="Body text")

        assert extract_structural_title(doc) == "Document Heading"

    def test_section_header_fallback(self):
        """First SECTION_HEADER is used when no TITLE exists."""
        doc = DoclingDocument(name="test")
        doc.add_text(label=DocItemLabel.SECTION_HEADER, text="Introduction")
        doc.add_text(label=DocItemLabel.SECTION_HEADER, text="Background")
        doc.add_text(label=DocItemLabel.PARAGRAPH, text="Body text")

        assert extract_structural_title(doc) == "Introduction"

    def test_no_title_or_headers(self):
        """Returns None when no TITLE or SECTION_HEADER exists."""
        doc = DoclingDocument(name="test")
        doc.add_text(label=DocItemLabel.PARAGRAPH, text="Just a paragraph")

        assert extract_structural_title(doc) is None

    def test_furniture_title_preferred_over_body_title(self):
        """FURNITURE TITLE takes priority over BODY TITLE."""
        doc = DoclingDocument(name="test")
        doc.add_text(
            label=DocItemLabel.TITLE,
            text="Body H1 Title",
            content_layer=ContentLayer.BODY,
        )
        doc.add_text(
            label=DocItemLabel.TITLE,
            text="HTML Page Title",
            content_layer=ContentLayer.FURNITURE,
        )

        assert extract_structural_title(doc) == "HTML Page Title"

    def test_whitespace_stripped(self):
        """Whitespace is stripped from extracted titles."""
        doc = DoclingDocument(name="test")
        doc.add_text(
            label=DocItemLabel.TITLE,
            text="  Padded Title  ",
            content_layer=ContentLayer.BODY,
        )

        assert extract_structural_title(doc) == "Padded Title"

    def test_empty_title_text_skipped(self):
        """Empty or whitespace-only TITLE text is skipped."""
        doc = DoclingDocument(name="test")
        doc.add_text(
            label=DocItemLabel.TITLE,
            text="   ",
            content_layer=ContentLayer.BODY,
        )
        doc.add_text(label=DocItemLabel.SECTION_HEADER, text="Actual Heading")

        assert extract_structural_title(doc) == "Actual Heading"


# =========================================================================
# resolve_title
# =========================================================================


class TestResolveTitle:
    @pytest.mark.asyncio
    async def test_auto_title_disabled_returns_none(self):
        """When auto_title is False, returns None (no title generation)."""
        doc = DoclingDocument(name="test")
        doc.add_text(label=DocItemLabel.TITLE, text="Structural Title")

        config = AppConfig(processing=ProcessingConfig(auto_title=False))
        result = await resolve_title(config, doc, "some content")
        assert result is None

    @pytest.mark.asyncio
    async def test_structural_title_extracted(self):
        """Structural title is extracted when auto_title is enabled."""
        doc = DoclingDocument(name="test")
        doc.add_text(label=DocItemLabel.TITLE, text="Auto Extracted Title")

        config = AppConfig(processing=ProcessingConfig(auto_title=True))
        result = await resolve_title(config, doc, "some content")
        assert result == "Auto Extracted Title"

    @pytest.mark.asyncio
    async def test_llm_failure_returns_none(self, monkeypatch):
        """LLM failure during ingestion returns None instead of raising."""
        doc = DoclingDocument(name="test")
        doc.add_text(label=DocItemLabel.PARAGRAPH, text="Just text")

        async def exploding_llm(config, content):
            raise RuntimeError("LLM is down")

        monkeypatch.setattr(
            "haiku.rag.client.titles.generate_title_with_llm", exploding_llm
        )
        config = AppConfig(processing=ProcessingConfig(auto_title=True))
        result = await resolve_title(config, doc, "some content")
        assert result is None


# =========================================================================
# Integration: create_document with auto_title
# =========================================================================


class TestCreateDocumentAutoTitle:
    @pytest.mark.asyncio
    async def test_auto_title_from_structural(self, temp_db_path):
        """create_document with auto_title=True extracts title from docling."""
        config = AppConfig(processing=ProcessingConfig(auto_title=True))
        async with HaikuRAG(temp_db_path, config=config, create=True) as client:
            doc = await client.create_document(
                "# My Document\n\nSome content here.", uri="test://auto-title"
            )
            assert doc.title == "My Document"

    @pytest.mark.asyncio
    async def test_auto_title_disabled(self, temp_db_path):
        """create_document with auto_title=False leaves title as None."""
        config = AppConfig(processing=ProcessingConfig(auto_title=False))
        async with HaikuRAG(temp_db_path, config=config, create=True) as client:
            doc = await client.create_document(
                "# My Document\n\nSome content here.", uri="test://no-auto-title"
            )
            assert doc.title is None

    @pytest.mark.asyncio
    async def test_explicit_title_not_overridden(self, temp_db_path):
        """Explicit title is never overridden by auto-generation."""
        config = AppConfig(processing=ProcessingConfig(auto_title=True))
        async with HaikuRAG(temp_db_path, config=config, create=True) as client:
            doc = await client.create_document(
                "# Auto Title\n\nSome content here.",
                uri="test://explicit-title",
                title="My Explicit Title",
            )
            assert doc.title == "My Explicit Title"


# =========================================================================
# Integration: import_document with auto_title
# =========================================================================


class TestImportDocumentAutoTitle:
    @pytest.mark.asyncio
    async def test_auto_title_from_structural(self, temp_db_path):
        """import_document with auto_title=True extracts title from docling."""
        config = AppConfig(processing=ProcessingConfig(auto_title=True))
        async with HaikuRAG(temp_db_path, config=config, create=True) as client:
            docling_doc = await client.convert("# Imported Doc\n\nContent here.")
            chunks = await client.chunk(docling_doc)
            doc = await client.import_document(
                docling_doc, chunks, uri="test://import-auto-title"
            )
            assert doc.title == "Imported Doc"

    @pytest.mark.asyncio
    async def test_explicit_title_preserved(self, temp_db_path):
        """import_document explicit title is not overridden."""
        config = AppConfig(processing=ProcessingConfig(auto_title=True))
        async with HaikuRAG(temp_db_path, config=config, create=True) as client:
            docling_doc = await client.convert("# Auto Title\n\nContent here.")
            chunks = await client.chunk(docling_doc)
            doc = await client.import_document(
                docling_doc,
                chunks,
                uri="test://import-explicit",
                title="Keep This Title",
            )
            assert doc.title == "Keep This Title"


# =========================================================================
# generate_title() public method
# =========================================================================


class TestGenerateTitle:
    @pytest.mark.asyncio
    async def test_structural_title(self, temp_db_path):
        """generate_title extracts structural title from document."""
        config = AppConfig(processing=ProcessingConfig(auto_title=True))
        async with HaikuRAG(temp_db_path, config=config, create=True) as client:
            doc = await client.create_document(
                "# Great Heading\n\nSome content.", uri="test://gen-title"
            )
            title = await client.generate_title(doc)
            assert title == "Great Heading"

    @pytest.mark.asyncio
    async def test_no_structural_title_no_llm(self, temp_db_path):
        """generate_title raises when no structural title and LLM unavailable."""
        async with HaikuRAG(temp_db_path, create=True) as client:
            doc = await client.create_document(
                "Just plain text without headings.",
                uri="test://gen-no-title",
                format="plain",
            )
            with pytest.raises(RuntimeError):
                await client.generate_title(doc)

    @pytest.mark.asyncio
    async def test_bypasses_auto_title_config(self, temp_db_path):
        """generate_title works even when auto_title is False."""
        config = AppConfig(processing=ProcessingConfig(auto_title=False))
        async with HaikuRAG(temp_db_path, config=config, create=True) as client:
            doc = await client.create_document(
                "# Heading Present\n\nBody text.", uri="test://gen-bypass"
            )
            title = await client.generate_title(doc)
            assert title == "Heading Present"


# =========================================================================
# rebuild --title-only
# =========================================================================


class TestRebuildTitleOnly:
    @pytest.mark.asyncio
    async def test_generates_titles_for_untitled_docs(self, temp_db_path):
        """TITLE_ONLY mode generates titles for documents without one."""
        from haiku.rag.client import RebuildMode

        config = AppConfig(processing=ProcessingConfig(auto_title=True))
        async with HaikuRAG(temp_db_path, config=config, create=True) as client:
            doc1 = await client.create_document(
                "# Doc With Heading\n\nContent.",
                uri="test://titled",
            )
            assert doc1.title == "Doc With Heading"

            # Simulate an untitled doc
            doc1.title = None
            await client.document_repository.update(doc1)

            doc2 = await client.create_document(
                "# Another Heading\n\nMore content.",
                uri="test://also-titled",
            )
            assert doc2.title == "Another Heading"

            processed_ids = []
            async for doc_id in client.rebuild_database(mode=RebuildMode.TITLE_ONLY):
                processed_ids.append(doc_id)

            assert len(processed_ids) == 1
            assert doc1.id in processed_ids

            updated_doc1 = await client.get_document_by_id(doc1.id)
            assert updated_doc1 is not None
            assert updated_doc1.title == "Doc With Heading"

    @pytest.mark.asyncio
    async def test_skips_already_titled_docs(self, temp_db_path):
        """TITLE_ONLY mode skips documents that already have titles."""
        from haiku.rag.client import RebuildMode

        config = AppConfig(processing=ProcessingConfig(auto_title=True))
        async with HaikuRAG(temp_db_path, config=config, create=True) as client:
            await client.create_document(
                "# Has Title\n\nContent.",
                uri="test://has-title",
            )

            processed_ids = []
            async for doc_id in client.rebuild_database(mode=RebuildMode.TITLE_ONLY):
                processed_ids.append(doc_id)

            assert len(processed_ids) == 0

    @pytest.mark.asyncio
    async def test_continues_on_per_document_failure(self, temp_db_path, monkeypatch):
        """TITLE_ONLY mode continues when generate_title fails for a document."""
        from haiku.rag.client import RebuildMode

        config = AppConfig(processing=ProcessingConfig(auto_title=True))
        async with HaikuRAG(temp_db_path, config=config, create=True) as client:
            doc1 = await client.create_document(
                "# First Heading\n\nContent.",
                uri="test://first",
            )
            doc2 = await client.create_document(
                "# Second Heading\n\nContent.",
                uri="test://second",
            )
            # Clear both titles
            doc1.title = None
            await client.document_repository.update(doc1)
            doc2.title = None
            await client.document_repository.update(doc2)

            # Make generate_title fail for the first doc, succeed for the second
            original = HaikuRAG.generate_title
            call_count = 0

            async def flaky_generate(self, document):
                nonlocal call_count
                call_count += 1
                if call_count == 1:
                    raise RuntimeError("LLM failed")
                return await original(self, document)

            monkeypatch.setattr(HaikuRAG, "generate_title", flaky_generate)

            processed_ids = []
            async for doc_id in client.rebuild_database(mode=RebuildMode.TITLE_ONLY):
                processed_ids.append(doc_id)

            # Only the second doc should have been processed
            assert len(processed_ids) == 1


@pytest.mark.asyncio
async def test_generate_title_with_llm_returns_model_output(monkeypatch):
    """The agent's output is stripped and returned."""
    from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart
    from pydantic_ai.models.function import AgentInfo, FunctionModel

    from haiku.rag.client.titles import generate_title_with_llm
    from haiku.rag.config import AppConfig

    def respond(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        return ModelResponse(parts=[TextPart("  A Generated Title  ")])

    monkeypatch.setattr(
        "haiku.rag.utils.get_model", lambda *a, **kw: FunctionModel(respond)
    )

    assert await generate_title_with_llm(AppConfig(), "body") == "A Generated Title"


@pytest.mark.asyncio
async def test_generate_title_with_llm_returns_none_for_blank_output(monkeypatch):
    from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart
    from pydantic_ai.models.function import AgentInfo, FunctionModel

    from haiku.rag.client.titles import generate_title_with_llm
    from haiku.rag.config import AppConfig

    def respond(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        return ModelResponse(parts=[TextPart("   ")])

    monkeypatch.setattr(
        "haiku.rag.utils.get_model", lambda *a, **kw: FunctionModel(respond)
    )

    assert await generate_title_with_llm(AppConfig(), "body") is None