utils.py

# utils.py — Document Sentinel v2.0
# Three-layer PII detection: Regex (structured) + NER/GLiNER (entities) + LLM (context-aware)
# Verification loop, per-entity explainability, synthetic evaluation, feedback system
# LLM: Groq (dev/speed) or OpenAI (jury/accuracy) via env switch

import re
import os
import json
import time
import uuid
import fitz
import pdfplumber
import pytesseract
from PIL import Image
from typing import List, Tuple, Dict, Optional
from groq import Groq
from openai import OpenAI
from dotenv import load_dotenv
from dataclasses import dataclass, field, asdict
from enum import Enum

load_dotenv()

# ── SpaCy (optional, graceful fallback for NER) ──────────────────────────────
try:
    import spacy
    nlp = spacy.load("en_core_web_sm")
    SPACY_AVAILABLE = True
except Exception:
    SPACY_AVAILABLE = False

# ── GLiNER (primary NER engine) ──────────────────────────────────────────────
GLINER_MODEL = None
GLINER_AVAILABLE = False

def _load_gliner():
    """Lazy-load GLiNER model on first use."""
    global GLINER_MODEL, GLINER_AVAILABLE
    if GLINER_MODEL is not None:
        return GLINER_MODEL
    try:
        from gliner import GLiNER
        # Try PII-specific model first, fall back to general multi-PII
        for model_name in [
            "urchade/gliner_multi_pii-v1",
            "knowledgator/gliner-pii-base-v1.0",
            "urchade/gliner_multi-v2.1",
        ]:
            try:
                GLINER_MODEL = GLiNER.from_pretrained(model_name)
                GLINER_AVAILABLE = True
                print(f"[NER] Loaded GLiNER model: {model_name}")
                return GLINER_MODEL
            except Exception:
                continue
        print("[NER] No GLiNER model available — falling back to spaCy/skip")
        return None
    except ImportError:
        print("[NER] GLiNER not installed — falling back to spaCy/skip")
        return None


# ── LLM Client Setup ─────────────────────────────────────────────────────────
LLM_PROVIDER   = os.getenv("LLM_PROVIDER", "groq").lower()
GROQ_API_KEY   = os.getenv("GROQ_API_KEY", "")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")

GROQ_MODEL   = "llama-3.1-8b-instant"
OPENAI_MODEL = "gpt-4o"
CHUNK_TOKENS = 800


# ── Detection Layer Enum ─────────────────────────────────────────────────────
class DetectionMethod(str, Enum):
    REGEX        = "regex"
    NER          = "ner"
    LLM          = "llm"
    VERIFICATION = "verification"


# ── In-memory stores for metrics/feedback ────────────────────────────────────
PIPELINE_METRICS = {
    "total_processed": 0,
    "total_time_ms": 0,
    "layer_counts": {"regex": 0, "ner": 0, "llm": 0, "verification": 0},
    "scores": [],
    "entity_type_counts": {},
}
PROCESSED_DOCS: Dict[str, Dict] = {}   # doc_id → {entities, masked_text, score, ...}
FEEDBACK_LOG: List[Dict] = []          # accumulated feedback for self-improvement


# ══════════════════════════════════════════════════════════════════════════════
#  REGEX PATTERNS — UNTOUCHED FROM ORIGINAL (except minor additions)
# ══════════════════════════════════════════════════════════════════════════════

# ──────────────────────────────────────────────────────────────────────────────
# REGEX RESPONSIBILITY: Only rigid, structural, validatable formats.
# High precision (>95%), near-zero false positives.
#
# REMOVED from regex → moved to NER or LLM:
#   CVV          → was matching every 3-4 digit number (73 false positives)
#   EXPIRY       → was matching every MM/YY date fragment (22 false positives)
#   GENDER       → needs context, not pattern ("Male" in a form vs "Male connector")
#   RELIGION     → needs context ("Christian" as name vs religion)
#   AGE          → needs context ("age 5" could be product version)
#   BLOOD_TYPE   → too short, high FP ("A+" as grade vs blood type)
#   PENALTY_AMOUNT → needs NDA context to distinguish from regular amounts
#   FINANCIAL_AMOUNT → too broad, LLM should decide if amount is sensitive
#   DOB/DOB_ISO  → bare dates are ambiguous; only keyword-anchored DOB stays
# ──────────────────────────────────────────────────────────────────────────────

PII_REGEX = {
    # ── TIER 1: Near-perfect precision (>99%) — unique structural formats ──
    "EMAIL":          r'\b[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}\b',
    "SSN":            r'\b\d{3}-\d{2}-\d{4}\b',
    "CREDIT_CARD":    r'\b\d{4}[\s\-]\d{4}[\s\-]\d{4}[\s\-]\d{4}\b',
    "IP_ADDRESS":     r'\b(?:(?:25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(?:25[0-5]|2[0-4]\d|[01]?\d\d?)\b',
    "MAC_ADDRESS":    r'\b([0-9A-Fa-f]{2}[:\-]){5}[0-9A-Fa-f]{2}\b',
    "URL":            r'(https?://[^\s<>"\']+)',

    # ── TIER 2: High precision (>95%) — keyword-anchored patterns ──
    "SESSION_TOKEN":  r'(user_session|session_id|auth_token|cookie)=[^\s;]+',
    "API_KEY":        r'(?i)(api[_\-]?key|secret|token)\s*[=:]\s*[\w\-]{16,}',
    "PHONE":          r'\b(\+?1?\s?)?(\(?\d{3}\)?[\s.\-]?\d{3}[\s.\-]?\d{4})\b',
    "AADHAR":         r'\b\d{4}[\s\-]?\d{4}[\s\-]?\d{4}\b',
    "PAN":            r'\b[A-Z]{5}[0-9]{4}[A-Z]\b',
    "SWIFT_BIC":      r'\b[A-Z]{4}[A-Z]{2}[A-Z0-9]{2}(?:[A-Z0-9]{3})?\b',
    "MED_RECORD":     r'\b\d{4}-\d{2}-\d{2}-\d{2}\b',

    # ── TIER 3: Keyword-anchored only (require context word nearby) ──
    "DOB":            r'(?i)(?:d\.?o\.?b\.?|date\s*of\s*birth|born|birthday)[\s:]*(\d{1,2}[/\-]\d{1,2}[/\-]\d{2,4})',
    "PASSPORT":       r'(?i)(?:passport)[\s#:]*([A-Z]\d{7,8})',
    "DRIVERS_LICENSE": r"(?i)(?:driver'?s?\s*(?:license|licence|lic))[\s#:]*([A-Z0-9]{5,15})",
    "BANK_ACCOUNT":   r'(?i)(?:account|acct|a/c)[\s#:]*(\d{8,17})',
    "ROUTING_NUMBER":  r'(?i)(?:routing|aba)[\s#:]*(\d{9})',
}

