src/cdms/product_catalog.py

"""
Product Catalog
================
Single source of truth for "which pesticide products do we actually have labels
for?"  Derived dynamically from the processed-documents database instead of a
hard-coded keyword list.

Why this exists
---------------
The old product-name extraction in ``cdms_label_tool.py`` matched the user's
question against a fixed list of six names
(``["roundup", "sevin", "2,4-d", "glyphosate", "carbaryl", "atrazine"]``).
Any product outside that list (Dauntless, Kozami, Megalodon, Acquit, ...) was not
recognised, so the pipeline fell through to fuzzy heuristics and ultimately ran a
global vector search.  Because one product (Roundup in this snapshot, "Trust" in
the deployed data) owns the majority of the chunks, that global search returned
*that* product's chunks for almost any question -- the "it keeps talking about
Trust / it answers from a different herbicide" bug reported by ISA.

This module resolves a query to a product that we can actually answer for, and
lets callers *abstain* when the requested product is not in the index.
"""

import re
from functools import lru_cache
from typing import Dict, List, Optional, Set, Tuple

from src.cdms.schema import DatabaseManager, Document, DocumentChunk
from src.config.paths import DB_PATH

try:
    from rapidfuzz import fuzz
    _RAPIDFUZZ = True
except ImportError:  # pragma: no cover - rapidfuzz is a declared dependency
    _RAPIDFUZZ = False

# Trailing content hash that pdf_downloader appends, e.g. "roundup_bdc94bbee383".
_HASH_SUFFIX = re.compile(r"_[0-9a-f]{8,}$", re.IGNORECASE)
# Filenames that are actually stored user questions, not products (a separate
# download-naming bug). Skip them so they never masquerade as a product.
_NON_PRODUCT_PREFIXES = ("is there any", "what ", "how ", "tell me")

# Minimum rapidfuzz score (0-100) for a query token span to count as a product hit.
DEFAULT_MATCH_THRESHOLD = 82


def normalize_filename(filename: str) -> str:
    """Turn a stored PDF filename into a human product name.

    ``roundup_bdc94bbee383.pdf`` -> ``roundup``
    ``Brandt_Nema_Q.pdf``        -> ``brandt nema q``
    ``24-d_fa1e6bdacae6.pdf``    -> ``24-d``
    """
    stem = re.sub(r"\.pdf$", "", filename, flags=re.IGNORECASE)
    stem = _HASH_SUFFIX.sub("", stem)
    return stem.replace("_", " ").strip().lower()


