app/services/canonical_rollout.py

"""Semantic scan and best-effort apply for canonical paragraph rollout across the library."""

from __future__ import annotations

import logging
import re
from pathlib import Path

from fastapi import HTTPException
from sqlalchemy.ext.asyncio import AsyncSession

from app.chunking.splitter import split_documents
from app.config import settings
from app.ingest.pipeline import ingest_document, load_raw_documents
from app.models.schemas import (
    CanonicalApplyRequest,
    CanonicalApplyResponse,
    CanonicalRolloutMatch,
    CanonicalRolloutRequest,
    CanonicalRolloutResponse,
)
from app.services.content_similarity import jaccard_similarity
from app.services.provenance_enrichment import fetch_doc_filenames
from app.vectorstore.factory import get_vectorstore

logger = logging.getLogger(__name__)

_MIN_CANON_LEN = 12
_FETCH_MULT = 8
_MAX_SNIPPET = 4000

# Numbers / units likely to be document-specific (conservative extraction)
_FACT_TOKEN_RE = re.compile(
    r"\b\d+(?:\.\d+)?(?:\s*(?:mm|cm|m|kN|kPa|%|°C|°F))?\b|\b(?:19|20)\d{2}\b",
    re.IGNORECASE,
)


def _preservation_note(chunk_text: str, canonical: str) -> str:
    """List fact-like tokens present in chunk but not in canonical (case-insensitive)."""
    cn = canonical.lower()
    found = [m.group(0) for m in _FACT_TOKEN_RE.finditer(chunk_text)]
    extra = [x for x in found if x.lower() not in cn]
    if not extra:
        return ""
    # Dedupe preserving order
    seen: set[str] = set()
    uniq = []
    for x in extra:
        k = x.lower()
        if k in seen:
            continue
        seen.add(k)
        uniq.append(x)
    return "Document-specific tokens in chunk (verify in replacement): " + ", ".join(uniq[:24])


def _compatibility_tier(relevance_pct: float, jac: float) -> str:
    if relevance_pct >= 55.0 and jac >= 0.14:
        return "high"
    if relevance_pct >= 28.0 and jac >= 0.08:
        return "review"
    return "low"


def _adapt_with_llm(canonical: str, chunk_text: str) -> str:
    from langchain_core.messages import HumanMessage, SystemMessage
    from langchain_openai import ChatOpenAI

    llm = ChatOpenAI(
        model=settings.chat_model,
        temperature=0.1,
        api_key=settings.openai_api_key,
        max_retries=2,
    )
    sys = SystemMessage(
        content=(
            "You merge an improved canonical paragraph with text from an indexed document chunk. "
            "Respond with exactly one replacement paragraph, plain text. "
            "Preserve every number, unit, date, and document-specific proper noun from the original chunk "
            "that does not already appear in the canonical text. Do not invent facts. "
            "If preservation is impossible without invention, return the canonical text verbatim."
        )
    )
    hum = HumanMessage(content=f"CANONICAL:\n{canonical}\n\nORIGINAL_CHUNK:\n{chunk_text}")
    out = llm.invoke([sys, hum])
    text = (getattr(out, "content", None) or "").strip()
    return text if text else canonical


async def _adapt_with_llm_async(canonical: str, chunk_text: str) -> str:
    from langchain_core.messages import HumanMessage, SystemMessage
    from langchain_openai import ChatOpenAI

    from app.llm.llm_throttle import throttled_llm_call

    llm = ChatOpenAI(
        model=settings.chat_model,
        temperature=0.1,
        api_key=settings.openai_api_key,
        max_retries=2,
    )
    sys = SystemMessage(
        content=(
            "You merge an improved canonical paragraph with text from an indexed document chunk. "
            "Respond with exactly one replacement paragraph, plain text. "
            "Preserve every number, unit, date, and document-specific proper noun from the original chunk "
            "that does not already appear in the canonical text. Do not invent facts. "
            "If preservation is impossible without invention, return the canonical text verbatim."
        )
    )
    hum = HumanMessage(content=f"CANONICAL:\n{canonical}\n\nORIGINAL_CHUNK:\n{chunk_text}")

    async def _call() -> str:
        out = await llm.ainvoke([sys, hum])
        return (getattr(out, "content", None) or "").strip()

    text = await throttled_llm_call(
        phase="canonical_adapt",
        section_id=None,
        cache_hit=None,
        call=_call,
    )
    return text if text else canonical