REGEX_ORDER = [
    # Tier 1: longest/most unique first to prevent partial overlaps
    "SESSION_TOKEN", "API_KEY",
    "EMAIL",
    "MAC_ADDRESS",
    "CREDIT_CARD", "SSN", "AADHAR", "MED_RECORD",
    "SWIFT_BIC",
    "PHONE",
    "DOB",
    "IP_ADDRESS",
    "PAN",
    "PASSPORT", "DRIVERS_LICENSE", "BANK_ACCOUNT", "ROUTING_NUMBER",
    "URL",
]

SWIFT_FALSE_POSITIVES = {
    "PERSONAL","INFORMATION","DETAILED","DESCRIPTION","INCLUDES","GUIDANCE",
    "ATTACHED","DOCUMENT","REQUEST","SUPPORT","TYPE","DATE","DAILY","FIRST",
    "LAST","NAME","GENDER","FAITH","ABILITY","BELIEF","LETTER","SPIRITUAL",
    "CHRISTIAN","MEDICAL","HISTORY","PHYSICAL","ADDRESS","NATIONAL","FEDERAL",
}

REASON_MAP = {
    "EMAIL":          "Direct contact identifier — links to a real individual",
    "CREDIT_CARD":    "Financial credential — enables fraud or unauthorised transactions",
    "PHONE":          "Personal contact number — enables direct reach",
    "SSN":            "US Social Security Number — highest identity theft risk",
    "AADHAR":         "Indian national ID — highest identity theft risk",
    "MED_RECORD":     "Medical record number — HIPAA protected health identifier",
    "DOB":            "Date of birth — core identity verification data",
    "MAC_ADDRESS":    "Hardware device identifier — enables physical device tracking",
    "IP_ADDRESS":     "Network address — reveals location and infrastructure",
    "SWIFT_BIC":      "Bank routing code — financial institution exposure",
    "SESSION_TOKEN":  "Active session credential — enables account takeover",
    "API_KEY":        "API credential — enables unauthorised system access",
    "PAN":            "Indian tax ID — financial identity exposure",
    "PASSPORT":       "Travel document identifier — identity theft risk",
    "DRIVERS_LICENSE": "Government ID — identity verification data",
    "BANK_ACCOUNT":   "Financial account number — enables unauthorized access",
    "ROUTING_NUMBER":  "Bank routing number — financial institution identifier",
    "URL":            "Web address — may reveal internal systems or personal accounts",
}


# ══════════════════════════════════════════════════════════════════════════════
#  HELPERS — UNTOUCHED FROM ORIGINAL
# ══════════════════════════════════════════════════════════════════════════════

def is_overlap(start: int, end: int, spans: List[Tuple[int, int]]) -> bool:
    return any(s < end and start < e for s, e in spans)


def _valid_swift(value: str) -> bool:
    if value.upper() in SWIFT_FALSE_POSITIVES:
        return False
    if len(value) not in (8, 11):
        return False
    if not re.search(r'\d', value):
        return False
    return True


def chunk_text(text: str, max_tokens: int = CHUNK_TOKENS) -> List[str]:
    words   = text.split()
    overlap = 50
    chunks  = []
    i = 0
    while i < len(words):
        chunk = words[i: i + max_tokens]
        chunks.append(" ".join(chunk))
        i += max_tokens - overlap
    return chunks


def remove_regex_spans(text, spans):
    masked = text
    for start, end in sorted(spans, reverse=True):
        masked = masked[:start] + " " * (end - start) + masked[end:]
    return masked


# ══════════════════════════════════════════════════════════════════════════════
#  LAYER 1: REGEX — UNTOUCHED LOGIC
# ══════════════════════════════════════════════════════════════════════════════

def regex_detect(text: str) -> Tuple[List[Dict], List[Tuple[int, int]]]:
    entities   = []
    used_spans = []

    for label in REGEX_ORDER:
        pattern = PII_REGEX.get(label)
        if not pattern:
            continue
        for m in re.finditer(pattern, text, re.IGNORECASE):
            start, end = m.start(), m.end()
            value = m.group().strip()

            if label == "SWIFT_BIC" and not _valid_swift(value):
                continue
            if not value or len(value) < 2:
                continue
            if is_overlap(start, end, used_spans):
                continue

            entities.append({
                "entity":     value,
                "label":      label,
                "method":     "regex",
                "confidence": 0.97,
                "reason":     REASON_MAP.get(label, "Sensitive pattern detected"),
                "start":      start,
                "end":        end,
            })
            used_spans.append((start, end))

    return entities, used_spans


# ══════════════════════════════════════════════════════════════════════════════
#  LAYER 2: NER (GLiNER + spaCy fallback) — NEW
# ══════════════════════════════════════════════════════════════════════════════

# GLiNER entity labels to scan for
# ──────────────────────────────────────────────────────────────────────────────
# NER RESPONSIBILITY: Named entities + context-dependent categories
# that regex can't handle without massive false positives.
#
# ADDED (moved from regex):
#   gender, religion, age, blood type, financial amount, date,
#   penalty amount, expiry date
# ──────────────────────────────────────────────────────────────────────────────

GLINER_PII_LABELS = [
    # ── Named entities (NER's core strength) ──
    "person",
    "organization",
    "physical address",
    "street address",
    "city",
    "state",
    "country",
    "company name",
    "legal entity",

    # ── Moved from regex (need context, not pattern) ──
    "gender",
    "religion",
    "age",
    "blood type",
    "financial amount",
    "penalty amount",
    "date",
    "expiry date",
    "nationality",
    "education",
    "medical condition",

    # ── Structural but NER as backup for regex misses ──
    "phone number",
    "email address",
    "date of birth",
    "passport number",
    "driver license number",
    "bank account number",
]

GLINER_LABEL_MAP = {
    "person":              "PERSON_NAME",
    "organization":        "ORGANIZATION",
    "phone number":        "PHONE",
    "email address":       "EMAIL",
    "physical address":    "ADDRESS",
    "street address":      "ADDRESS",
    "city":                "LOCATION",
    "state":               "LOCATION",
    "country":             "LOCATION",
    "date of birth":       "DOB",
    "passport number":     "PASSPORT",
    "driver license number":"DRIVERS_LICENSE",
    "bank account number": "BANK_ACCOUNT",
    "company name":        "ORGANIZATION",
    "legal entity":        "ORGANIZATION",
    "gender":              "GENDER",
    "religion":            "RELIGION",
    "age":                 "AGE",
    "blood type":          "BLOOD_TYPE",
    "financial amount":    "FINANCIAL_AMOUNT",
    "penalty amount":      "PENALTY_AMOUNT",
    "date":                "DATE",
    "expiry date":         "EXPIRY",
    "nationality":         "NATIONALITY",
    "education":           "EDUCATION",
    "medical condition":   "MEDICAL_INFO",
}

