app/tests/test_generation_service.py

"""Tests for app.services.generation — all three AI modes.

The full stack (DB, retriever, reranker, LLM adapter, style cache) is mocked
so these tests run without any network calls or a real database write path.

Covers:
- _generate_section_text: RAG pipeline returns (text, provenance, confidence, top_results).
- _run_generate: persists section with correct mode + style profile metadata.
- _run_proofread: strips [Editor notes: ...] before passing text to the LLM.
- _run_enhance: strips [Editor notes: ...] before passing text to the LLM.
- _run_proofread / _run_enhance fallback: no existing text triggers seed generation.
- run_generation dispatch: correct sub-pipeline called for each mode string.
"""

import json
import uuid
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
import pytest_asyncio
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine

from app.db.database import Base
from app.db.models import Document, Report, ReportSection, ReportStatus
from app.generator.adapter import MockLLMAdapter
from app.models.schemas import GenerationMode, WritingStyleProfile
from app.services.generation import (
    _generate_section_text,
    _run_enhance,
    _run_generate,
    _run_proofread,
    run_generation,
)

# ── Shared fixtures ───────────────────────────────────────────────────────────

_MOCK_PROFILE = WritingStyleProfile(
    tone="formal",
    formality_level="professional",
    avg_sentence_complexity="moderate",
    vocabulary_level="technical",
    common_phrases=[],
    structural_patterns=[],
    writing_style_summary="Professional RICS surveyor style.",
)


class _FakeResult:
    """Minimal stand-in for a retrieval/rerank result object."""
    def __init__(self, idx: int = 0) -> None:
        self.text = f"Retrieved snippet {idx}."
        self.doc_id = f"doc-{idx}"
        self.chunk_id = f"chunk-{idx}"
        self.score = 0.9
        self.rerank_score = 0.85


@pytest_asyncio.fixture
async def test_db_session() -> AsyncSession:
    """Isolated in-memory SQLite session (same pattern as conftest.test_db)."""
    engine = create_async_engine("sqlite+aiosqlite:///:memory:", echo=False)
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    factory = async_sessionmaker(engine, expire_on_commit=False)
    async with factory() as session:
        yield session
    await engine.dispose()


@pytest_asyncio.fixture
async def seeded_report(test_db_session: AsyncSession) -> Report:
    """Persist a Document + Report row and return the Report ORM object."""
    doc = Document(
        id=str(uuid.uuid4()),
        tenant_id="tenant_test",
        filename="test.pdf",
        file_path="/tmp/test.pdf",
    )
    report = Report(
        id=str(uuid.uuid4()),
        tenant_id="tenant_test",
        document_id=doc.id,
        status=ReportStatus.generating,
    )
    test_db_session.add(doc)
    test_db_session.add(report)
    await test_db_session.commit()
    return report


# ── _generate_section_text ────────────────────────────────────────────────────

@pytest.mark.asyncio
async def test_generate_section_text_returns_text_provenance_confidence(
    seeded_report: Report,
) -> None:
    """_generate_section_text must return a non-empty 4-tuple."""
    fake_results = [_FakeResult(i) for i in range(2)]

    with (
        patch("app.services.generation.settings.hierarchical_rag_enabled", True),
        patch("app.services.generation.get_vectorstore", return_value=MagicMock(search=MagicMock(return_value=fake_results))),
        patch("app.services.generation.retrieve_document_level_context", return_value=[]),
        patch("app.services.generation.retrieve_for_report", return_value=fake_results),
        patch("app.services.generation.rerank", return_value=fake_results),
        patch("app.llm.generation_facade.get_llm_adapter", return_value=MockLLMAdapter()),
        patch("app.services.generation.get_template", return_value=MagicMock(skeleton="[D]: [content].")),
    ):
        text, provenance, confidence, top_results, doc_ctx, metrics = await _generate_section_text(
            tenant_id="tenant_test",
            template_id="D",
            bullets=["Semi-detached, NW3"],
            style_profile=_MOCK_PROFILE,
            retrieval_level="paragraph",
        )

    assert isinstance(text, str) and text
    assert isinstance(provenance, list)
    assert isinstance(confidence, float)
    assert isinstance(top_results, list)
    assert doc_ctx == []
    assert isinstance(metrics, dict)


