core/document_processor.py

"""
Advanced document preprocessing pipeline.

Goes beyond basic text extraction to handle:
  - Table extraction from PDFs and DOCX (preserves structure)
  - Code block detection and special handling
  - Section/heading hierarchy extraction
  - Automatic document summarization at ingest time
  - Language detection and filtering
  - PII detection warnings (emails, phone numbers, SSNs)
  - Document fingerprinting for change detection

These preprocessing steps dramatically improve retrieval quality by
ensuring chunks have clean, well-structured text without garbled
table cells or lost context from code blocks.
"""

from __future__ import annotations

import hashlib
import logging
import re
from collections.abc import Callable
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)


# ── Document analysis result ──────────────────────────────────────────────────


@dataclass
class DocumentAnalysis:
    """Rich metadata extracted during document preprocessing."""

    source: str
    total_chars: int = 0
    total_words: int = 0
    language: str = "unknown"
    has_tables: bool = False
    has_code_blocks: bool = False
    detected_sections: list[str] = field(default_factory=list)
    pii_warnings: list[str] = field(default_factory=list)
    summary: str = ""
    document_fingerprint: str = ""
    quality_score: float = 1.0  # 0-1, penalized for garbled text


# ── Text quality assessment ───────────────────────────────────────────────────


def assess_text_quality(text: str) -> float:
    """
    Score text quality from 0.0 to 1.0.

    Penalizes:
      - High ratio of non-ASCII characters (garbled PDF extraction)
      - Excessive whitespace / line breaks
      - Very short paragraphs (table artifacts)
      - Repeated characters (OCR noise)

    Returns:
        Quality score (1.0 = clean, 0.0 = likely garbled)
    """
    if not text:
        return 0.0

    total = len(text)
    non_ascii = sum(1 for c in text if ord(c) > 127)
    ascii_ratio = 1.0 - (non_ascii / total)

    lines = text.splitlines()
    non_empty = [ln for ln in lines if ln.strip()]
    avg_line_len = sum(len(ln) for ln in non_empty) / max(len(non_empty), 1)
    line_quality = min(1.0, avg_line_len / 40)  # penalize very short lines

    # Detect repeated character sequences (OCR artifacts)
    repeat_pattern = re.compile(r"(.)\1{5,}")  # same char 6+ times
    repeats = len(repeat_pattern.findall(text))
    repeat_penalty = max(0.0, 1.0 - repeats * 0.05)

    score = ascii_ratio * 0.5 + line_quality * 0.3 + repeat_penalty * 0.2
    return round(min(1.0, max(0.0, score)), 3)


# ── PII detection ─────────────────────────────────────────────────────────────


_PII_PATTERNS = {
    "email": re.compile(r"[a-zA-Z0-9._%+\-]+@[a-zA-Z0-9.\-]+\.[a-zA-Z]{2,}"),
    "phone_us": re.compile(r"\b(?:\+1[-.\s]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b"),
    "ssn": re.compile(r"\b\d{3}-\d{2}-\d{4}\b"),
    "credit_card": re.compile(r"\b(?:\d{4}[-\s]?){3}\d{4}\b"),
}


def detect_pii(text: str) -> list[str]:
    """
    Detect potential PII in document text.

    Returns a list of warning strings (types found), not the actual values.
    Use this to warn users before ingesting sensitive documents.
    """
    warnings = []
    for pii_type, pattern in _PII_PATTERNS.items():
        if pattern.search(text):
            count = len(pattern.findall(text))
            warnings.append(f"{pii_type}: ~{count} occurrences")
    return warnings


# ── Section/heading extraction ────────────────────────────────────────────────


def extract_sections(text: str) -> list[str]:
    """
    Extract section headings from the document.

    Detects:
      - Markdown headings (# Heading)
      - ALL-CAPS lines (common in PDFs and legal docs)
      - Numbered sections (1. Introduction, 2.1 Background)

    Returns:
        List of detected section titles
    """
    sections = []
    patterns = [
        re.compile(r"^#{1,4}\s+(.+)$", re.MULTILINE),  # Markdown
        re.compile(r"^([A-Z][A-Z\s]{5,60})$", re.MULTILINE),  # ALL-CAPS
        re.compile(r"^\d+(?:\.\d+)*\.?\s+([A-Z][^\n]{5,60})$", re.MULTILINE),  # Numbered
    ]
    for pattern in patterns:
        for match in pattern.finditer(text):
            title = match.group(1).strip()
            if title and title not in sections:
                sections.append(title)
    return sections[:50]  # cap at 50


# ── Table extraction ──────────────────────────────────────────────────────────


def extract_markdown_tables(text: str) -> list[str]:
    """
    Extract markdown-formatted tables from text.

    Returns each table as a clean string block.
    """
    table_pattern = re.compile(
        r"(\|[^\n]+\|\n\|[-:| ]+\|\n(?:\|[^\n]+\|\n)+)",
        re.MULTILINE,
    )
    return [m.group(1).strip() for m in table_pattern.finditer(text)]


def extract_code_blocks(text: str) -> list[tuple[str, str]]:
    """
    Extract code blocks from markdown/text.

    Returns list of (language, code) tuples.
    """
    pattern = re.compile(r"```(\w*)\n(.*?)```", re.DOTALL)
    return [(m.group(1) or "text", m.group(2).strip()) for m in pattern.finditer(text)]