SPACY_LABEL_MAP = {
    "PERSON":  "PERSON_NAME",
    "ORG":     "ORGANIZATION",
    "GPE":     "LOCATION",
    "LOC":     "LOCATION",
    "DATE":    "DATE",
    "MONEY":   "FINANCIAL_AMOUNT",
    "FAC":     "LOCATION",
    "NORP":    "NATIONALITY",
}

NER_REASON_MAP = {
    "PERSON_NAME":     "Named individual — direct personal identifier",
    "ORGANIZATION":    "Organisation name — may reveal confidential business relationships",
    "LOCATION":        "Geographic reference — may narrow identity or reveal jurisdiction",
    "ADDRESS":         "Physical address — direct location identifier",
    "NATIONALITY":     "National origin — protected characteristic, discrimination risk",
    "EDUCATION":       "Educational background — personal demographic that aids re-identification",
    "MEDICAL_INFO":    "Health/medical information — HIPAA protected, severe privacy risk",
    "DATE":            "Date reference — may be personally identifying in context",
    "GENDER":          "Gender identity — protected characteristic under anti-discrimination law",
    "RELIGION":        "Religious belief — protected characteristic, discrimination risk if exposed",
    "AGE":             "Age identifier — personal demographic enabling re-identification",
    "BLOOD_TYPE":      "Blood type — protected health information under HIPAA",
    "FINANCIAL_AMOUNT":"Monetary value — may reveal salary, transaction, or contractual terms",
    "PENALTY_AMOUNT":  "Penalty clause amount — corporate confidential NDA term",
    "EXPIRY":          "Expiry date — financial credential component",
    "DOB":             "Date of birth — core identity verification data",
}


def ner_detect(text: str, used_spans: List[Tuple[int, int]]) -> Tuple[List[Dict], List[Tuple[int, int]]]:
    """
    Layer 2: NER-based PII detection (OPTIMIZED).
    - GLiNER primary, spaCy fallback
    - Larger chunks (fewer model calls)
    - Tiered labels (core first, extended only if needed)
    - Fast overlap check via sorted intervals
    """
    entities = []
    new_spans = list(used_spans)

    model = _load_gliner()
    if model is not None:
        entities, new_spans = _gliner_detect(text, model, new_spans)
    elif SPACY_AVAILABLE:
        entities, new_spans = _spacy_detect(text, new_spans)
    else:
        print("[NER] No NER model available — skipping Layer 2")

    return entities, new_spans


# ── Label tiers: Core runs ALWAYS, Extended runs only for longer docs ─────────

GLINER_LABELS_CORE = [
    "person",
    "organization",
    "physical address",
    "city",
    "country",
    "date",
    "financial amount",
    "phone number",
]

GLINER_LABELS_EXTENDED = [
    "gender",
    "religion",
    "age",
    "blood type",
    "penalty amount",
    "expiry date",
    "nationality",
    "education",
    "medical condition",
    "street address",
    "state",
    "company name",
    "legal entity",
    "email address",
    "date of birth",
    "passport number",
    "driver license number",
    "bank account number",
]


def _gliner_detect(
    text: str,
    model,
    used_spans: List[Tuple[int, int]],
) -> Tuple[List[Dict], List[Tuple[int, int]]]:
    """
    Optimized GLiNER detection:
    1. Single chunk up to 8K chars (GLiNER handles ~512 tokens well, that's ~6-8K chars)
    2. One call with core labels; second call with extended labels only if doc > 500 chars
    3. Fast overlap via sorted span set
    """
    entities = []

    # ── Build fast overlap checker from existing spans ────────────────────
    span_set = _SpanSet(used_spans)

    # ── Chunk text — use LARGER chunks to reduce model calls ─────────────
    # GLiNER's internal tokenizer handles up to ~512 tokens (~2500 words)
    # Use 6000 char chunks — most documents fit in 1-2 chunks
    chunks = _chunk_for_ner(text, max_chars=6000, overlap=100)

    for chunk_offset, chunk_text in chunks:
        # ── TIER 1: Core labels (always run) ─────────────────────────
        _run_gliner_batch(model, chunk_text, chunk_offset,
                          GLINER_LABELS_CORE, 0.35,
                          entities, span_set)

        # ── TIER 2: Extended labels (only for docs with enough content) ──
        if len(chunk_text) > 500:
            _run_gliner_batch(model, chunk_text, chunk_offset,
                              GLINER_LABELS_EXTENDED, 0.40,  # slightly higher threshold
                              entities, span_set)

    print(f"[NER] GLiNER detected {len(entities)} entities")
    return entities, span_set.all_spans()


def _run_gliner_batch(
    model,
    chunk_text: str,
    chunk_offset: int,
    labels: List[str],
    threshold: float,
    entities: List[Dict],
    span_set: "_SpanSet",
):
    """Run GLiNER predict on one chunk with given labels. Appends to entities list."""
    try:
        predictions = model.predict_entities(chunk_text, labels, threshold=threshold)
    except Exception as e:
        print(f"[NER] GLiNER error: {e}")
        return

    for pred in predictions:
        pred_text = pred.get("text", "").strip()
        if not pred_text or len(pred_text) < 2:
            continue

        raw_label = pred.get("label", "unknown")
        score     = pred.get("score", 0.5)
        label     = GLINER_LABEL_MAP.get(raw_label, "OTHER_PII")

        rel_start = pred.get("start", 0)
        rel_end   = pred.get("end", rel_start + len(pred_text))
        abs_start = chunk_offset + rel_start
        abs_end   = chunk_offset + rel_end

        # Fast overlap check
        if span_set.overlaps(abs_start, abs_end):
            continue

        entities.append({
            "entity":     pred_text,
            "label":      label,
            "method":     "ner",
            "confidence": round(score, 3),
            "reason":     NER_REASON_MAP.get(label, f"NER model identified as '{raw_label}'"),
            "start":      abs_start,
            "end":        abs_end,
        })
        span_set.add(abs_start, abs_end)