# ── _run_generate ─────────────────────────────────────────────────────────────

@pytest.mark.asyncio
async def test_run_generate_persists_section_and_marks_complete(
    test_db_session: AsyncSession,
    seeded_report: Report,
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """_run_generate should write a ReportSection row and set report.status=complete."""
    monkeypatch.setattr("app.cache.section_cache.settings.cache_dir", tmp_path)
    monkeypatch.setattr("app.config.settings.cache_dir", tmp_path)

    fake_results = [_FakeResult(0)]

    with (
        patch("app.services.generation.retrieve_document_level_context", return_value=[]),
        patch("app.services.generation.retrieve_for_report", return_value=fake_results),
        patch("app.services.generation.rerank", return_value=fake_results),
        patch("app.llm.generation_facade.get_llm_adapter", return_value=MockLLMAdapter()),
        patch("app.services.generation.get_template", return_value=MagicMock(skeleton="[E4]: [content].")),
        patch("app.services.generation._get_or_build_style_profile", new=AsyncMock(return_value=_MOCK_PROFILE)),
    ):
        await _run_generate(
            db=test_db_session,
            report=seeded_report,
            tenant_id="tenant_test",
            template_id="E4",
            bullets=["Solid brick walls 275mm, slight crack NW corner"],
            force_regenerate=False,
        )

    await test_db_session.refresh(seeded_report)
    assert seeded_report.status == ReportStatus.complete


@pytest.mark.asyncio
async def test_run_generate_stores_mode_in_provenance(
    test_db_session: AsyncSession,
    seeded_report: Report,
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """The persisted provenance JSON must include mode='generate' in meta."""
    monkeypatch.setattr("app.cache.section_cache.settings.cache_dir", tmp_path)
    monkeypatch.setattr("app.config.settings.cache_dir", tmp_path)

    fake_results = [_FakeResult(0)]

    with (
        patch("app.services.generation.retrieve_document_level_context", return_value=[]),
        patch("app.services.generation.retrieve_for_report", return_value=fake_results),
        patch("app.services.generation.rerank", return_value=fake_results),
        patch("app.llm.generation_facade.get_llm_adapter", return_value=MockLLMAdapter()),
        patch("app.services.generation.get_template", return_value=MagicMock(skeleton="[E2]: [content].")),
        patch("app.services.generation._get_or_build_style_profile", new=AsyncMock(return_value=_MOCK_PROFILE)),
    ):
        await _run_generate(
            db=test_db_session,
            report=seeded_report,
            tenant_id="tenant_test",
            template_id="E2",
            bullets=["Main roof hip slate, one slipped tile"],
            force_regenerate=False,
        )

    from sqlalchemy import select
    result = await test_db_session.execute(
        select(ReportSection).where(ReportSection.report_id == seeded_report.id)
    )
    section = result.scalars().first()
    assert section is not None
    raw = json.loads(section.provenance)
    assert raw["meta"]["mode"] == "generate"
    assert "style_profile" in raw["meta"]


# ── _run_proofread — notes stripping ──────────────────────────────────────────

@pytest.mark.asyncio
async def test_run_proofread_strips_editor_notes_before_llm(
    test_db_session: AsyncSession,
    seeded_report: Report,
) -> None:
    """_run_proofread must send only the clean body text to the adapter, not old notes."""
    clean_body = "The property is a semi-detached house in good condition."
    dirty_text = clean_body + "\n\n[Editor notes: Fix grammar on line 2. Improve clarity.]"

    # Pre-persist a section with notes already appended
    existing_section = ReportSection(
        id=str(uuid.uuid4()),
        report_id=seeded_report.id,
        section_code="E4",
        text=dirty_text,
        confidence=0.8,
    )
    test_db_session.add(existing_section)
    await test_db_session.commit()

    received_texts: list[str] = []

    class _CapturingAdapter(MockLLMAdapter):
        def proofread(
            self,
            text: str,
            bullets: list[str],
            style_profile=None,
            temperature: float = 0.15,
            creativity_hint: str = "",
        ) -> str:
            received_texts.append(text)
            return text + "\n---NOTES---\nNo changes needed."

    with patch("app.llm.generation_facade.get_llm_adapter", return_value=_CapturingAdapter()), \
         patch("app.services.generation._get_or_build_style_profile", new=AsyncMock(return_value=_MOCK_PROFILE)):
        await _run_proofread(
            db=test_db_session,
            report=seeded_report,
            tenant_id="tenant_test",
            template_id="E4",
            bullets=["Solid brick walls 275mm, DPC visible"],
        )

    assert len(received_texts) == 1
    # The LLM must have received the clean body without the editor notes block
    assert "[Editor notes:" not in received_texts[0]
    assert clean_body in received_texts[0]


@pytest.mark.asyncio
async def test_run_proofread_persists_mode_proofread(
    test_db_session: AsyncSession,
    seeded_report: Report,
) -> None:
    """Provenance meta must record mode='proofread' after a proofread run."""
    existing_section = ReportSection(
        id=str(uuid.uuid4()),
        report_id=seeded_report.id,
        section_code="E4",
        text="Clean original text.",
        confidence=0.8,
    )
    test_db_session.add(existing_section)
    await test_db_session.commit()

    with patch("app.llm.generation_facade.get_llm_adapter", return_value=MockLLMAdapter()), \
         patch("app.services.generation._get_or_build_style_profile", new=AsyncMock(return_value=_MOCK_PROFILE)):
        await _run_proofread(
            db=test_db_session,
            report=seeded_report,
            tenant_id="tenant_test",
            template_id="E4",
            bullets=["Solid brick walls, DPC visible"],
        )

    from sqlalchemy import select
    result = await test_db_session.execute(
        select(ReportSection).where(
            ReportSection.report_id == seeded_report.id,
            ReportSection.section_code == "E4",
        )
    )
    section = result.scalars().first()
    raw = json.loads(section.provenance)
    assert raw["meta"]["mode"] == "proofread"


@pytest.mark.asyncio
async def test_run_proofread_generates_seed_when_no_existing_text(
    test_db_session: AsyncSession,
    seeded_report: Report,
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """If no existing section exists, proofread must generate seed text first."""
    monkeypatch.setattr("app.cache.section_cache.settings.cache_dir", tmp_path)
    monkeypatch.setattr("app.config.settings.cache_dir", tmp_path)

    fake_results = [_FakeResult(0)]

    with (
        patch("app.services.generation.retrieve_document_level_context", return_value=[]),
        patch("app.services.generation.retrieve_for_report", return_value=fake_results),
        patch("app.services.generation.rerank", return_value=fake_results),
        patch("app.llm.generation_facade.get_llm_adapter", return_value=MockLLMAdapter()),
        patch("app.services.generation.get_template", return_value=MagicMock(skeleton="[E2]: [content].")),
        patch("app.services.generation._get_or_build_style_profile", new=AsyncMock(return_value=_MOCK_PROFILE)),
    ):
        # No pre-existing section — the function must not raise
        await _run_proofread(
            db=test_db_session,
            report=seeded_report,
            tenant_id="tenant_test",
            template_id="E2",
            bullets=["Main roof hip slate, some slipped tiles"],
        )

    await test_db_session.refresh(seeded_report)
    assert seeded_report.status == ReportStatus.complete


# ── _run_enhance — notes stripping ────────────────────────────────────────────

@pytest.mark.asyncio
async def test_run_enhance_strips_editor_notes_before_llm(
    test_db_session: AsyncSession,
    seeded_report: Report,
) -> None:
    """_run_enhance must strip [Editor notes: ...] before passing text to the adapter."""
    clean_body = "The roof covering is of felt tiles."
    dirty_text = clean_body + "\n\n[Editor notes: Improve sentence structure.]"

    existing_section = ReportSection(
        id=str(uuid.uuid4()),
        report_id=seeded_report.id,
        section_code="E2",
        text=dirty_text,
        confidence=0.8,
    )
    test_db_session.add(existing_section)
    await test_db_session.commit()

    received_texts: list[str] = []

    class _CapturingAdapter(MockLLMAdapter):
        def enhance(
            self,
            text: str,
            bullets: list[str],
            snippets: list[str],
            style_profile=None,
            temperature: float = 0.2,
            creativity_hint: str = "",
        ) -> str:
            received_texts.append(text)
            return text + " [Enhanced.]"

    fake_results = [_FakeResult(0)]

    with (
        patch("app.services.generation.retrieve_for_report", return_value=fake_results),
        patch("app.services.generation.rerank", return_value=fake_results),
        patch("app.llm.generation_facade.get_llm_adapter", return_value=_CapturingAdapter()),
        patch("app.services.generation._get_or_build_style_profile", new=AsyncMock(return_value=_MOCK_PROFILE)),
    ):
        await _run_enhance(
            db=test_db_session,
            report=seeded_report,
            tenant_id="tenant_test",
            template_id="E2",
            bullets=["Roof: hip slate, some slipped tiles"],
            force_regenerate=False,
        )

    assert len(received_texts) == 1
    assert "[Editor notes:" not in received_texts[0]
    assert clean_body in received_texts[0]


@pytest.mark.asyncio
async def test_run_enhance_persists_mode_enhance(
    test_db_session: AsyncSession,
    seeded_report: Report,
) -> None:
    """Provenance meta must record mode='enhance' after an enhance run."""
    existing_section = ReportSection(
        id=str(uuid.uuid4()),
        report_id=seeded_report.id,
        section_code="F3",
        text="Existing text.",
        confidence=0.8,
    )
    test_db_session.add(existing_section)
    await test_db_session.commit()

    fake_results = [_FakeResult(0)]

    with (
        patch("app.services.generation.retrieve_for_report", return_value=fake_results),
        patch("app.services.generation.rerank", return_value=fake_results),
        patch("app.llm.generation_facade.get_llm_adapter", return_value=MockLLMAdapter()),
        patch("app.services.generation._get_or_build_style_profile", new=AsyncMock(return_value=_MOCK_PROFILE)),
    ):
        await _run_enhance(
            db=test_db_session,
            report=seeded_report,
            tenant_id="tenant_test",
            template_id="F3",
            bullets=["Internal walls plasterboard, some damp to NW corner"],
            force_regenerate=False,
        )

    from sqlalchemy import select
    result = await test_db_session.execute(
        select(ReportSection).where(
            ReportSection.report_id == seeded_report.id,
            ReportSection.section_code == "F3",
        )
    )
    section = result.scalars().first()
    raw = json.loads(section.provenance)
    assert raw["meta"]["mode"] == "enhance"


# ── run_generation dispatch ───────────────────────────────────────────────────

@pytest.mark.asyncio
async def test_run_generation_dispatches_to_proofread(
    seeded_report: Report,
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """run_generation with mode='proofread' must call _run_proofread, not _run_generate."""
    monkeypatch.setattr("app.cache.section_cache.settings.cache_dir", tmp_path)
    monkeypatch.setattr("app.config.settings.cache_dir", tmp_path)

    called = {"proofread": False, "generate": False, "enhance": False}

    async def _mock_proofread(*args, **kwargs) -> None:  # type: ignore[override]
        called["proofread"] = True

    async def _mock_generate(*args, **kwargs) -> None:  # type: ignore[override]
        called["generate"] = True

    with (
        patch("app.services.generation._run_proofread", side_effect=_mock_proofread),
        patch("app.services.generation._run_generate", side_effect=_mock_generate),
        patch(
            "app.services.generation.mark_report_generation_complete_if_still_generating",
            new_callable=AsyncMock,
        ),
        patch(
            "app.services.generation.finalize_generation_status_if_stuck",
            new_callable=AsyncMock,
        ),
        patch(
            "app.services.generation.abort_generation_if_report_invalid",
            new_callable=AsyncMock,
            return_value=True,
        ),
        patch("app.services.generation.get_session_factory") as mock_factory,
    ):
        # Build a minimal async context-manager that yields a real-enough session mock
        mock_session = MagicMock()
        mock_session.__aenter__ = AsyncMock(return_value=mock_session)
        mock_session.__aexit__ = AsyncMock(return_value=False)

        # Fake report lookup
        mock_report = MagicMock()
        mock_report.status = ReportStatus.generating
        mock_report.tenant_id = "tenant_test"
        mock_session.get = AsyncMock(return_value=mock_report)
        mock_session.commit = AsyncMock()

        mock_ctx = MagicMock()
        mock_ctx.__aenter__ = AsyncMock(return_value=mock_session)
        mock_ctx.__aexit__ = AsyncMock(return_value=False)
        mock_factory.return_value.return_value = mock_ctx

        await run_generation(
            report_id=seeded_report.id,
            tenant_id="tenant_test",
            template_id="E4",
            bullets=["Solid brick walls, DPC visible"],
            mode=GenerationMode.proofread,
        )

    assert called["proofread"] is True
    assert called["generate"] is False


@pytest.mark.asyncio
async def test_run_generation_dispatches_to_enhance(
    seeded_report: Report,
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """run_generation with mode='enhance' must call _run_enhance, not _run_generate."""
    monkeypatch.setattr("app.cache.section_cache.settings.cache_dir", tmp_path)
    monkeypatch.setattr("app.config.settings.cache_dir", tmp_path)

    called = {"enhance": False, "generate": False}

    async def _mock_enhance(*args, **kwargs) -> None:  # type: ignore[override]
        called["enhance"] = True

    async def _mock_generate(*args, **kwargs) -> None:  # type: ignore[override]
        called["generate"] = True

    with (
        patch("app.services.generation._run_enhance", side_effect=_mock_enhance),
        patch("app.services.generation._run_generate", side_effect=_mock_generate),
        patch(
            "app.services.generation.mark_report_generation_complete_if_still_generating",
            new_callable=AsyncMock,
        ),
        patch(
            "app.services.generation.finalize_generation_status_if_stuck",
            new_callable=AsyncMock,
        ),
        patch("app.services.generation.get_session_factory") as mock_factory,
    ):
        mock_session = MagicMock()
        mock_session.__aenter__ = AsyncMock(return_value=mock_session)
        mock_session.__aexit__ = AsyncMock(return_value=False)
        mock_report = MagicMock()
        mock_report.status = ReportStatus.generating
        mock_report.tenant_id = "tenant_test"
        mock_session.get = AsyncMock(return_value=mock_report)
        mock_session.commit = AsyncMock()
        mock_ctx = MagicMock()
        mock_ctx.__aenter__ = AsyncMock(return_value=mock_session)
        mock_ctx.__aexit__ = AsyncMock(return_value=False)
        mock_factory.return_value.return_value = mock_ctx

        await run_generation(
            report_id=seeded_report.id,
            tenant_id="tenant_test",
            template_id="F3",
            bullets=["Internal walls plasterboard, some damp NW corner"],
            mode=GenerationMode.enhance,
        )

    assert called["enhance"] is True
    assert called["generate"] is False