src/services/markdown_formatter.py

"""
markdown_formatter.py
─────────────────────
Converts raw PaddleOCR bounding-box results into clean, structured Markdown.

PaddleOCR returns a list of items, each shaped as:
    [[x1,y1], [x2,y2], [x3,y3], [x4,y4]], (text, confidence)

We use the bounding boxes to reconstruct reading order, detect key-value
pairs, infer headings, and build Markdown tables — without any extra models.
"""

from __future__ import annotations

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

# ─────────────────────────────────────────────
# Tunables
# ─────────────────────────────────────────────
CONFIDENCE_THRESHOLD = 0.55   # Drop lines below this score (garbage characters)
HEADING_MAX_WORDS    = 6      # Lines with ≤ this many words may be headings
HEADING_MIN_CONF     = 0.80   # Headings must be recognised with high confidence
PARA_GAP_MULTIPLIER  = 1.6    # Y-gap > avg_line_height × this → new paragraph
TABLE_X_TOLERANCE    = 40     # Pixels: columns within this X range are "aligned"
TABLE_MIN_ROWS       = 2      # Need at least this many rows to form a table


@dataclass
class OcrLine:
    """One detected text line with all spatial & recognition metadata."""
    text: str
    confidence: float
    x_min: float
    x_max: float
    y_min: float
    y_max: float

    @property
    def height(self) -> float:
        return self.y_max - self.y_min

    @property
    def word_count(self) -> int:
        return len(self.text.split())

    @property
    def x_center(self) -> float:
        return (self.x_min + self.x_max) / 2


# ─────────────────────────────────────────────
# Public entry point
# ─────────────────────────────────────────────

def format_ocr_result(raw_result) -> Tuple[str, str]:
    """
    Convert raw PaddleOCR output into both a plain-text string and a
    structured Markdown string.

    Args:
        raw_result: The list returned by ocr.ocr(img, cls=True)[0]

    Returns:
        (plain_text, markdown_text) — both are complete strings
    """
    lines = _parse_raw(raw_result)
    lines = _filter_low_confidence(lines)

    if not lines:
        return "", ""

    plain_text = "\n".join(l.text for l in lines)
    markdown   = _build_markdown(lines)

    return plain_text, markdown


# ─────────────────────────────────────────────
# Parsing
# ─────────────────────────────────────────────

def _parse_raw(raw_result) -> List[OcrLine]:
    """Turn PaddleOCR raw list into OcrLine objects, sorted top-to-bottom."""
    lines: List[OcrLine] = []
    if not raw_result:
        return lines

    for item in raw_result:
        try:
            box, (text, conf) = item
            # box is 4 corner points — extract bounding rectangle
            xs = [p[0] for p in box]
            ys = [p[1] for p in box]
            lines.append(OcrLine(
                text=text.strip(),
                confidence=conf,
                x_min=min(xs),
                x_max=max(xs),
                y_min=min(ys),
                y_max=max(ys),
            ))
        except Exception:
            continue

    # Sort by vertical position so we read top-to-bottom
    lines.sort(key=lambda l: l.y_min)
    return lines


def _filter_low_confidence(lines: List[OcrLine]) -> List[OcrLine]:
    """Remove lines whose OCR confidence is below the threshold."""
    return [l for l in lines if l.confidence >= CONFIDENCE_THRESHOLD]


# ─────────────────────────────────────────────
# Markdown builder
# ─────────────────────────────────────────────

def _build_markdown(lines: List[OcrLine]) -> str:
    """
    Walk through sorted lines and emit Markdown tokens.

    Strategy:
      1. Group lines into paragraphs using Y-gap analysis.
      2. Within each group, attempt to detect a Markdown table.
      3. Otherwise, classify each line as: heading, key-value pair, or body text.
    """
    if not lines:
        return ""

    avg_height = sum(l.height for l in lines) / len(lines)
    para_gap   = avg_height * PARA_GAP_MULTIPLIER

    # Split into paragraph groups
    groups: List[List[OcrLine]] = []
    current_group: List[OcrLine] = [lines[0]]

    for prev, cur in zip(lines, lines[1:]):
        gap = cur.y_min - prev.y_max
        if gap > para_gap:
            groups.append(current_group)
            current_group = [cur]
        else:
            current_group.append(cur)
    groups.append(current_group)

    md_parts: List[str] = []

    for group in groups:
        table_md = _try_build_table(group)
        if table_md:
            md_parts.append(table_md)
        else:
            for line in group:
                md_parts.append(_classify_line(line))

    return "\n\n".join(p for p in md_parts if p.strip())


