backend/core/paragraph_retriever.py

"""Retrieve MASTER or REFERENCE paragraphs for mapping messy field notes."""

from __future__ import annotations

import re
from typing import Literal

from backend.config import settings
from backend.core.rag_store import SearchHit, get_rag_store
from backend.models.schema import TemplateSchema

InterferenceLevel = Literal["minimum", "medium", "maximum"]
RetrievalLevel = Literal["paragraph", "section", "document"]

_STOP = frozenset({
    "about", "with", "from", "this", "that", "level", "section", "inside",
    "outside", "other", "property", "report", "survey", "your", "the", "and",
    "were", "was", "has", "have", "been", "noted", "inspected", "ground",
})


def build_retrieval_query(
    section_label: str,
    observations: list[str],
    *,
    section_id: str = "",
) -> str:
    """Semantic query anchored to the target section ID + observations."""
    sid = (section_id or "").strip().upper()
    label = section_label.strip()
    prefix = f"{sid} {label}".strip() if sid else label
    parts = [o.strip() for o in observations if o and o.strip()]
    if not parts:
        return prefix
    ranked = sorted(parts, key=len, reverse=True)
    return f"{prefix}: " + " ".join(ranked[:5])


def _content_tokens(text: str) -> set[str]:
    return {
        w
        for w in re.findall(r"[a-z]{4,}", text.lower())
        if w not in _STOP
    }


def rerank_hits_by_observations(
    hits: list[SearchHit],
    observations: list[str],
) -> list[SearchHit]:
    if not hits or not observations:
        return hits
    obs_tokens = _content_tokens(" ".join(observations))
    if not obs_tokens:
        return hits
    boost = settings.retrieval_lexical_boost

    def rank_key(h: SearchHit) -> float:
        para_tokens = _content_tokens(h.text)
        overlap = len(obs_tokens & para_tokens)
        return h.score + overlap * boost

    return sorted(hits, key=rank_key, reverse=True)


def rerank_hits_for_comprehensive_baseline(hits: list[SearchHit]) -> list[SearchHit]:
    """Among similarly scored REFERENCE hits, prefer the longest / fullest text block."""
    if len(hits) < 2:
        return hits
    top_score = hits[0].score
    score_floor = top_score - max(0.06, top_score * 0.08)

    def sort_key(h: SearchHit) -> tuple[float, int, int]:
        return (
            h.score,
            len(h.text.split()),
            -(h.paragraph_index or 0),
        )

    close = [h for h in hits if h.score >= score_floor]
    tail = [h for h in hits if h.score < score_floor]
    close.sort(key=sort_key, reverse=True)
    return close + tail


def _filter_hits_to_section(
    hits: list[SearchHit],
    paragraph_section_id: str,
) -> list[SearchHit]:
    """Keep only chunks tagged with the target section (case-insensitive)."""
    want = (paragraph_section_id or "").strip().upper()
    if not want:
        return hits
    return [
        h for h in hits
        if (h.section_id or "").strip().upper() == want
    ]


def _complete_section_group(
    tenant_id: str,
    source_key: str | None,
    section_id: str,
    *,
    tier: str | None,
    allowed_doc_keys: frozenset[str] | None,
    property_context: dict | None,
) -> list[SearchHit]:
    """Pull EVERY stored chunk for a section (optionally pinned to one source).

    Mirrors the property guard applied to similarity hits so a section-complete
    baseline cannot reintroduce wrong-property content the ranked path would drop.
    """
    from backend.core.rag_store import TIER_REFERENCE, get_rag_store

    use_tier = tier or TIER_REFERENCE
    chunks = get_rag_store().fetch_section_chunks(
        tenant_id,
        tier=use_tier,
        section_id=section_id,
        source_key=source_key,
        allowed_doc_keys=allowed_doc_keys,
    )
    if property_context and chunks:
        from backend.core.property_blocklist import filter_hits_by_property

        chunks = filter_hits_by_property(chunks, property_context)
    return chunks