def _build_proposed(
    canonical: str,
    chunk_text: str,
    adapt_with_llm: bool,
) -> tuple[str, str]:
    note = _preservation_note(chunk_text, canonical)
    if adapt_with_llm and settings.openai_api_key.strip():
        try:
            proposed = _adapt_with_llm(canonical, chunk_text)
            return proposed, note
        except Exception as exc:  # noqa: BLE001
            logger.warning("LLM adaptation failed: %s", exc)
            return canonical, note + " (LLM adaptation failed; proposed text is canonical.)"
    if adapt_with_llm and not settings.openai_api_key.strip():
        return canonical, note + " (Set OPENAI_API_KEY to enable adapt_with_llm.)"
    return canonical, note


async def _build_proposed_async(
    canonical: str,
    chunk_text: str,
    adapt_with_llm: bool,
) -> tuple[str, str]:
    note = _preservation_note(chunk_text, canonical)
    if adapt_with_llm and settings.openai_api_key.strip():
        try:
            proposed = await _adapt_with_llm_async(canonical, chunk_text)
            return proposed, note
        except Exception as exc:  # noqa: BLE001
            logger.warning("LLM adaptation failed: %s", exc)
            return canonical, note + " (LLM adaptation failed; proposed text is canonical.)"
    if adapt_with_llm and not settings.openai_api_key.strip():
        return canonical, note + " (Set OPENAI_API_KEY to enable adapt_with_llm.)"
    return canonical, note


async def scan_canonical_rollout(
    db: AsyncSession,
    tenant_id: str,
    body: CanonicalRolloutRequest,
) -> CanonicalRolloutResponse:
    """Find semantically similar chunks; propose replacements (canonical or LLM-adapted)."""
    canonical = body.canonical_text.strip()
    if len(canonical) < _MIN_CANON_LEN:
        return CanonicalRolloutResponse(
            matches=[],
            message="canonical_text is too short — use at least 12 characters for a stable embedding match.",
        )

    from app.retrieval.vector_search import async_vs_search

    vs = get_vectorstore()
    exclude_docs = set(body.exclude_document_ids or [])
    scope_docs = set(body.document_ids) if body.document_ids else None

    fetch_k = min(max(body.limit * _FETCH_MULT, 32), 200)
    raw = await async_vs_search(vs, canonical, tenant_id, k=fetch_k)
    if not raw:
        return CanonicalRolloutResponse(
            matches=[],
            message="No indexed chunks for this tenant, or nothing similar to the canonical paragraph.",
        )

    mx = max(r.score for r in raw)
    doc_ids = {r.doc_id for r in raw if r.doc_id}
    filenames = await fetch_doc_filenames(db, tenant_id, doc_ids)

    matches: list[CanonicalRolloutMatch] = []
    seen_chunk: set[str] = set()

    for r in raw:
        if not r.doc_id or r.doc_id in exclude_docs:
            continue
        if scope_docs is not None and r.doc_id not in scope_docs:
            continue
        if r.chunk_id in seen_chunk:
            continue

        chunk_text = (r.text or "").strip()
        jac = jaccard_similarity(canonical, chunk_text)
        if jac < body.min_jaccard_vs_canonical:
            continue

        pct = 100.0 * r.score / mx if mx > 0 else 0.0
        if pct < body.min_relevance_percent:
            continue

        seen_chunk.add(r.chunk_id)
        compat = _compatibility_tier(pct, jac)
        proposed, note = await _build_proposed_async(
            canonical, chunk_text, body.adapt_with_llm
        )

        snip = chunk_text if len(chunk_text) <= _MAX_SNIPPET else chunk_text[: _MAX_SNIPPET - 1] + "…"

        matches.append(
            CanonicalRolloutMatch(
                chunk_id=r.chunk_id,
                document_id=r.doc_id,
                filename=filenames.get(r.doc_id),
                original_snippet=snip,
                relevance_percent=round(pct, 1),
                jaccard_vs_canonical=round(jac, 4),
                compatibility=compat,
                proposed_replacement=proposed,
                preservation_note=note,
            )
        )
        if len(matches) >= body.limit:
            break

    msg = ""
    if not matches and raw:
        msg = (
            "No chunks passed your filters (relevance, Jaccard vs canonical, or document scope). "
            "Lower min_relevance_percent or min_jaccard_vs_canonical, or widen document_ids."
        )

    return CanonicalRolloutResponse(matches=matches, message=msg)