class _SpanSet:
    """
    Fast overlap detection using a sorted list of non-overlapping intervals.
    O(log n) overlap check instead of O(n) linear scan.
    """
    __slots__ = ("_spans",)

    def __init__(self, initial_spans: List[Tuple[int, int]] = None):
        self._spans = sorted(initial_spans or [], key=lambda s: s[0])

    def overlaps(self, start: int, end: int) -> bool:
        # Binary search for the insertion point
        lo, hi = 0, len(self._spans)
        while lo < hi:
            mid = (lo + hi) // 2
            if self._spans[mid][1] <= start:
                lo = mid + 1
            else:
                hi = mid
        # Check the span at lo and lo-1
        if lo < len(self._spans) and self._spans[lo][0] < end:
            return True
        if lo > 0 and self._spans[lo - 1][1] > start:
            return True
        return False

    def add(self, start: int, end: int):
        # Insert in sorted position (bisect)
        lo, hi = 0, len(self._spans)
        while lo < hi:
            mid = (lo + hi) // 2
            if self._spans[mid][0] < start:
                lo = mid + 1
            else:
                hi = mid
        self._spans.insert(lo, (start, end))

    def all_spans(self) -> List[Tuple[int, int]]:
        return list(self._spans)


def _spacy_detect(
    text: str,
    used_spans: List[Tuple[int, int]],
) -> Tuple[List[Dict], List[Tuple[int, int]]]:
    """Fallback NER using spaCy — optimized with truncation and fast overlap."""
    entities = []
    span_set = _SpanSet(used_spans)

    # spaCy is fast but don't feed it a novel — cap at 20K chars
    doc = nlp(text[:20000])

    for ent in doc.ents:
        label = SPACY_LABEL_MAP.get(ent.label_)
        if not label:
            continue
        if len(ent.text.strip()) < 2:
            continue
        if span_set.overlaps(ent.start_char, ent.end_char):
            continue

        entities.append({
            "entity":     ent.text.strip(),
            "label":      label,
            "method":     "ner",
            "confidence": 0.80,
            "reason":     NER_REASON_MAP.get(label, f"NER: {ent.label_}"),
            "start":      ent.start_char,
            "end":        ent.end_char,
        })
        span_set.add(ent.start_char, ent.end_char)

    print(f"[NER] spaCy detected {len(entities)} entities")
    return entities, span_set.all_spans()


def _chunk_for_ner(
    text: str, max_chars: int = 4000, overlap: int = 200
) -> List[Tuple[int, str]]:
    """Split text into overlapping chunks. Returns (offset, chunk_text) pairs."""
    if len(text) <= max_chars:
        return [(0, text)]
    chunks = []
    start = 0
    while start < len(text):
        end = min(start + max_chars, len(text))
        # Break at paragraph/sentence boundary
        if end < len(text):
            for sep in ["\n\n", "\n", ". ", "! ", "? "]:
                brk = text.rfind(sep, max(start + max_chars - 300, start), end)
                if brk > start + max_chars // 2:
                    end = brk + len(sep)
                    break
        chunks.append((start, text[start:end]))
        start = end - overlap if end < len(text) else end
    return chunks


# ══════════════════════════════════════════════════════════════════════════════
#  LAYER 3: LLM — SINGLE UNIFIED CALL
#  Replaces 3 separate calls (detect + reason + verify) with 1 call.
#  Before: detect(2-4s) + reason(2-3s) + verify(2-3s) = 6-10s
#  After:  unified(2-4s) = 2-4s total
# ══════════════════════════════════════════════════════════════════════════════

LLM_UNIFIED_PROMPT = """You are a privacy and confidentiality expert AI performing THREE tasks in ONE pass.

The document has ALREADY been scanned by:
  - Regex (caught: emails, SSNs, credit cards, phones, IPs, URLs, API keys)
  - NER model (caught: person names, organisations, addresses, locations, dates, gender, religion, ages, financial amounts)

You will receive:
  A) The original document text
  B) List of entities already detected by regex and NER

You must do ALL THREE tasks and return ONE JSON response:

━━━ TASK 1: DETECT MISSED ENTITIES ━━━
Find sensitive information that regex and NER missed:
- NDA/contract specifics: party roles, jurisdiction, governing law, arbitration
- Implicit identifiers: unique attribute combos that could identify someone
- Written-out numbers: "two million dollars", "fifteen years"
- Relationship references: "his wife", "her employer"
- Medical context: diagnoses, treatments, prescriptions
- Proprietary info: trade secrets, codenames
- Contextual sensitivity: data sensitive ONLY because of surrounding text

━━━ TASK 2: EXPLAIN EVERY ENTITY ━━━
For ALL entities (both already-detected AND your new ones), provide a
context-specific reason WHY it is sensitive in THIS document.
BAD: "This is a phone number and phone numbers are PII"
GOOD: "This phone number is listed as the primary contact for the NDA signatory, directly linking it to a named party in a confidential agreement"

━━━ TASK 3: VERIFY COMPLETENESS ━━━
After considering all entities, rate the sanitization completeness.

━━━ RESPONSE FORMAT ━━━
Return ONLY this JSON structure:
{
  "new_entities": [
    {"entity": "<exact text>", "label": "<SHORT_LABEL>", "reason": "<context reason>", "confidence": 0.0-1.0}
  ],
  "entity_reasons": [
    {"entity": "<text>", "label": "<LABEL>", "reason": "<context-specific reason>"}
  ],
  "verification": {
    "sanitization_score": 0-100,
    "is_clean": true/false,
    "issues": ["any remaining concerns"]
  }
}

If no new entities found, use empty array for new_entities.
entity_reasons should cover ALL entities — the ones already detected AND your new ones."""


def _get_client():
    """Get reusable LLM client + model name."""
    if LLM_PROVIDER == "groq":
        return Groq(api_key=GROQ_API_KEY), GROQ_MODEL
    else:
        return OpenAI(api_key=OPENAI_API_KEY), OPENAI_MODEL