def assemble_reference_baseline(
    hits: list[SearchHit],
    *,
    paragraph_section_id: str = "",
    tenant_id: str | None = None,
    tier: str | None = None,
    allowed_doc_keys: frozenset[str] | None = None,
    property_context: dict | None = None,
) -> tuple[str, list[SearchHit]]:
    """Build the fullest REFERENCE baseline from ranked hits (same doc + section).

    When ``tenant_id`` is supplied and section-complete assembly is enabled, the
    best source is chosen by similarity (as before) but the baseline is then built
    from EVERY chunk that source holds for the section — in document order — rather
    than only the top-K chunks similarity surfaced. This is what makes the whole
    past-report section get mapped instead of a partial fragment. A character budget
    bounds the assembled length so the mapping prompt stays sane.
    """
    from backend.core.paragraph_merge import combine_reference_blocks

    if not hits:
        return "", []

    by_source: dict[str, list[SearchHit]] = {}
    for hit in hits:
        key = hit.source_filename or hit.doc_id or "unknown"
        by_source.setdefault(key, []).append(hit)

    sid = (paragraph_section_id or "").upper()

    def source_rank(key: str) -> tuple[float, int]:
        group = by_source[key]
        max_score = max(h.score for h in group)
        matched = [
            h for h in group if sid and (h.section_id or "").upper() == sid
        ] or group
        word_count = sum(len(h.text.split()) for h in matched)
        return max_score + min(word_count / 400.0, 0.2), word_count

    best_source = max(by_source, key=source_rank)
    group = by_source[best_source]

    # Section-complete expansion: replace the top-K group with the chosen source's
    # ENTIRE section so no paragraph the surveyor expects is left unmapped.
    section_complete = False
    if (
        tenant_id
        and sid
        and settings.reference_section_complete_enabled
    ):
        complete = _complete_section_group(
            tenant_id,
            best_source,
            sid,
            tier=tier,
            allowed_doc_keys=allowed_doc_keys,
            property_context=property_context,
        )
        if complete:
            group = complete
            section_complete = True

    section_hits = [
        h for h in group if sid and (h.section_id or "").upper() == sid
    ]
    if not section_hits:
        return "", []

    section_hits.sort(
        key=lambda h: (h.paragraph_index or 0, -len(h.text.split()), -h.score),
    )

    # Only the expanded path is char-budgeted; the bounded top-K path is untouched.
    budget = (
        settings.reference_section_complete_max_chars if section_complete else None
    )
    texts: list[str] = []
    contributing: list[SearchHit] = []
    seen_para: set[int] = set()
    used = 0
    for hit in section_hits:
        para_idx = hit.paragraph_index or 0
        if para_idx and para_idx in seen_para:
            continue
        body = hit.text.strip()
        if not body:
            continue
        if budget is not None and texts and used + len(body) > budget:
            break
        if para_idx:
            seen_para.add(para_idx)
        if body not in texts:
            texts.append(body)
            contributing.append(hit)
            used += len(body)

    if not texts:
        best = max(section_hits, key=lambda h: (h.score, len(h.text.split())))
        return best.text.strip(), [best]

    baseline = combine_reference_blocks(texts[0], texts[1:])
    return baseline, contributing


def fetch_complete_section_baseline(
    tenant_id: str,
    *,
    paragraph_section_id: str,
    tier: str | None = None,
    allowed_doc_keys: frozenset[str] | None = None,
    property_context: dict | None = None,
) -> tuple[str, list[SearchHit]]:
    """Metadata-only baseline for a section (no similarity query).

    Fallback for when similarity search surfaced nothing for a section that
    nonetheless exists in the index (weak query match, alias drift). Returns the
    fullest single-source section, or ``("", [])`` when the section is genuinely
    absent — in which case the caller correctly degrades to a notes-only section.
    """
    if not settings.reference_section_complete_enabled:
        return "", []
    chunks = _complete_section_group(
        tenant_id,
        None,
        paragraph_section_id,
        tier=tier,
        allowed_doc_keys=allowed_doc_keys,
        property_context=property_context,
    )
    if not chunks:
        return "", []
    return assemble_reference_baseline(
        chunks,
        paragraph_section_id=paragraph_section_id,
        tenant_id=tenant_id,
        tier=tier,
        allowed_doc_keys=allowed_doc_keys,
        property_context=property_context,
    )


