"""Heuristic privacy layout detector for images.

Uses OCR word-box geometry (from pytesseract image_to_data) and OCR text
patterns to classify regions of an image as containing sensitive document
structures, without requiring an ML layout model.

Three heuristics are implemented:

1. **MRZ (Machine Readable Zone)** — passport / travel document
   - Looks for lines where ≥70% of characters are from the MRZ alphabet
     (A-Z, 0-9, <) and the line spans nearly full image width.
   - MRZ lines are 44 chars (TD3 passport) or 36 chars (TD2) wide.

2. **ID-card layout** — national identity cards
   - Detects a dense text block on the right half of the image
     combined with a blank/sparse region on the left (photo placeholder).
   - Also checks for ID-specific keyword density in OCR text.

3. **Table / bank-statement layout** — structured tabular PII
   - Detects a grid of horizontally aligned word clusters (≥3 rows
     of ≥2 words with similar x-alignment and consistent row spacing).
   - Also checks for financial keywords (IBAN, amount, balance, etc.)
     in the OCR text near the suspected table region.

Each heuristic returns a ``LayoutSignal`` dataclass describing:
- ``layout_type``: "mrz" | "id_card" | "table"
- ``confidence``: 0.0–1.0
- ``bounding_box``: (x, y, w, h) of the detected region in image pixels
- ``ocr_text_spans``: char offsets into the OCR text that are within the region
"""
from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple


@dataclass
class LayoutSignal:
    layout_type: str            # "mrz" | "id_card" | "table"
    confidence: float           # 0.0–1.0
    bounding_box: Tuple[int, int, int, int]  # (x, y, w, h)
    ocr_text_spans: List[Tuple[int, int]] = field(default_factory=list)  # [(start, end), …]


# ---------------------------------------------------------------------------
# MRZ detection
# ---------------------------------------------------------------------------

_MRZ_ALPHA = re.compile(r"^[A-Z0-9<\s]+$")
# MRZ lines are typically ≥30 characters of MRZ alphabet
_MRZ_MIN_LEN = 25


def _mrz_line_score(line_text: str) -> float:
    """Return fraction of MRZ-alphabet chars in *line_text*."""
    if not line_text:
        return 0.0
    mrz_chars = sum(1 for c in line_text if c in "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789< ")
    return mrz_chars / len(line_text)


def detect_mrz(word_boxes: List[Dict], ocr_text: str, image_width: int) -> Optional[LayoutSignal]:
    """Detect MRZ zone from word boxes.

    Groups word boxes by line (same top±5px), checks each line for MRZ
    characteristics, looks for 2–3 consecutive MRZ lines.
    """
    if not word_boxes or image_width == 0:
        return None

    # Group boxes by approximate line (cluster by top coordinate)
    lines: Dict[int, List[Dict]] = {}
    for box in word_boxes:
        top = box["top"]
        # Round to nearest 10px bucket
        bucket = round(top / 10) * 10
        lines.setdefault(bucket, []).append(box)

    mrz_lines = []
    for top_bucket in sorted(lines):
        line_boxes = sorted(lines[top_bucket], key=lambda b: b["left"])
        # Reconstruct line text from word box texts (use char offsets if available)
        line_text = " ".join(b.get("word", "") for b in line_boxes).upper()
        if len(line_text.replace(" ", "")) < _MRZ_MIN_LEN:
            continue
        score = _mrz_line_score(line_text.replace(" ", ""))
        # Require ≥80% MRZ chars AND line spans most of image width
        if score >= 0.80:
            leftmost = min(b["left"] for b in line_boxes)
            rightmost = max(b["left"] + b["width"] for b in line_boxes)
            width_fraction = (rightmost - leftmost) / image_width
            if width_fraction >= 0.55:
                mrz_lines.append((top_bucket, line_boxes, score))

    if len(mrz_lines) < 2:
        return None

    # Collect bounding box of all MRZ lines
    all_boxes = [b for _, bs, _ in mrz_lines for b in bs]
    x = min(b["left"] for b in all_boxes)
    y = min(b["top"] for b in all_boxes)
    r = max(b["left"] + b["width"] for b in all_boxes)
    bot = max(b["top"] + b["height"] for b in all_boxes)

    # Collect OCR text spans
    spans = [(b["char_start"], b["char_end"]) for b in all_boxes if "char_start" in b]

    avg_score = sum(s for _, _, s in mrz_lines) / len(mrz_lines)
    confidence = min(1.0, avg_score * (len(mrz_lines) / 3.0))

    return LayoutSignal(
        layout_type="mrz",
        confidence=round(confidence, 3),
        bounding_box=(x, y, r - x, bot - y),
        ocr_text_spans=spans,
    )


# ---------------------------------------------------------------------------
# ID-card layout detection
# ---------------------------------------------------------------------------

_ID_KEYWORDS = re.compile(
    r"(?i)\b(?:surname|given\s+names?|nationality|date\s+of\s+birth|"
    r"place\s+of\s+birth|expiry|personal\s+no|id\s*no|document\s+no|"
    r"sex|height|issuing\s+authority|passeport|numéro|nombre|nascimento|"
    r"cognome|nome|polnisch|geburtsdatum|staatsangehörigkeit)\b"
)


