src/agents.py

"""
agents.py — Specialist document extraction agents for UK Motor Insurance.

Architecture
────────────
  PDF path
    → docling  (PDF → Markdown)
    → PIIMasker.mask()
    → InsuranceExtractionAgents.classify_document()  [LLM: llama-3.1-8b-instant]
    → extract_schedule() | extract_certificate()     [LLM: llama-4-scout-17b]
    → UKMotorGoldenRecord (with source_document provenance)
"""
from __future__ import annotations

import json
import logging
import os
import time
from pathlib import Path
from typing import Any

import instructor
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions
from docling.document_converter import DocumentConverter, PdfFormatOption
from groq import Groq
from pydantic import ValidationError

from privacy import PIIMasker
from prompts import PromptRegistry
from schema import DocumentType, SourceMetadata, UKMotorGoldenRecord
from settings import settings

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Groq clients — extraction (instructor-wrapped) + classifier (raw Groq)
# ---------------------------------------------------------------------------


def _build_extraction_client() -> instructor.Instructor:
    api_key = os.environ.get("GROQ_API_KEY")
    if not api_key:
        raise EnvironmentError(
            "GROQ_API_KEY environment variable is not set. "
            "Export it before running the pipeline."
        )
    return instructor.from_groq(Groq(api_key=api_key), mode=instructor.Mode.JSON)


def _build_groq_client() -> Groq:
    api_key = os.environ.get("GROQ_API_KEY")
    if not api_key:
        raise EnvironmentError(
            "GROQ_API_KEY environment variable is not set. "
            "Export it before running the pipeline."
        )
    return Groq(api_key=api_key)


# Models resolved at import time from settings.yaml / env vars
_EXTRACTION_MODEL: str = settings.llm.model
_CLASSIFIER_MODEL: str = settings.llm.classifier_model


def _build_docling_converter() -> DocumentConverter:
    """Build a DocumentConverter configured from settings.docling."""
    opts = PdfPipelineOptions()
    opts.do_ocr = settings.docling.do_ocr
    opts.do_table_structure = settings.docling.do_table_structure
    return DocumentConverter(
        format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=opts)}
    )

# ---------------------------------------------------------------------------
# Document type classifier (keyword heuristic — fast, zero API calls)
# ---------------------------------------------------------------------------

_CLASSIFICATION_KEYWORDS: dict[DocumentType, list[str]] = {
    DocumentType.SCHEDULE: [
        # Phrases that only appear in a Schedule, not in a Certificate
        "policy schedule",
        "schedule of insurance",
        "schedule number",
        "premium payable",
        "compulsory excess",
        "voluntary excess",
        "no claims bonus",
        "ncb",
        "windscreen excess",
    ],
    DocumentType.CERTIFICATE: [
        # Phrases that are definitive for a Certificate document
        "certificate of motor insurance",
        "motor insurance certificate",
        "certificate number",
        "persons entitled to drive",
        "class of use",
        "road traffic act",
        "this is to certify",
    ],
    DocumentType.STATEMENT_OF_FACT: [
        "statement of fact",
        "statement of insurance",
        "proposal form",
        "claims history",
        "motoring convictions",
        "annual mileage",
    ],
    DocumentType.POLICY_BOOKLET: [
        "policy booklet",
        "policy wording",
        "terms and conditions",
        "what is covered",
        "general conditions",
        "complaints procedure",
    ],
}


def _keyword_classify(text: str) -> str:
    """Keyword heuristic fallback classifier. Returns DocumentType.value string."""
    lower = text.lower()
    scores: dict[DocumentType, int] = {dt: 0 for dt in _CLASSIFICATION_KEYWORDS}

    for doc_type, keywords in _CLASSIFICATION_KEYWORDS.items():
        for kw in keywords:
            if kw in lower:
                scores[doc_type] += 1

    best_type, best_score = max(scores.items(), key=lambda kv: kv[1])
    return best_type.value if best_score > 0 else DocumentType.UNKNOWN.value


def _str_to_doc_type(s: str) -> DocumentType:
    """Convert a string to DocumentType, falling back to UNKNOWN."""
    try:
        return DocumentType(s)
    except ValueError:
        return DocumentType.UNKNOWN


# ---------------------------------------------------------------------------
# Extraction failure sentinel
# ---------------------------------------------------------------------------


class ExtractionFailedError(RuntimeError):
    """
    Raised when the LLM fails to produce a valid UKMotorGoldenRecord after
    exhausting all retries.  Callers should treat the document as failed and
    skip it rather than propagating an empty record silently.
    """