def _uses_reference_tier(interference_level: InterferenceLevel) -> bool:
    return interference_level in ("minimum", "medium", "maximum")


def _retrieval_params(
    retrieval_level: RetrievalLevel,
    *,
    paragraph_section_id: str,
    k: int,
) -> tuple[bool, int]:
    level = (retrieval_level or "paragraph").lower()
    if level == "section":
        return bool(paragraph_section_id), max(k, 8)
    if level == "document":
        return bool(paragraph_section_id), max(k, 15)
    return bool(paragraph_section_id), k


def _focus_document_hits(hits: list[SearchHit], k: int) -> list[SearchHit]:
    """Prefer chunks from the single best-matching past report."""
    if not hits:
        return hits
    by_source: dict[str, list[SearchHit]] = {}
    for hit in hits:
        key = hit.source_filename or hit.doc_id or "unknown"
        by_source.setdefault(key, []).append(hit)
    best_key = max(by_source, key=lambda name: max(h.score for h in by_source[name]))
    focused = sorted(by_source[best_key], key=lambda h: h.score, reverse=True)
    return focused[:k]


def _search_tier(
    tenant_id: str,
    query: str,
    *,
    paragraph_section_id: str,
    k: int,
    interference_level: InterferenceLevel,
    retrieval_level: RetrievalLevel = "paragraph",
    allowed_doc_keys: frozenset[str] | None = None,
) -> list[SearchHit]:
    store = get_rag_store()
    if _uses_reference_tier(interference_level):
        search = store.search_for_reference_mapping
    else:
        search = store.search_for_generation

    strict, fetch_k = _retrieval_params(
        retrieval_level,
        paragraph_section_id=paragraph_section_id,
        k=k,
    )
    search_kwargs = {
        "section_id": paragraph_section_id or None,
        "allowed_doc_keys": allowed_doc_keys,
    }
    # Never widen to cross-section retrieval — empty is safer than wrong-section bleed.
    use_strict = strict or bool(paragraph_section_id)
    hits = search(
        tenant_id,
        query,
        top_k=max(fetch_k * 3, fetch_k),
        section_strict=use_strict,
        **search_kwargs,
    )
    hits = _filter_hits_to_section(hits, paragraph_section_id)
    if (retrieval_level or "paragraph").lower() == "document":
        hits = _focus_document_hits(hits, fetch_k)
    return hits


def retrieve_paragraphs_for_mapping(
    tenant_id: str,
    *,
    section_label: str,
    paragraph_section_id: str,
    observations: list[str],
    interference_level: InterferenceLevel = "minimum",
    retrieval_level: RetrievalLevel = "paragraph",
    top_k: int | None = None,
    allowed_doc_keys: frozenset[str] | None = None,
    property_context: dict | None = None,
) -> list[SearchHit]:
    """Ranked paragraphs from uploaded past reports for all interference levels.

    When ``property_context`` is supplied, paragraphs whose terminology is
    incompatible with the confirmed property type (e.g. flat-only content in a
    house report) are dropped before ranking — the next compatible candidate is
    used rather than leaving a wrong-property baseline.
    """
    k = top_k or settings.reference_baseline_top_k
    if interference_level == "medium":
        k = max(k, settings.reference_baseline_top_k)
    elif interference_level == "maximum":
        k = max(k, settings.reference_baseline_top_k + 2)
    query = build_retrieval_query(
        section_label,
        observations,
        section_id=paragraph_section_id,
    )
    hits = _search_tier(
        tenant_id,
        query,
        paragraph_section_id=paragraph_section_id,
        k=k,
        interference_level=interference_level,
        retrieval_level=retrieval_level,
        allowed_doc_keys=allowed_doc_keys,
    )
    if property_context:
        from backend.core.property_blocklist import filter_hits_by_property

        hits = filter_hits_by_property(hits, property_context)
    hits = rerank_hits_by_observations(hits, observations)
    hits = rerank_hits_for_comprehensive_baseline(hits)
    return hits[:k]