def llm_unified_call(
    text: str,
    existing_entities: List[Dict],
    used_spans: List[Tuple[int, int]],
) -> Tuple[List[Dict], Dict[Tuple[str,str], str], Dict]:
    """
    SINGLE LLM call that replaces detect + reason + verify.

    Returns:
      new_entities  — list of newly detected entities (by LLM)
      reason_map    — {(entity_text, label): reason} for ALL entities
      verification  — {"sanitization_score": N, "is_clean": bool, "issues": [...]}
    """
    if not GROQ_API_KEY and not OPENAI_API_KEY:
        print("[LLM] No API key — skipping LLM layer")
        return [], {}, {"sanitization_score": 75, "is_clean": True, "issues": ["LLM skipped: no API key"]}

    client, model = _get_client()

    # Build existing entity summary (compact)
    existing_summary = ""
    if existing_entities:
        seen = set()
        lines = []
        for e in existing_entities:
            key = f"{e['label']}: {e['entity']}"
            if key not in seen:
                seen.add(key)
                lines.append(f"  [{e['method']}] {key}")
        existing_summary = "\n".join(lines[:40])  # Cap at 40 to save tokens

    # Truncate text for token limits
    max_text = 8000
    text_sample = text[:max_text]

    user_message = f"""Already detected entities:
{existing_summary if existing_summary else "None yet."}

--- DOCUMENT TEXT ---
{text_sample}
--- END ---

Perform all 3 tasks: detect missed entities, explain every entity, verify completeness. Return JSON only."""

    try:
        kwargs = {
            "model": model,
            "messages": [
                {"role": "system", "content": LLM_UNIFIED_PROMPT},
                {"role": "user", "content": user_message},
            ],
            "temperature": 0.1,
            "max_tokens": 2000,
        }
        # OpenAI supports response_format
        if LLM_PROVIDER != "groq":
            kwargs["response_format"] = {"type": "json_object"}

        response = client.chat.completions.create(**kwargs)
        raw = response.choices[0].message.content.strip()
        raw = raw.replace("```json", "").replace("```", "").strip()
        parsed = json.loads(raw)

    except Exception as e:
        print(f"[LLM UNIFIED ERROR]: {e}")
        return [], {}, {"sanitization_score": 75, "is_clean": True, "issues": [str(e)]}

    # ── Parse new entities ────────────────────────────────────────────────
    new_entities = []
    span_set = _SpanSet(used_spans)
    seen_vals = {e["entity"].lower() for e in existing_entities}

    for item in parsed.get("new_entities", []):
        val = (item.get("entity") or "").strip()
        if not val or len(val) < 3 or val.lower() in seen_vals:
            continue

        idx = text.find(val)
        if idx == -1:
            idx = text.lower().find(val.lower())
            if idx == -1:
                continue
        end_idx = idx + len(val)

        if span_set.overlaps(idx, end_idx):
            continue

        label = (item.get("label") or "SENSITIVE").upper().replace(" ", "_")
        new_entities.append({
            "entity":     val,
            "label":      label,
            "method":     "llm",
            "confidence": float(item.get("confidence", 0.80)),
            "reason":     item.get("reason", "Detected by contextual AI analysis"),
            "start":      idx,
            "end":        end_idx,
        })
        span_set.add(idx, end_idx)
        seen_vals.add(val.lower())

    # ── Parse reason map ──────────────────────────────────────────────────
    reason_map = {}
    for item in parsed.get("entity_reasons", []):
        ent_text = (item.get("entity") or "").strip()
        label    = (item.get("label") or "").strip()
        reason   = (item.get("reason") or "").strip()
        if ent_text and reason:
            reason_map[(ent_text, label)] = reason

    # ── Parse verification ────────────────────────────────────────────────
    v = parsed.get("verification", {})
    verification = {
        "sanitization_score": float(v.get("sanitization_score", 80)),
        "is_clean":           v.get("is_clean", True),
        "missed":             [],  # Already captured in new_entities
        "issues":             v.get("issues", []),
    }

    print(f"[LLM] Unified call: {len(new_entities)} new entities, "
          f"{len(reason_map)} reasons, score={verification['sanitization_score']}")

    return new_entities, reason_map, verification


# ── Keep old functions as thin wrappers for backward compatibility ─────────

def llm_detect(chunk: str) -> List[Dict]:
    """Legacy wrapper — use llm_unified_call instead."""
    try:
        client, model = _get_client()
        messages = [
            {"role": "system", "content": LLM_UNIFIED_PROMPT.split("━━━ TASK 1")[1].split("━━━ TASK 2")[0]},
            {"role": "user", "content": f"Find sensitive data:\n\n{chunk}"},
        ]
        response = client.chat.completions.create(
            model=model, messages=messages, temperature=0.1, max_tokens=800,
        )
        raw = response.choices[0].message.content.strip()
        parsed = json.loads(raw)
        if isinstance(parsed, list): return parsed
        for val in parsed.values():
            if isinstance(val, list): return val
        return []
    except Exception as e:
        print(f"[LLM ERROR]: {e}")
        return []


def llm_detect_full(text, used_spans):
    """Legacy wrapper — called from evaluate.py when running without unified call."""
    clean_text = remove_regex_spans(text, used_spans)
    chunks = chunk_text(clean_text)
    all_ents = []
    new_spans = list(used_spans)
    seen = set()
    for chunk in chunks:
        for e in llm_detect(chunk):
            val = (e.get("entity") or "").strip()
            if not val or val in seen or len(val) < 3:
                continue
            idx = text.find(val)
            if idx == -1:
                idx = text.lower().find(val.lower())
                if idx == -1: continue
            end_idx = idx + len(val)
            if is_overlap(idx, end_idx, new_spans):
                continue
            seen.add(val)
            label = (e.get("label") or "SENSITIVE").upper().replace(" ", "_")
            all_ents.append({
                "entity": val, "label": label, "method": "llm",
                "confidence": float(e.get("confidence", 0.80)),
                "reason": e.get("reason", "Detected by AI"),
                "start": idx, "end": end_idx,
            })
            new_spans.append((idx, end_idx))
    print(f"[LLM] Detected {len(all_ents)} entities")
    return all_ents, new_spans


def llm_generate_reasons(entities, document_text=""):
    """Legacy no-op — reasoning is now done inside unified call."""
    return []


def verify_redaction(masked_text, entities):
    """Legacy no-op — verification is now done inside unified call."""
    return {"sanitization_score": 85, "is_clean": True, "missed": [], "issues": []}


# ══════════════════════════════════════════════════════════════════════════════
#  DOCUMENT CLASSIFICATION — NEW
# ══════════════════════════════════════════════════════════════════════════════

def classify_document(text: str) -> str:
    """Classify document as nda/transcript/paperwork/unknown using heuristics + LLM."""
    lower = text.lower()
    nda_kw = ["non-disclosure", "confidential", "nda", "agreement",
              "governing law", "jurisdiction", "liquidated damages", "party"]
    transcript_kw = ["speaker", "interviewer", "q:", "a:", "said", "replied",
                     "conversation", "transcript", "recording"]

    nda_score = sum(1 for kw in nda_kw if kw in lower)
    transcript_score = sum(1 for kw in transcript_kw if kw in lower)

    if nda_score >= 3:
        return "nda"
    if transcript_score >= 3:
        return "transcript"
    return "paperwork"


# ══════════════════════════════════════════════════════════════════════════════
#  MERGE ALL LAYERS — UPGRADED
# ══════════════════════════════════════════════════════════════════════════════