def detect_id_card(
    word_boxes: List[Dict], ocr_text: str, image_width: int, image_height: int
) -> Optional[LayoutSignal]:
    """Detect ID-card document structure.

    Heuristics:
    - Right half of image contains significantly more text than left half
      (photo is on the left, text fields on the right).
    - OCR text contains ≥2 ID-specific keywords.
    """
    if not word_boxes or image_width == 0:
        return None

    mid = image_width // 2
    left_boxes = [b for b in word_boxes if b["left"] + b["width"] // 2 < mid]
    right_boxes = [b for b in word_boxes if b["left"] + b["width"] // 2 >= mid]

    # Right side should have significantly more text boxes
    if len(right_boxes) == 0 or len(right_boxes) < len(left_boxes) * 0.5:
        return None

    keyword_matches = len(_ID_KEYWORDS.findall(ocr_text))
    if keyword_matches < 2:
        return None

    # Confidence: ratio of keyword hits + side asymmetry
    asymmetry = len(right_boxes) / max(len(left_boxes), 1)
    confidence = min(1.0, 0.5 + (keyword_matches * 0.1) + (min(asymmetry, 3) * 0.1))

    # Bounding box: right half region containing text
    if right_boxes:
        x = min(b["left"] for b in right_boxes)
        y = min(b["top"] for b in right_boxes)
        r = max(b["left"] + b["width"] for b in right_boxes)
        bot = max(b["top"] + b["height"] for b in right_boxes)
        bbox = (x, y, r - x, bot - y)
    else:
        bbox = (0, 0, image_width, image_height)

    spans = [(b["char_start"], b["char_end"]) for b in right_boxes if "char_start" in b]

    return LayoutSignal(
        layout_type="id_card",
        confidence=round(confidence, 3),
        bounding_box=bbox,
        ocr_text_spans=spans,
    )


# ---------------------------------------------------------------------------
# Table / bank-statement layout detection
# ---------------------------------------------------------------------------

_FINANCIAL_KEYWORDS = re.compile(
    r"(?i)\b(?:iban|swift|bic|account|balance|amount|debit|credit|"
    r"transaction|statement|reference|invoice|total|due|payment|"
    r"solde|montant|konto|betrag|saldo|importo|saldo|saldo)\b"
)
_MIN_TABLE_ROWS = 3
_MIN_COLS_PER_ROW = 2


def detect_table(
    word_boxes: List[Dict], ocr_text: str
) -> Optional[LayoutSignal]:
    """Detect table/grid structure from word box alignment.

    Groups word boxes by row (similar top coordinates), then checks
    whether multiple rows share consistent column x-positions.
    """
    if not word_boxes:
        return None

    # Require at least some financial keywords for the table to matter privacy-wise
    kw_count = len(_FINANCIAL_KEYWORDS.findall(ocr_text))
    if kw_count < 1:
        return None

    # Group boxes into rows by top coordinate (10px tolerance)
    rows: Dict[int, List[Dict]] = {}
    for box in word_boxes:
        bucket = round(box["top"] / 12) * 12
        rows.setdefault(bucket, []).append(box)

    # Keep rows that have ≥2 word boxes (candidate table rows)
    table_rows = {k: sorted(v, key=lambda b: b["left"])
                  for k, v in rows.items() if len(v) >= _MIN_COLS_PER_ROW}

    if len(table_rows) < _MIN_TABLE_ROWS:
        return None

    # Check vertical regularity: row tops should be evenly spaced
    sorted_tops = sorted(table_rows.keys())
    if len(sorted_tops) >= 2:
        gaps = [sorted_tops[i + 1] - sorted_tops[i] for i in range(len(sorted_tops) - 1)]
        avg_gap = sum(gaps) / len(gaps)
        gap_variance = sum((g - avg_gap) ** 2 for g in gaps) / len(gaps)
        # Regular table: variance should be modest
        if avg_gap == 0 or gap_variance / (avg_gap ** 2) > 1.5:
            return None  # too irregular

    # Collect all boxes in table rows
    all_table_boxes = [b for bs in table_rows.values() for b in bs]
    x = min(b["left"] for b in all_table_boxes)
    y = min(b["top"] for b in all_table_boxes)
    r = max(b["left"] + b["width"] for b in all_table_boxes)
    bot = max(b["top"] + b["height"] for b in all_table_boxes)

    row_count = len(table_rows)
    confidence = min(1.0, 0.4 + (kw_count * 0.1) + (min(row_count, 10) * 0.04))
    spans = [(b["char_start"], b["char_end"]) for b in all_table_boxes if "char_start" in b]

    return LayoutSignal(
        layout_type="table",
        confidence=round(confidence, 3),
        bounding_box=(x, y, r - x, bot - y),
        ocr_text_spans=spans,
    )


# ---------------------------------------------------------------------------
# Combined analysis
# ---------------------------------------------------------------------------

def analyse_layout(
    word_boxes: List[Dict],
    ocr_text: str,
    image_width: int,
    image_height: int,
) -> List[LayoutSignal]:
    """Run all layout heuristics and return detected signals (may be empty)."""
    signals: List[LayoutSignal] = []

    mrz = detect_mrz(word_boxes, ocr_text, image_width)
    if mrz:
        signals.append(mrz)

    id_card = detect_id_card(word_boxes, ocr_text, image_width, image_height)
    if id_card:
        signals.append(id_card)

    table = detect_table(word_boxes, ocr_text)
    if table:
        signals.append(table)

    return signals