def find_paragraph_by_topic(
    tenant_id: str,
    observations: list[str],
    *,
    interference_level: InterferenceLevel = "minimum",
    top_k: int = 3,
    allowed_doc_keys: frozenset[str] | None = None,
    paragraph_section_id: str = "",
) -> list[SearchHit]:
    """Topic search scoped to ``paragraph_section_id`` when provided."""
    parts = [o.strip() for o in observations if o and o.strip()]
    if not parts:
        return []
    query = " ".join(sorted(parts, key=len, reverse=True)[:5])
    store = get_rag_store()
    sid = (paragraph_section_id or "").strip() or None
    use_strict = bool(sid)
    if _uses_reference_tier(interference_level):
        hits = store.search_for_reference_mapping(
            tenant_id,
            query,
            section_id=sid,
            top_k=max(top_k * 5, top_k),
            section_strict=use_strict,
            allowed_doc_keys=allowed_doc_keys,
        )
    else:
        hits = store.search_for_generation(
            tenant_id,
            query,
            section_id=sid,
            top_k=max(top_k * 5, top_k),
            section_strict=use_strict,
        )
    hits = _filter_hits_to_section(hits, paragraph_section_id)
    ranked = rerank_hits_by_observations(hits, parts)
    return rerank_hits_for_comprehensive_baseline(ranked)[:top_k]


def report_section_for_paragraph_id(
    schema: TemplateSchema,
    paragraph_section_id: str,
) -> str | None:
    if not paragraph_section_id:
        return None
    for report_id, para_id in schema.section_alias_map.items():
        if para_id == paragraph_section_id:
            return report_id
    if schema.get_section(paragraph_section_id):
        return paragraph_section_id
    return None


def guess_report_section_from_topic(
    schema: TemplateSchema,
    observations: list[str],
    reference_text: str = "",
) -> str | None:
    """Best-effort section when REFERENCE chunks lack section_id metadata."""
    corpus = " ".join(observations) + " " + reference_text
    words = _content_tokens(corpus)
    if not words:
        return None

    best_id: str | None = None
    best_score = 0
    for sec in schema.ordered_sections():
        title_words = _content_tokens(f"{sec.id} {sec.title}")
        overlap = len(words & title_words)
        if overlap > best_score:
            best_score = overlap
            best_id = sec.id
    return best_id if best_score > 0 else None


# Backward-compatible alias — always searches MASTER tier.
def retrieve_master_paragraphs(
    tenant_id: str,
    *,
    section_label: str,
    paragraph_section_id: str,
    observations: list[str],
    top_k: int | None = None,
) -> list[SearchHit]:
    k = top_k or settings.retrieval_top_k
    query = build_retrieval_query(
        section_label,
        observations,
        section_id=paragraph_section_id,
    )
    store = get_rag_store()
    hits = store.search_for_generation(
        tenant_id,
        query,
        section_id=paragraph_section_id or None,
        top_k=max(k * 3, k),
        section_strict=bool(paragraph_section_id),
    )
    hits = _filter_hits_to_section(hits, paragraph_section_id)
    hits = rerank_hits_by_observations(hits, observations)
    hits = rerank_hits_for_comprehensive_baseline(hits)
    return hits[:k]