def merge_entities(
    regex_ents: List[Dict],
    ner_ents: List[Dict],
    llm_ents: List[Dict],
    text: str,
) -> List[Dict]:
    """
    Merge entities from all three layers.
    Priority: Regex > NER > LLM for overlapping spans.
    All entities get character positions.
    """
    seen_values = set()
    merged = []

    # Layer priority order
    for source_ents in [regex_ents, ner_ents, llm_ents]:
        for ent in source_ents:
            val = ent["entity"].strip()
            val_lower = val.lower()
            if val_lower in seen_values:
                continue
            if len(val) < 2:
                continue

            # Ensure positions exist
            if "start" not in ent or ent["start"] is None:
                idx = text.find(val)
                if idx == -1:
                    continue
                ent["start"] = idx
                ent["end"]   = idx + len(val)

            merged.append(ent)
            seen_values.add(val_lower)

    merged.sort(key=lambda x: x.get("start", 0))
    return merged


# ══════════════════════════════════════════════════════════════════════════════
#  METRICS COMPUTATION — NEW
# ══════════════════════════════════════════════════════════════════════════════

def compute_metrics(gt: List[Dict], pred: List[Dict]) -> Dict:
    """Compute precision, recall, F1 comparing predictions to ground truth."""
    def norm(e):
        return (e["entity"].strip().lower(), e["label"].upper())

    gt_set   = set(norm(e) for e in gt)
    pred_set = set(norm(e) for e in pred)

    TP = len(gt_set & pred_set)
    FP = len(pred_set - gt_set)
    FN = len(gt_set - pred_set)

    precision = TP / (TP + FP) if TP + FP else 0
    recall    = TP / (TP + FN) if TP + FN else 0
    f1        = 2 * precision * recall / (precision + recall) if precision + recall else 0

    return {
        "precision": round(precision, 4),
        "recall":    round(recall, 4),
        "f1":        round(f1, 4),
        "tp": TP, "fp": FP, "fn": FN,
    }


def compute_layer_metrics(entities: List[Dict]) -> Dict:
    """Compute per-layer contribution breakdown."""
    layer_counts = {"regex": 0, "ner": 0, "llm": 0, "verification": 0}
    layer_types  = {"regex": {}, "ner": {}, "llm": {}, "verification": {}}
    layer_conf   = {"regex": [], "ner": [], "llm": [], "verification": []}

    for e in entities:
        method = e.get("method", "unknown")
        if method in layer_counts:
            layer_counts[method] += 1
            label = e.get("label", "UNKNOWN")
            layer_types[method][label] = layer_types[method].get(label, 0) + 1
            layer_conf[method].append(e.get("confidence", 0.0))

    total = sum(layer_counts.values()) or 1
    result = {}
    for layer in ["regex", "ner", "llm", "verification"]:
        confs = layer_conf[layer]
        result[layer] = {
            "count":           layer_counts[layer],
            "percentage":      round(layer_counts[layer] / total * 100, 1),
            "entity_types":    layer_types[layer],
            "avg_confidence":  round(sum(confs) / max(len(confs), 1), 3),
        }

    return result


# ══════════════════════════════════════════════════════════════════════════════
#  SYNTHETIC DATA GENERATOR — NEW (for evaluation without annotations)
# ══════════════════════════════════════════════════════════════════════════════

def generate_synthetic_docs(n: int = 20) -> List[Dict]:
    """
    Generate synthetic documents with known PII for evaluation.
    Returns list of {text, doc_type, ground_truth: [{entity, label}]}.
    """
    import random
    try:
        from faker import Faker
        fake = Faker()
        Faker.seed(42)
    except ImportError:
        print("[EVAL] Faker not installed — cannot generate synthetic data")
        return []

    random.seed(42)
    docs = []

    for i in range(n):
        doc_type = random.choice(["nda", "transcript", "paperwork"])
        if doc_type == "nda":
            docs.append(_synth_nda(fake, random))
        elif doc_type == "transcript":
            docs.append(_synth_transcript(fake, random))
        else:
            docs.append(_synth_paperwork(fake, random))

    return docs


def _synth_nda(fake, random) -> Dict:
    party_a    = fake.company()
    party_b    = fake.company()
    person_a   = fake.name()
    person_b   = fake.name()
    state      = fake.state()
    penalty    = f"${random.randint(10,500) * 1000:,}"
    eff_date   = fake.date_this_year().strftime("%m/%d/%Y")
    email_a    = fake.email()
    phone_a    = f"({random.randint(200,999)}) {random.randint(200,999)}-{random.randint(1000,9999)}"

    text = f"""NON-DISCLOSURE AGREEMENT

This Non-Disclosure Agreement is entered into as of {eff_date} by and between:

Party A: {party_a}, represented by {person_a}
Email: {email_a}  Phone: {phone_a}

Party B: {party_b}, represented by {person_b}

1. CONFIDENTIAL INFORMATION
All proprietary information shared shall be considered Confidential.

2. REMEDIES
Breach penalty: liquidated damages of {penalty}.

3. GOVERNING LAW
Governed by the laws of the State of {state}.

Signed: {person_a}, CEO of {party_a}
Signed: {person_b}, Director of {party_b}
Date: {eff_date}
"""

    gt = []
    for val, label in [
        (party_a, "ORGANIZATION"), (party_b, "ORGANIZATION"),
        (person_a, "PERSON_NAME"), (person_b, "PERSON_NAME"),
        (state, "JURISDICTION"), (penalty, "PENALTY_AMOUNT"),
        (email_a, "EMAIL"), (phone_a, "PHONE"),
        (eff_date, "DATE"),
    ]:
        if val in text:
            gt.append({"entity": val, "label": label})

    return {"text": text, "doc_type": "nda", "ground_truth": gt}


def _synth_transcript(fake, random) -> Dict:
    name_a  = fake.name()
    name_b  = fake.name()
    phone   = f"({random.randint(200,999)}) {random.randint(200,999)}-{random.randint(1000,9999)}"
    email   = fake.email()
    ssn     = f"{random.randint(100,999)}-{random.randint(10,99)}-{random.randint(1000,9999)}"
    dob     = fake.date_of_birth(minimum_age=20, maximum_age=70).strftime("%m/%d/%Y")
    company = fake.company()
    city    = fake.city()

    text = f"""TRANSCRIPT

{name_a}: Hello, my name is {name_a}, calling from {company}.
{name_b}: Hi {name_a.split()[0]}, this is {name_b}. How can I help?
{name_a}: I need to update my phone to {phone}.
{name_b}: Can you confirm your date of birth?
{name_a}: It's {dob}.
{name_b}: And your SSN for verification?
{name_a}: {ssn}
{name_b}: Email on file?
{name_a}: {email}
{name_b}: And you're located in {city}?
{name_a}: That's correct.
"""

    gt = []
    for val, label in [
        (name_a, "PERSON_NAME"), (name_b, "PERSON_NAME"),
        (phone, "PHONE"), (email, "EMAIL"), (ssn, "SSN"),
        (dob, "DOB"), (company, "ORGANIZATION"), (city, "LOCATION"),
    ]:
        if val in text:
            gt.append({"entity": val, "label": label})

    return {"text": text, "doc_type": "transcript", "ground_truth": gt}


