backend/tests/test_paragraph_matching.py

"""Tests that messy notes retrieve and map to the correct MASTER paragraph."""

from __future__ import annotations

from docx import Document

from backend.config import settings
from backend.core import ingest, section_mapper, template_discoverer
from backend.core.paragraph_retriever import build_retrieval_query, retrieve_master_paragraphs
from backend.core.rag_store import TIER_MASTER, Chunk, get_rag_store


def _write_cross_id_bundle(tmp_path):
    """Canonical-aligned template + standard paragraphs (D1/D2)."""
    tpl = tmp_path / "template.docx"
    doc = Document()
    doc.add_paragraph("D1 Chimney stacks")
    doc.add_paragraph("D2 Roof coverings")
    doc.save(str(tpl))

    paras = tmp_path / "paragraphs.docx"
    doc2 = Document()
    doc2.add_paragraph("D1\nChimney Stacks")
    doc2.add_paragraph(
        "The chimney stacks were inspected from ground level using binoculars."
    )
    doc2.add_paragraph("D2\nRoof Coverings")
    doc2.add_paragraph(
        "The roof covering comprises slate laid to pitched timber rafters."
    )
    doc2.save(str(paras))
    return tpl, paras


def test_build_retrieval_query_uses_observations():
    q = build_retrieval_query(
        "Chimney Stacks",
        ["stack cracked", "mortar open joints at crown"],
        section_id="D1",
    )
    assert "D1 Chimney Stacks:" in q
    assert "mortar" in q


def test_chimney_notes_retrieve_d1_not_d2(tmp_path, monkeypatch):
    monkeypatch.setattr(settings, "data_dir", str(tmp_path / "data"))
    tpl, paras = _write_cross_id_bundle(tmp_path)
    tenant = "t-match"
    ingest.ingest_report_template(tenant, tpl)
    ingest.ingest_standard_paragraphs(tenant, paras)

    schema = template_discoverer.load_schema(tenant)
    assert schema.section_alias_map.get("D1") == "D1"

    hits = retrieve_master_paragraphs(
        tenant,
        section_label="Chimney stacks",
        paragraph_section_id=schema.paragraph_section_id("D1"),
        observations=["stack cracked mortar missing"],
    )
    assert hits
    assert hits[0].section_id == "D1"
    assert "chimney" in hits[0].text.lower()
    assert "slate" not in hits[0].text.lower()


def test_roof_notes_retrieve_d2_not_d1(tmp_path, monkeypatch):
    monkeypatch.setattr(settings, "data_dir", str(tmp_path / "data"))
    tpl, paras = _write_cross_id_bundle(tmp_path)
    tenant = "t-roof"
    ingest.ingest_report_template(tenant, tpl)
    ingest.ingest_standard_paragraphs(tenant, paras)

    schema = template_discoverer.load_schema(tenant)
    hits = retrieve_master_paragraphs(
        tenant,
        section_label="Roof coverings",
        paragraph_section_id=schema.paragraph_section_id("D2"),
        observations=["slate tile slipped on south slope"],
    )
    assert hits
    assert hits[0].section_id == "D2"
    assert "slate" in hits[0].text.lower()
    assert "chimney" not in hits[0].text.lower()


def test_section_strict_excludes_other_sections():
    store = get_rag_store()
    store.clear_tier("t-strict", TIER_MASTER)
    store.ingest_document(
        "t-strict",
        "p",
        [
            Chunk(text="Chimney stack mortar decay noted.", section_id="D1", tier=TIER_MASTER),
            Chunk(text="Roof slate covering slipped tiles.", section_id="D2", tier=TIER_MASTER),
        ],
        tier=TIER_MASTER,
    )
    hits = store.search_for_generation(
        "t-strict",
        "slate roof tile slipped",
        section_id="D2",
        top_k=3,
        section_strict=True,
    )
    assert hits
    assert all(h.section_id == "D2" for h in hits)


def test_orphan_note_routed_by_rag_topic(tmp_path, monkeypatch):
    """Notes with no section header still match a MASTER paragraph by topic."""
    monkeypatch.setattr(settings, "data_dir", str(tmp_path / "data"))
    tpl, paras = _write_cross_id_bundle(tmp_path)
    tenant = "t-orphan"
    ingest.ingest_report_template(tenant, tpl)
    ingest.ingest_standard_paragraphs(tenant, paras)
    ingest.ingest_reference(tenant, paras)

    result = section_mapper.generate_report(
        tenant,
        "slate tile slipped on the south slope of the roof",
        interference_level="maximum",
    )
    d3 = next(s for s in result.sections if s.section_id == "D2")
    assert d3.status == "OK"
    assert "slate" in d3.text.lower()
    assert "south" in d3.text.lower()


def test_end_to_end_d1_notes_use_d1_rag_source(tmp_path, monkeypatch):
    monkeypatch.setattr(settings, "data_dir", str(tmp_path / "data"))
    tpl, paras = _write_cross_id_bundle(tmp_path)
    tenant = "t-e2e"
    ingest.ingest_report_template(tenant, tpl)
    ingest.ingest_standard_paragraphs(tenant, paras)
    ingest.ingest_reference(tenant, paras)

    result = section_mapper.generate_report(
        tenant,
        "D1: stack cracked mortar open\n\nD2: slate slipped south slope",
        interference_level="maximum",
    )
    d1 = next(s for s in result.sections if s.section_id == "D1")
    d3 = next(s for s in result.sections if s.section_id == "D2")

    assert d1.status != "NO_RAG_MATCH"
    assert d3.status != "NO_RAG_MATCH"
    assert d1.reference_sources or d1.rag_sources
    assert d3.reference_sources or d3.rag_sources
    assert any(
        "paragraphs.docx" in src.report_filename for src in d1.reference_sources
    ) or any("paragraphs.docx" in s for s in d1.rag_sources)
    assert "chimney" in d1.text.lower() or "stack" in d1.text.lower()
    assert "slate" in d3.text.lower() or "roof" in d3.text.lower()
    assert "slate" not in d1.text.lower() or "chimney" in d1.text.lower()