def clean_pdf_text(text: str) -> str:
    """
    Clean common PDF extraction artifacts.

    Fixes:
      - Hyphenated line breaks (re-join words split across lines)
      - Excessive whitespace
      - Form feed characters
      - Ligature replacements (fi → fi, fl → fl)
    """
    # Rejoin hyphenated words at line breaks
    text = re.sub(r"(\w)-\n(\w)", r"\1\2", text)
    # Replace form feeds and null bytes
    text = text.replace("\x0c", "\n\n").replace("\x00", "")
    # Fix ligatures
    ligatures = {"fi": "fi", "fl": "fl", "ff": "ff", "ffi": "ffi", "ffl": "ffl"}
    for lig, rep in ligatures.items():
        text = text.replace(lig, rep)
    # Normalize multiple spaces
    text = re.sub(r" {3,}", "  ", text)
    # Normalize multiple blank lines
    text = re.sub(r"\n{4,}", "\n\n\n", text)
    return text.strip()


# ── Document fingerprinting ───────────────────────────────────────────────────


def fingerprint_document(text: str) -> str:
    """
    Generate a document-level SHA-256 fingerprint for change detection.

    Useful for detecting when a document has been updated and needs re-ingestion.
    """
    return hashlib.sha256(text.encode("utf-8")).hexdigest()[:32]


# ── Language detection (simple, no external deps) ────────────────────────────


def detect_language(text: str) -> str:
    """
    Lightweight language detection using character frequency analysis.

    Returns ISO 639-1 language code ("en", "es", "fr", etc.) or "unknown".
    Only detects major languages — use langdetect for production accuracy.
    """
    sample = text[:500].lower()

    # Common function words by language
    language_hints = {
        "en": ["the", "and", "is", "in", "of", "to", "a", "that"],
        "es": ["el", "la", "de", "que", "y", "en", "los", "se"],
        "fr": ["le", "la", "de", "et", "en", "les", "des", "que"],
        "de": ["der", "die", "das", "und", "ist", "in", "den", "von"],
        "pt": ["o", "a", "de", "que", "e", "do", "da", "em"],
    }

    words = re.findall(r"\b\w+\b", sample)
    word_set = set(words)
    scores = {
        lang: sum(1 for w in hints if w in word_set) for lang, hints in language_hints.items()
    }
    best = max(scores, key=scores.get)
    return best if scores[best] >= 3 else "unknown"


# ── Document summarization ────────────────────────────────────────────────────


def summarize_document(
    text: str,
    source: str,
    llm_fn: Callable[[str], str],
    max_summary_length: int = 300,
) -> str:
    """
    Generate a concise document summary at ingestion time.

    The summary is stored alongside chunks in ChromaDB metadata, enabling
    "collection overview" queries and improving routing accuracy.

    Args:
        text: full document text (truncated to first 3000 chars for efficiency)
        source: source filename for context
        llm_fn: LLM callable
        max_summary_length: target summary length in words

    Returns:
        Summary string, or empty string if generation fails
    """
    # Use first 3000 chars to keep cost low
    sample = text[:3000]
    prompt = (
        f"Write a concise summary (under {max_summary_length} words) of the following document. "
        f"Include: main topic, key points, and any important entities or figures.\n\n"
        f"Document ({source}):\n{sample}\n\nSummary:"
    )
    try:
        summary = llm_fn(prompt).strip()
        logger.debug("Generated summary for '%s': %d chars", source, len(summary))
        return summary
    except Exception as e:
        logger.warning("Document summarization failed for '%s': %s", source, e)
        return ""


# ── Full analysis pipeline ────────────────────────────────────────────────────


def analyze_document(
    text: str,
    source: str,
    llm_fn: Callable[[str], str] | None = None,
    generate_summary: bool = True,
) -> DocumentAnalysis:
    """
    Run the full document analysis pipeline.

    Performs quality assessment, PII detection, section extraction,
    language detection, and optional LLM-generated summary.

    Args:
        text: raw document text
        source: source filename or URL
        llm_fn: optional LLM callable for summarization
        generate_summary: whether to generate an LLM summary

    Returns:
        DocumentAnalysis with all extracted metadata
    """
    cleaned = clean_pdf_text(text)

    analysis = DocumentAnalysis(
        source=source,
        total_chars=len(cleaned),
        total_words=len(cleaned.split()),
        language=detect_language(cleaned),
        has_tables=bool(extract_markdown_tables(cleaned)),
        has_code_blocks=bool(extract_code_blocks(cleaned)),
        detected_sections=extract_sections(cleaned),
        pii_warnings=detect_pii(cleaned),
        document_fingerprint=fingerprint_document(cleaned),
        quality_score=assess_text_quality(cleaned),
    )

    if analysis.pii_warnings:
        logger.warning("PII detected in '%s': %s", source, ", ".join(analysis.pii_warnings))

    if generate_summary and llm_fn and len(cleaned) > 200:
        analysis.summary = summarize_document(cleaned, source, llm_fn)

    logger.info(
        "Document analysis: '%s' | %d words | lang=%s | quality=%.2f | sections=%d",
        source,
        analysis.total_words,
        analysis.language,
        analysis.quality_score,
        len(analysis.detected_sections),
    )
    return analysis