def _classify_line(line: OcrLine) -> str:
    """
    Classify a single line and return its Markdown representation.

    Priority:
      heading > key-value > body text
    """
    # Heading detection: short, high-confidence, uppercase-ish or title-case
    if (
        line.word_count <= HEADING_MAX_WORDS
        and line.confidence >= HEADING_MIN_CONF
        and _looks_like_heading(line.text)
    ):
        return f"## {line.text}"

    # Key-value detection: "Label: Value" or "Label - Value"
    kv = _extract_key_value(line.text)
    if kv:
        key, val = kv
        return f"**{key}:** {val}"

    # Default: plain body text
    return line.text


def _looks_like_heading(text: str) -> bool:
    """
    Heuristic: a line looks like a heading if it is mostly uppercase,
    or in title case, and contains no sentence-ending punctuation mid-string.
    """
    # Strip trailing punctuation for comparison
    core = text.rstrip(".:,;")
    if not core:
        return False

    words = core.split()
    # Mostly uppercase words (like "GOVERNMENT OF INDIA")
    upper_words = sum(1 for w in words if w.isupper() and len(w) > 1)
    if upper_words / len(words) >= 0.6:
        return True

    # Title case (like "Birth Certificate")
    title_words = sum(1 for w in words if w.istitle())
    if title_words / len(words) >= 0.7:
        return True

    return False


def _extract_key_value(text: str) -> Optional[Tuple[str, str]]:
    """
    Detect common key-value patterns in OCR'd text.
    Examples:
        "DOB: 01/06/2007"        → ("DOB", "01/06/2007")
        "Name of Father - Kumar" → ("Name of Father", "Kumar")
        "Aadhaar no.issued:..."  → ("Aadhaar no.issued", "...")
    """
    # Pattern: anything before ":" or " - " then a value
    m = re.match(r'^([A-Za-z][^:\-]{1,40}?)\s*[:\-]\s*(.+)$', text)
    if m:
        key, val = m.group(1).strip(), m.group(2).strip()
        # Reject if key is suspiciously long or value is empty
        if val and len(key.split()) <= 6:
            return key, val
    return None


# ─────────────────────────────────────────────
# Table detection
# ─────────────────────────────────────────────

def _try_build_table(group: List[OcrLine]) -> Optional[str]:
    """
    If lines in a group appear to form columns (aligned X-centers),
    reconstruct them as a Markdown table.

    Returns the table string or None if no table structure is found.
    """
    if len(group) < TABLE_MIN_ROWS:
        return None

    # Collect all unique X-centers and cluster them into columns
    x_centers = [l.x_center for l in group]
    columns = _cluster_values(x_centers, TABLE_X_TOLERANCE)

    # Need at least 2 distinct columns to form a table
    if len(columns) < 2:
        return None

    # Assign each line to a column by nearest center
    # Build a row-major dict: row_y_key → {col_idx: text}
    row_map: dict = {}
    for line in group:
        col_idx = _nearest_cluster(line.x_center, columns)
        row_key  = round(line.y_min / 15)  # bucket by ~15px rows
        if row_key not in row_map:
            row_map[row_key] = {}
        row_map[row_key][col_idx] = line.text

    rows = [row_map[k] for k in sorted(row_map.keys())]
    n_cols = len(columns)

    # Build markdown table
    header = rows[0] if rows else {}
    header_cells = [header.get(i, "") for i in range(n_cols)]
    separator    = ["---"] * n_cols

    md_rows = [
        "| " + " | ".join(header_cells) + " |",
        "| " + " | ".join(separator)    + " |",
    ]
    for row in rows[1:]:
        cells = [row.get(i, "") for i in range(n_cols)]
        md_rows.append("| " + " | ".join(cells) + " |")

    return "\n".join(md_rows)


def _cluster_values(values: List[float], tolerance: float) -> List[float]:
    """Group a list of floats into clusters separated by at least `tolerance`."""
    if not values:
        return []
    sorted_vals = sorted(set(values))
    clusters: List[float] = [sorted_vals[0]]
    for v in sorted_vals[1:]:
        if v - clusters[-1] > tolerance:
            clusters.append(v)
        else:
            # Merge into cluster centroid
            clusters[-1] = (clusters[-1] + v) / 2
    return clusters


def _nearest_cluster(value: float, clusters: List[float]) -> int:
    """Return the index of the nearest cluster to `value`."""
    return min(range(len(clusters)), key=lambda i: abs(clusters[i] - value))