def _synth_paperwork(fake, random) -> Dict:
    name    = fake.name()
    ssn     = f"{random.randint(100,999)}-{random.randint(10,99)}-{random.randint(1000,9999)}"
    dob     = fake.date_of_birth(minimum_age=18, maximum_age=80).strftime("%m/%d/%Y")
    email   = fake.email()
    phone   = f"({random.randint(200,999)}) {random.randint(200,999)}-{random.randint(1000,9999)}"
    addr    = fake.street_address()
    company = fake.company()
    ip      = f"{random.randint(1,255)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}"
    cc      = f"{random.randint(4000,4999)} {random.randint(1000,9999)} {random.randint(1000,9999)} {random.randint(1000,9999)}"
    gender  = random.choice(["Male", "Female"])

    text = f"""PERSONAL INFORMATION FORM

Name: {name}
DOB: {dob}
Gender: {gender}
SSN: {ssn}
Email: {email}
Phone: {phone}
Address: {addr}
Employer: {company}
Credit Card: {cc}
IP Address: {ip}
"""

    gt = []
    for val, label in [
        (name, "PERSON_NAME"), (ssn, "SSN"), (dob, "DOB"),
        (email, "EMAIL"), (phone, "PHONE"), (addr, "ADDRESS"),
        (company, "ORGANIZATION"), (cc, "CREDIT_CARD"),
        (ip, "IP_ADDRESS"),
    ]:
        if val in text:
            gt.append({"entity": val, "label": label})

    return {"text": text, "doc_type": "paperwork", "ground_truth": gt}


def run_evaluation(n_docs: int = 20, use_llm: bool = True) -> Dict:
    """
    Full evaluation pipeline:
    1. Generate synthetic docs with known ground truth
    2. Run detection pipeline on each
    3. Compute per-entity-type and per-layer P/R/F1
    Returns detailed evaluation report.
    """
    docs = generate_synthetic_docs(n_docs)
    if not docs:
        return {"error": "Could not generate synthetic data"}

    all_preds  = []
    all_truths = []
    per_doc    = []
    total_time = 0

    for i, doc in enumerate(docs):
        print(f"[EVAL] Processing doc {i+1}/{n_docs} ({doc['doc_type']})...")
        start = time.time()

        # Run detection (without masking/pdf — just text)
        regex_ents, used_spans = regex_detect(doc["text"])
        ner_ents, used_spans   = ner_detect(doc["text"], used_spans)

        if use_llm:
            llm_ents, used_spans = llm_detect_full(doc["text"], used_spans)
        else:
            llm_ents = []

        merged = merge_entities(regex_ents, ner_ents, llm_ents, doc["text"])

        elapsed = (time.time() - start) * 1000
        total_time += elapsed

        # Per-doc metrics
        doc_metrics = compute_metrics(doc["ground_truth"], merged)
        doc_metrics["doc_type"]   = doc["doc_type"]
        doc_metrics["time_ms"]    = round(elapsed, 0)
        doc_metrics["num_pred"]   = len(merged)
        doc_metrics["num_truth"]  = len(doc["ground_truth"])
        doc_metrics["layers"]     = compute_layer_metrics(merged)
        per_doc.append(doc_metrics)

        all_preds.extend(merged)
        all_truths.extend(doc["ground_truth"])

    # Overall metrics
    overall = compute_metrics(all_truths, all_preds)
    overall_layers = compute_layer_metrics(all_preds)

    # Per entity-type breakdown
    type_metrics = {}
    for etype in set(e["label"] for e in all_truths):
        type_gt   = [e for e in all_truths if e["label"] == etype]
        type_pred = [e for e in all_preds if e.get("label") == etype]
        type_metrics[etype] = compute_metrics(type_gt, type_pred)

    return {
        "overall":          overall,
        "per_entity_type":  type_metrics,
        "layer_breakdown":  overall_layers,
        "per_document":     per_doc,
        "summary": {
            "total_docs":       n_docs,
            "total_time_ms":    round(total_time, 0),
            "avg_time_per_doc": round(total_time / n_docs, 0),
            "layers_used":      ["regex", "ner"] + (["llm"] if use_llm else []),
        },
    }


# ══════════════════════════════════════════════════════════════════════════════
#  MAIN ENTRY POINT — UPGRADED (3 layers + verification + feedback)
# ══════════════════════════════════════════════════════════════════════════════