class ProductCatalog:
    """Read-only view of the products we have processed labels for."""

    def __init__(self, db_path: str = None):
        self.db_path = db_path if db_path is not None else str(DB_PATH)

    def _rows(self) -> List[Tuple[str, int]]:
        """Return (filename, chunk_count) for every processed document."""
        db = DatabaseManager(db_path=self.db_path)
        session = db.get_session()
        try:
            rows = []
            for doc in session.query(Document).all():
                n = (
                    session.query(DocumentChunk)
                    .filter(DocumentChunk.document_id == doc.id)
                    .count()
                )
                rows.append((doc.filename, n))
            return rows
        finally:
            session.close()

    def catalog(self) -> Dict[str, int]:
        """Map normalized product name -> total chunk count across its PDFs."""
        catalog: Dict[str, int] = {}
        for filename, n_chunks in self._rows():
            product = normalize_filename(filename)
            if not product or product.startswith(_NON_PRODUCT_PREFIXES):
                continue
            catalog[product] = catalog.get(product, 0) + n_chunks
        return catalog

    def known_products(self) -> Set[str]:
        """Every product with a document row (even if text extraction failed)."""
        return set(self.catalog().keys())

    def available_products(self) -> Set[str]:
        """Products we can actually answer about (have at least one chunk)."""
        return {p for p, n in self.catalog().items() if n > 0}

    def is_available(self, product_name: str) -> bool:
        return normalize_filename(product_name) in self.available_products()

    def resolve(
        self,
        text: str,
        threshold: int = DEFAULT_MATCH_THRESHOLD,
    ) -> Optional[str]:
        """Resolve free text (a user question) to a known product name.

        Returns the best-matching *available* product, or ``None`` when nothing
        clears the threshold -- the caller should then abstain rather than run a
        global search that would surface the dominant product.
        """
        products = self.available_products()
        if not products:
            return None

        text_lower = text.lower()

        # 1) Exact substring wins (handles multi-word names like "brandt nema q").
        #    Prefer the longest match so "24-d" doesn't shadow a longer name.
        substring_hits = [p for p in products if p and p in text_lower]
        if substring_hits:
            return max(substring_hits, key=len)

        # 1b) Alphanumeric-only match so punctuation/spacing differences resolve:
        #     the label file is "24-d" but users type "2,4-D" -> both become "24d".
        def _alnum(s: str) -> str:
            return re.sub(r"[^a-z0-9]", "", s.lower())

        text_alnum = _alnum(text)
        alnum_hits = [p for p in products if _alnum(p) and _alnum(p) in text_alnum]
        if alnum_hits:
            return max(alnum_hits, key=len)

        if not _RAPIDFUZZ:
            return None

        # 2) Fuzzy match each product against the best-aligned span of the query.
        best_product, best_score = None, 0.0
        for product in products:
            score = fuzz.partial_ratio(product, text_lower)
            if score > best_score:
                best_product, best_score = product, score

        return best_product if best_score >= threshold else None


@lru_cache(maxsize=1)
def get_catalog() -> ProductCatalog:
    """Process-wide singleton (the DB is small and read-mostly)."""
    return ProductCatalog()


def diversify_by_product(
    results: List[dict],
    limit: int,
    max_per_product: int = 2,
) -> List[dict]:
    """Round-robin results across products so one product can't monopolise top-k.

    On a *general* query (no product filter) the vector search is dominated by
    whichever product owns most of the index (Roundup = 71% of chunks here), so
    all top-k hits come from it. This re-ranks by taking the best chunks from
    each product in turn -- preserving score order within a product -- so the
    answer draws on a variety of labels instead of a single dominant one.

    ``results`` must be dicts with ``source_file`` and (ideally) ``score``.
    Assumes ``results`` is already sorted best-first.
    """
    if not results:
        return []

    # Group by product, preserving incoming (score) order within each group.
    by_product: Dict[str, List[dict]] = {}
    for r in results:
        product = normalize_filename(r.get("source_file", "")) or "_unknown"
        by_product.setdefault(product, []).append(r)

    # Cap per product, then round-robin across products until we hit `limit`.
    queues = [items[:max_per_product] for items in by_product.values()]
    out: List[dict] = []
    idx = 0
    while len(out) < limit and any(idx < len(q) for q in queues):
        for q in queues:
            if idx < len(q):
                out.append(q[idx])
                if len(out) >= limit:
                    break
        idx += 1
    return out[:limit]


def cross_product_abstention(
    user_question: str,
    chunk_source_files: List[str],
    catalog: Optional[ProductCatalog] = None,
) -> Optional[str]:
    """Decide whether to refuse answering from a different product's chunks.

    Returns the requested product name (meaning: ABSTAIN, answer nothing) when
    the user clearly asked about an indexed product but *none* of the retrieved
    chunks come from it. Returns ``None`` when it is safe to proceed.

    This is the core guard against the ISA-reported bug where a question about
    herbicide X was answered from the label of a different (dominant) product.
    """
    catalog = catalog or get_catalog()
    requested = catalog.resolve(user_question)
    if not requested or requested not in catalog.available_products():
        return None  # nothing specific requested (or we don't have it anyway)
    chunk_products = {normalize_filename(s) for s in chunk_source_files if s}
    return requested if requested not in chunk_products else None