# ---------------------------------------------------------------------------
# InsuranceExtractionAgents
# ---------------------------------------------------------------------------


class InsuranceExtractionAgents:
    """
    Specialist extraction agents for UK Motor Insurance documents.

    Uses two LLM models:
    - llama-3.1-8b-instant     — fast document type classification
    - llama-4-scout-17b-16e    — deep structured extraction (Schedule / Certificate)

    Parameters
    ----------
    masker : PIIMasker | None
    max_retries : int
    prompt_registry : PromptRegistry | None
    debug_dir : Path | None
    """

    def __init__(
        self,
        masker: PIIMasker | None = None,
        max_retries: int = settings.llm.max_retries,
        prompt_registry: PromptRegistry | None = None,
        debug_dir: Path | None = None,
    ) -> None:
        self._client = _build_extraction_client()
        self._groq = _build_groq_client()
        self._masker = masker or PIIMasker()
        self._max_retries = max_retries
        self._prompts = prompt_registry or PromptRegistry()
        self._converter = _build_docling_converter()
        self._debug_dir = debug_dir

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def classify_document(self, markdown_text: str) -> str:
        """
        Use llama-3.1-8b-instant to classify the document type.

        The LLM is the primary classifier.  If it fails or returns an invalid
        label, the keyword heuristic is used as a fallback.  A discrepancy
        between the two is logged as a warning to flag low-confidence cases.

        Returns one of: "Schedule", "Certificate", "StatementOfFact",
        "PolicyBooklet", "Unknown".
        """
        keyword_result = _keyword_classify(markdown_text)

        system_prompt = (
            "You are a UK motor insurance document classifier.\n"
            "Given the document text, respond with EXACTLY one word from:\n"
            "Schedule | Certificate | StatementOfFact | PolicyBooklet | Unknown\n\n"
            "- Schedule: Policy Schedule \u2014 excess figures, premium, NCB, "
            "vehicle details, driver ages/DOBs.\n"
            "- Certificate: Certificate of Motor Insurance \u2014 Road Traffic Act, "
            "'persons entitled to drive', 'class of use'.\n"
            "- StatementOfFact: Statement of Fact / Proposal \u2014 claims history, "
            "convictions, annual mileage.\n"
            "- PolicyBooklet: Policy Booklet / Wording \u2014 terms and conditions, "
            "'what is covered', complaints.\n"
            "- Unknown: Cannot determine.\n\n"
            "Respond with ONLY the single classification word. No punctuation."
        )
        try:
            response = self._groq.chat.completions.create(
                model=_CLASSIFIER_MODEL,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {
                        "role": "user",
                        "content": "Classify this document:\n\n" + markdown_text[:4000],
                    },
                ],
                max_tokens=10,
                temperature=0,
            )
            llm_result = response.choices[0].message.content.strip().split()[0]
            valid = {"Schedule", "Certificate", "StatementOfFact", "PolicyBooklet", "Unknown"}
            if llm_result in valid:
                if llm_result != keyword_result:
                    logger.warning(
                        "Classifier discrepancy: LLM=%s, keyword=%s "
                        "(using LLM result — verify document type)",
                        llm_result, keyword_result,
                    )
                else:
                    logger.debug("Classifier agreement: LLM=%s \u2713", llm_result)
                return llm_result
            logger.warning(
                "LLM classifier returned '%s' \u2014 falling back to keyword heuristic", llm_result
            )
        except Exception as exc:  # noqa: BLE001
            logger.warning(
                "LLM classifier failed (%s) \u2014 falling back to keyword heuristic", exc
            )
        return keyword_result

    def extract_schedule(self, markdown_text: str, filename: str) -> UKMotorGoldenRecord:
        """
        Extract all financial, vehicle, and driver risk data from a Policy Schedule.

        Instructs the LLM to:
        - Compute total_accidental_damage = standard_compulsory + voluntary
        - Extract driver DOBs and distinguish Full UK vs UK Provisional licence types
        - Separate fire excess from theft excess (they can differ)
        - Extract own_repairer_additional_excess if present
        - Extract premium breakdown and optional extras (float if purchased,
          "Not Selected" if not)
        """
        return self._extract(
            markdown_text,
            filename,
            DocumentType.SCHEDULE,
            self._prompts.get(DocumentType.SCHEDULE),
        )

    def extract_certificate(self, markdown_text: str, filename: str) -> UKMotorGoldenRecord:
        """
        Extract legal permissions from a Certificate of Motor Insurance.

        Instructs the LLM to:
        - Extract the exact "Limitations as to use" / class_of_use clause verbatim
        - Extract the policy_number for cross-reference
        - Record driving_other_cars entitlement (true/false)
        - Leave all financial fields (excess, premium, NCB) as null
        """
        return self._extract(
            markdown_text,
            filename,
            DocumentType.CERTIFICATE,
            self._prompts.get(DocumentType.CERTIFICATE),
        )

    def process(self, pdf_path: str | Path) -> tuple[UKMotorGoldenRecord, str]:
        """
        Full pipeline for one PDF: PDF → Markdown → PII mask → classify → extract.

        Returns
        -------
        tuple[UKMotorGoldenRecord, str]
            The extracted record and the document type string (e.g. "Schedule").

        Raises
        ------
        ExtractionFailedError
            When the LLM fails to extract a valid record after all retries.
        """
        record, doc_type_str, _ = self._process_internal(Path(pdf_path), build_corpus=False)
        return record, doc_type_str

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _process_internal(
        self,
        pdf_path: Path,
        build_corpus: bool,
    ) -> tuple[UKMotorGoldenRecord, str, Any]:
        """
        Unified core pipeline: PDF → Markdown → PII mask → classify → extract,
        optionally building a ProvenanceCorpus from the raw Docling IR.

        Parameters
        ----------
        pdf_path : Path
        build_corpus : bool
            When True, builds a ProvenanceCorpus before PII masking so the
            original text is available for fuzzy matching.

        Returns
        -------
        tuple[UKMotorGoldenRecord, str, ProvenanceCorpus | None]
            (record, doc_type_str, corpus_or_None)

        Raises
        ------
        ExtractionFailedError
            Propagated from _extract() when the LLM fails after all retries.
        """
        from provenance import ProvenanceCorpus  # local import — avoids circular dep

        logger.info("Processing%s: %s", " (with provenance)" if build_corpus else "", pdf_path.name)

        # Pre-classify from filename for page-cap selection (no API call)
        pre_type_str = _keyword_classify(pdf_path.stem)
        pre_doc_type = _str_to_doc_type(pre_type_str)
        logger.debug("  Pre-classified from filename: %s", pre_type_str)

        # PDF → Markdown + raw DoclingDocument
        markdown, raw_doc = self._pdf_to_markdown_and_doc(pdf_path, pre_doc_type)

        # Build corpus from original text BEFORE masking (critical for accurate fuzzy match)
        corpus: Any = None
        if build_corpus:
            corpus = ProvenanceCorpus(source_filename=pdf_path.name, doc_type=pre_type_str)
            corpus.add_from_docling(raw_doc, pdf_path.name)
            logger.debug("  Provenance corpus: %d items", len(corpus.items))

        if self._debug_dir and settings.debug.save_markdown:
            _write_debug(self._debug_dir, f"{pdf_path.name}.md", markdown)

        # PII mask
        masked_markdown, _token_map = self._masker.mask(markdown)
        if self._debug_dir and settings.debug.save_masked_markdown:
            _write_debug(self._debug_dir, f"{pdf_path.name}.masked.md", masked_markdown)

        # Classify
        t0 = time.monotonic()
        doc_type_str = self.classify_document(masked_markdown)
        logger.info("  Classified as: %s", doc_type_str)

        # Route to specialist extractor
        if doc_type_str == "Schedule":
            record = self.extract_schedule(masked_markdown, pdf_path.name)
        elif doc_type_str == "Certificate":
            record = self.extract_certificate(masked_markdown, pdf_path.name)
        else:
            logger.info("  Non-primary type '%s' — running generic extraction", doc_type_str)
            record = self._extract(
                masked_markdown,
                pdf_path.name,
                _str_to_doc_type(doc_type_str),
                self._prompts.get(_str_to_doc_type(doc_type_str)),
            )

        elapsed = round(time.monotonic() - t0, 3)

        record.source_document = SourceMetadata(
            document_type=_str_to_doc_type(doc_type_str),
            filename=pdf_path.name,
        )

        if self._debug_dir and settings.debug.save_extraction_json:
            _write_debug(
                self._debug_dir,
                f"{pdf_path.name}.extraction.json",
                record.model_dump_json(indent=2),
            )
            fc = getattr(record, "field_citations", None) or {}
            logger.info("  field_citations populated by LLM: %d entries", len(fc))
            if fc:
                import json as _json
                _write_debug(
                    self._debug_dir,
                    f"{pdf_path.name}.field_citations.json",
                    _json.dumps(fc, indent=2, ensure_ascii=False),
                )

        if self._debug_dir and settings.debug.save_metrics:
            metrics: dict = {
                "filename": pdf_path.name,
                "doc_type": doc_type_str,
                "extraction_model": _EXTRACTION_MODEL,
                "classifier_model": _CLASSIFIER_MODEL,
                "response_time_seconds": elapsed,
            }
            if corpus is not None:
                metrics["corpus_items"] = len(corpus.items)
            _append_metrics(self._debug_dir, metrics)

        return record, doc_type_str, corpus

    def _pdf_to_markdown(
        self, pdf_path: Path, doc_type: DocumentType = DocumentType.UNKNOWN
    ) -> str:
        """Convert a PDF to Markdown using docling, respecting per-doc-type page caps."""
        markdown, _ = self._pdf_to_markdown_and_doc(pdf_path, doc_type)
        return markdown

    def _pdf_to_markdown_and_doc(
        self, pdf_path: Path, doc_type: DocumentType = DocumentType.UNKNOWN
    ) -> tuple[str, Any]:
        """Convert PDF to Markdown and also return the raw DoclingDocument for provenance."""
        # Apply page cap during conversion (not just in Markdown export) to prevent
        # Docling's layout model from running out of memory on large PDFs (Policy Booklet).
        max_pg = settings.docling.max_pages.get(doc_type.value)
        convert_kwargs: dict[str, Any] = {}
        if max_pg is not None:
            convert_kwargs["max_num_pages"] = max_pg

        result = self._converter.convert(str(pdf_path), **convert_kwargs)
        doc = result.document
        markdown = doc.export_to_markdown()

        if max_pg is not None:
            separator = "\n---\n"
            parts = markdown.split(separator)
            if len(parts) > max_pg:
                logger.info(
                    "  Page cap applied: %s capped at %d/%d pages",
                    pdf_path.name, max_pg, len(parts),
                )
                markdown = separator.join(parts[:max_pg])

        return markdown, doc

    def process_with_provenance(
        self, pdf_path: str | Path
    ) -> tuple[UKMotorGoldenRecord, str, Any]:
        """
        Like process() but also returns a ProvenanceCorpus built from the Docling IR.

        The corpus is constructed *before* PII masking so that the original text
        strings (not masked tokens) are available for fuzzy matching.

        Returns
        -------
        tuple[UKMotorGoldenRecord, str, ProvenanceCorpus]
            (record, doc_type_str, corpus)

        Raises
        ------
        ExtractionFailedError
            When the LLM fails to extract a valid record after all retries.
        """
        return self._process_internal(Path(pdf_path), build_corpus=True)  # type: ignore[return-value]

    def _extract(
        self,
        text: str,
        filename: str,
        doc_type: DocumentType,
        system_prompt: str,
    ) -> UKMotorGoldenRecord:
        """Call Groq via instructor to extract a UKMotorGoldenRecord."""
        user_message = (
            "Extract all motor insurance data from the following document text. "
            "Return a JSON object that strictly conforms to the UKMotorGoldenRecord schema.\n\n"
            f"--- DOCUMENT TEXT ---\n{text}\n--- END ---"
        )
        try:
            record: UKMotorGoldenRecord = self._client.chat.completions.create(
                model=_EXTRACTION_MODEL,
                response_model=UKMotorGoldenRecord,
                max_retries=self._max_retries,
                messages=[
                    {"role": "system", "content": system_prompt.strip()},
                    {"role": "user", "content": user_message},
                ],
            )
        except (ValidationError, Exception) as exc:
            raise ExtractionFailedError(
                f"Extraction failed for {doc_type.value!r} document '{filename}' "
                f"after {self._max_retries} retries: {exc}"
            ) from exc
        return record


# ---------------------------------------------------------------------------
# Debug helpers (module-level so they can be unit-tested independently)
# ---------------------------------------------------------------------------


def _write_debug(debug_dir: Path, filename: str, content: str) -> None:
    """Write a debug artifact to disk, silently skipping on any I/O error."""
    try:
        (debug_dir / filename).write_text(content, encoding="utf-8")
        logger.debug("Debug artifact saved: %s", filename)
    except OSError as exc:
        logger.warning("Could not write debug artifact %s: %s", filename, exc)


def _append_metrics(debug_dir: Path, metrics: dict) -> None:
    """Append a metrics dict as a JSONL line to extraction_metrics.jsonl."""
    try:
        with (debug_dir / "extraction_metrics.jsonl").open("a", encoding="utf-8") as fh:
            fh.write(json.dumps(metrics) + "\n")
    except OSError as exc:
        logger.warning("Could not write metrics: %s", exc)