def get_chunk_text_from_file(file_path: Path, doc_id: str, chunk_id: str) -> str:
    """Recompute chunk text from the file using the same split as ingestion."""
    if not chunk_id.startswith(doc_id + "_"):
        msg = f"chunk_id {chunk_id!r} does not belong to document {doc_id}"
        raise ValueError(msg)
    idx_str = chunk_id.rsplit("_", 1)[-1]
    try:
        idx = int(idx_str)
    except ValueError as exc:
        raise ValueError(f"invalid chunk index in chunk_id={chunk_id!r}") from exc

    raw = load_raw_documents(file_path)
    chunks = split_documents(raw)
    if idx < 0 or idx >= len(chunks):
        raise ValueError(f"chunk index {idx} out of range (file has {len(chunks)} chunks)")
    return (chunks[idx].page_content or "").strip()


def _write_plain_docx(path: Path, full_text: str) -> None:
    """Rebuild a .docx from plain text (one paragraph per line). Destroys complex layout."""
    from docx import Document as DocxDocument

    doc = DocxDocument()
    lines = full_text.split("\n")
    if not lines:
        doc.add_paragraph("")
    else:
        for line in lines:
            doc.add_paragraph(line)
    doc.save(str(path))


async def apply_canonical_replacement(
    db: AsyncSession,
    tenant_id: str,
    body: CanonicalApplyRequest,
) -> CanonicalApplyResponse:
    """Replace one chunk inside a .docx by editing extracted text, then re-ingest."""
    if not body.confirm_destructive_docx:
        raise HTTPException(
            status_code=400,
            detail="Set confirm_destructive_docx=true to acknowledge that the .docx will be rebuilt as plain paragraphs.",
        )

    from app.db.models import Document as DBDocument

    doc = await db.get(DBDocument, body.document_id)
    if doc is None or doc.tenant_id != tenant_id:
        raise HTTPException(status_code=404, detail="Document not found")

    path = Path(doc.file_path)
    if path.suffix.lower() != ".docx":
        raise HTTPException(
            status_code=422,
            detail="Automatic apply supports .docx only. Replace PDF content offline and re-upload.",
        )

    if not path.is_file():
        raise HTTPException(status_code=404, detail="File missing on disk")

    try:
        chunk_text = get_chunk_text_from_file(path, doc.id, body.chunk_id)
    except ValueError as exc:
        raise HTTPException(status_code=422, detail=str(exc)) from exc

    raw_docs = load_raw_documents(path)
    if len(raw_docs) != 1:
        raise HTTPException(
            status_code=422,
            detail="Apply supports a single continuous text body per .docx; edit multi-part files manually.",
        )
    full = raw_docs[0].page_content or ""

    if chunk_text not in full:
        raise HTTPException(
            status_code=422,
            detail="Chunk text not found verbatim in extracted file text — file may have changed; re-ingest and retry.",
        )

    new_full = full.replace(chunk_text, body.replacement_text, 1)
    _write_plain_docx(path, new_full)

    vs = get_vectorstore()
    vs.delete_document(doc.id)
    await ingest_document(doc.id, path)

    from app.retrieval.semantic_cache import invalidate_semantic_cache_for_tenant

    await invalidate_semantic_cache_for_tenant(doc.tenant_id)

    n = vs.count_for_doc(doc.id)
    return CanonicalApplyResponse(
        document_id=doc.id,
        chunk_id=body.chunk_id,
        filename=doc.filename,
        detail="Document rewritten with replacement text and re-indexed.",
        chunks_indexed=n,
    )