def detect_and_mask(
    text: str,
    enable_verification: bool = True,
    enable_ner: bool = True,
) -> Tuple[str, List[Dict], Dict]:
    """
    Full three-layer hybrid pipeline with unified LLM call.

    OPTIMIZED: Detection + Reasoning + Verification in 1 LLM call.
    Before: 3 sequential calls = 6-10s
    After:  1 unified call     = 2-4s

    Returns:
      masked_text  — original text with PII replaced by [LABEL] placeholders
      entity_list  — list of all detected entities with label, reason, method, confidence
      report       — dict with layer_breakdown, verification result, doc_type, metrics
    """
    start_time = time.time()
    doc_id     = str(uuid.uuid4())[:8]

    # ── Classify document ─────────────────────────────────────────────────
    doc_type = classify_document(text)
    print(f"[{doc_id}] Document type: {doc_type}")

    # ── Layer 1: Regex ────────────────────────────────────────────────────
    regex_ents, used_spans = regex_detect(text)
    print(f"[{doc_id}] Regex: {len(regex_ents)} entities")

    # ── Layer 2: NER ──────────────────────────────────────────────────────
    ner_ents = []
    if enable_ner:
        ner_ents, used_spans = ner_detect(text, used_spans)
        print(f"[{doc_id}] NER: {len(ner_ents)} entities")

    # ── Layer 3: UNIFIED LLM call (detect + reason + verify) ──────────────
    # Merge regex+ner first so LLM sees what's already found
    pre_merge = merge_entities(regex_ents, ner_ents, [], text)

    llm_ents, reason_map, verification = llm_unified_call(
        text, pre_merge, used_spans
    )
    print(f"[{doc_id}] LLM: {len(llm_ents)} new entities")

    # ── Final merge (all 3 layers) ────────────────────────────────────────
    all_entities = merge_entities(regex_ents, ner_ents, llm_ents, text)
    print(f"[{doc_id}] Total after merge: {len(all_entities)} entities")

    # ── Apply LLM reasons to all entities ─────────────────────────────────
    clean_entities = []
    for e in all_entities:
        # Try exact match first, then fuzzy key match
        llm_reason = reason_map.get((e["entity"], e["label"]))
        if not llm_reason:
            # Try matching just by entity text (label might differ slightly)
            for (ent_text, _), reason in reason_map.items():
                if ent_text.lower() == e["entity"].lower():
                    llm_reason = reason
                    break

        clean_entities.append({
            "entity":     e["entity"],
            "label":      e["label"],
            "method":     e["method"],
            "confidence": e.get("confidence", 0.9),
            "reason":     llm_reason if llm_reason else e.get("reason", "Sensitive data"),
        })

    # ── Build masked text ─────────────────────────────────────────────────
    masked = text
    positioned = []
    for e in all_entities:
        idx = masked.find(e["entity"])
        if idx != -1:
            positioned.append((idx, idx + len(e["entity"]), e))

    positioned.sort(key=lambda x: x[0], reverse=True)
    for start, end, e in positioned:
        masked = masked[:start] + f"[{e['label']}]" + masked[end:]

    # ── Handle verification misses (from unified call) ────────────────────
    if verification.get("missed"):
        for missed in verification["missed"]:
            val = missed.get("entity", "").strip()
            if not val or len(val) < 2:
                continue
            if val.lower() in {e["entity"].lower() for e in all_entities}:
                continue
            idx = text.find(val)
            if idx == -1:
                continue
            new_ent = {
                "entity":     val,
                "label":      missed.get("label", "SENSITIVE"),
                "method":     "verification",
                "confidence": 0.85,
                "reason":     missed.get("reason", "Caught by verification"),
            }
            clean_entities.append(new_ent)
            # Re-mask
            m_idx = masked.find(val)
            if m_idx != -1:
                masked = masked[:m_idx] + f"[{new_ent['label']}]" + masked[m_idx + len(val):]

    # ── Compute layer metrics ─────────────────────────────────────────────
    layer_breakdown = compute_layer_metrics(clean_entities)
    elapsed_ms = (time.time() - start_time) * 1000

    # ── Build report ──────────────────────────────────────────────────────
    report = {
        "document_id":       doc_id,
        "document_type":     doc_type,
        "total_entities":    len(clean_entities),
        "extraction_method": "text",
        "layer_breakdown":   layer_breakdown,
        "verification":      verification,
        "processing_time_ms": round(elapsed_ms, 0),
        "sanitization_score": verification.get("sanitization_score", 85.0),
    }

    # ── Update global metrics ─────────────────────────────────────────────
    _update_global_metrics(clean_entities, elapsed_ms, verification)

    # ── Store for feedback retrieval ──────────────────────────────────────
    PROCESSED_DOCS[doc_id] = {
        "entities":    clean_entities,
        "masked_text": masked,
        "report":      report,
    }

    print(f"[{doc_id}] Pipeline complete: {len(clean_entities)} entities in {elapsed_ms:.0f}ms")

    return masked, clean_entities, report

    return masked, clean_entities, report


def _update_global_metrics(entities, elapsed_ms, verification):
    """Update running pipeline metrics."""
    PIPELINE_METRICS["total_processed"] += 1
    PIPELINE_METRICS["total_time_ms"] += elapsed_ms
    if verification:
        PIPELINE_METRICS["scores"].append(verification.get("sanitization_score", 80))
    for e in entities:
        method = e.get("method", "unknown")
        if method in PIPELINE_METRICS["layer_counts"]:
            PIPELINE_METRICS["layer_counts"][method] += 1
        label = e.get("label", "UNKNOWN")
        PIPELINE_METRICS["entity_type_counts"][label] = \
            PIPELINE_METRICS["entity_type_counts"].get(label, 0) + 1


def get_pipeline_metrics() -> Dict:
    """Get aggregated pipeline metrics."""
    m = PIPELINE_METRICS
    total = max(m["total_processed"], 1)
    total_ents = sum(m["layer_counts"].values()) or 1

    return {
        "total_documents_processed": m["total_processed"],
        "avg_sanitization_score": round(
            sum(m["scores"]) / max(len(m["scores"]), 1), 1
        ),
        "avg_processing_time_ms": round(m["total_time_ms"] / total, 0),
        "layer_contribution": {
            layer: f"{count / total_ents * 100:.1f}%"
            for layer, count in m["layer_counts"].items()
        },
        "entity_type_distribution": m["entity_type_counts"],
        "verification_pass_rate": round(
            sum(1 for s in m["scores"] if s >= 90) / max(len(m["scores"]), 1) * 100, 1
        ),
    }


def get_feedback(doc_id: str) -> Dict:
    """Get per-entity feedback for a processed document."""
    if doc_id not in PROCESSED_DOCS:
        return {"error": "Document not found"}
    doc = PROCESSED_DOCS[doc_id]
    return {
        "document_id": doc_id,
        "total_entities": len(doc["entities"]),
        "feedback": doc["entities"],
        "report": doc["report"],
    }


# ══════════════════════════════════════════════════════════════════════════════
#  PDF HELPERS — COMPLETELY UNTOUCHED FROM ORIGINAL
# ══════════════════════════════════════════════════════════════════════════════

def extract_text_from_pdf(pdf_path: str) -> str:
    """Primary: pdfplumber. Fallback: Tesseract OCR."""
    text = ""
    try:
        with pdfplumber.open(pdf_path) as pdf:
            text = "\n".join(p.extract_text() or "" for p in pdf.pages)
    except Exception:
        pass

    if len(text.strip()) < 30:
        try:
            doc   = fitz.open(pdf_path)
            parts = []
            for page in doc:
                pix  = page.get_pixmap(dpi=300)
                img  = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                parts.append(pytesseract.image_to_string(img))
            text = "\n".join(parts)
        except Exception:
            pass

    return text.strip()


def redact_pdf(input_path: str, output_path: str, entities: List[Dict]) -> None:
    """
    White-out each entity in the original PDF and
    write its [LABEL] placeholder in dark red at the same position.
    """
    doc = fitz.open(input_path)

    for page in doc:
        span_map = []
        for block in page.get_text("dict")["blocks"]:
            if block["type"] != 0:
                continue
            for line in block["lines"]:
                for span in line["spans"]:
                    span_map.append((fitz.Rect(span["bbox"]), span["size"]))

        def font_at(rect: fitz.Rect) -> float:
            for sr, sz in span_map:
                if sr.intersects(rect):
                    return sz
            return 9.0

        for ent in entities:
            value = (ent.get("entity") or "").strip()
            if not value or len(value) < 2:
                continue

            for rect in page.search_for(value):
                fs = min(max(font_at(rect) - 1, 6), 9)
                page.draw_rect(rect, color=(1,1,1), fill=(1,1,1))
                page.insert_text(
                    (rect.x0, rect.y1 - 1),
                    f"[{ent['label']}]",
                    fontsize=fs,
                    color=(0.85, 0, 0),
                )

    doc.save(output_path)
    doc.close()