app.py

"""
GLM-OCR Hugging Face Space app for PDF/image OCR with header inclusion
and table-structure stabilization for downstream bank-statement pipelines.
Hard-coded knobs (no environment variables required).
Primary goals for reconcile rate:
- preserve right-most columns (often "Balance") by higher DPI render + right padding
- keep tables as tables (convert markdown pipe tables -> HTML table)
- return ---page-separator--- between pages
- normalize HTML tables generically so downstream parsing/classification is stable:
  1) expand colspan/rowspan into a rectangular grid
  2) drop truly-empty columns (common in summary tables)
  3) merge "blank header" columns that contain text into the left column (common when DESCRIPTION is split)
  4) clean common header artifacts (e.g. "DESCRIPTIONBeginning Balance", "BALANCE$3,447.10")
  5) recover fused first data row generically (two-signal guard: text fusion + money fusion must both fire)
  6) promote misplaced header row: when real column headers land in a data row, restructure the grid
"""

# Patch asyncio first (before Gradio imports it) to suppress Python 3.13 cleanup noise
import asyncio
try:
    _orig_close = asyncio.BaseEventLoop.close
    def _safe_close(self):
        try:
            _orig_close(self)
        except (ValueError, OSError):
            pass
    asyncio.BaseEventLoop.close = _safe_close
except Exception:
    pass

import logging
import os
import re
import html
import tempfile
from typing import List, Tuple
from collections import defaultdict
from html.parser import HTMLParser

import yaml
import gradio as gr
import glmocr

log = logging.getLogger("glmocr_app")
logging.basicConfig(level=logging.INFO)

GLMOCR_BASE = os.path.dirname(glmocr.__file__)
CONFIG_PATH = os.path.join(GLMOCR_BASE, "config.yaml")
FORMATTER_PATH = os.path.join(GLMOCR_BASE, "postprocess", "result_formatter.py")

# ============================================================
# HARD-CODED SETTINGS (edit these numbers to tune quality/speed)
# ============================================================

GLMOCR_API_KEY = "e2b138b2005a41cb9d87dd18805838aa.lyd51L23rcDbsw0w"

# Higher = better OCR for small/right-aligned digits; slower
RENDER_SCALE = 2.2  # try 2.5 if right-side numbers are missed

# Padding to protect columns near edges (Balance is often right-most)
PAD_LEFT_FRAC = 0.02
PAD_RIGHT_FRAC = 0.06  # try 0.10 if right-most balances are missing
PAD_TOP_FRAC = 0.01
PAD_BOTTOM_FRAC = 0.01

ENABLE_CONTRAST = True

DEFAULT_ZONE_FRAC = 0.12
PDF_HEADER_BAND_FRAC = 0.10

ENABLE_FOOTER_OCR = False
PDF_FOOTER_BAND_FRAC = 0.88

MIN_CROP_HEIGHT = 112
MIN_CROP_PIXELS = 112 * 112

# ============================================================

_parser = None


def get_parser():
    global _parser
    if _parser is None:
        from glmocr import GlmOcr
        _parser = GlmOcr(api_key=GLMOCR_API_KEY, mode="maas")
    return _parser


# ---------------------------------------------------------------------------
# Best-effort config tweaks (safe to fail on read-only HF env)
# ---------------------------------------------------------------------------
try:
    with open(CONFIG_PATH, "r") as f:
        config = yaml.safe_load(f)
    config["pipeline"]["maas"]["enabled"] = True
    config["pipeline"]["maas"]["api_key"] = GLMOCR_API_KEY
    with open(CONFIG_PATH, "w") as f:
        yaml.dump(config, f, default_flow_style=False, sort_keys=False)
except Exception:
    pass

# Best-effort formatter tweak: avoid stripping header/footer labels
try:
    with open(FORMATTER_PATH, "r") as f:
        source = f.read()
    for label in (
        '"header"', "'header'",
        '"footer"', "'footer'",
        '"doc_header"', "'doc_header'",
        '"doc_footer"', "'doc_footer'",
    ):
        source = re.sub(r",\s*" + re.escape(label), "", source)
        source = re.sub(re.escape(label) + r"\s*,", "", source)
        source = re.sub(re.escape(label), "", source)
    with open(FORMATTER_PATH, "w") as f:
        f.write(source)
except Exception:
    pass


# --------------------------
# Header/footer helpers
# --------------------------
def get_header_footer_zones(regions, norm_height=1000):
    if not regions:
        return None, None
    y_tops, y_bottoms = [], []
    for r in regions:
        bbox = r.get("bbox_2d") if isinstance(r, dict) else getattr(r, "bbox_2d", None)
        if bbox and len(bbox) >= 4:
            y_tops.append(bbox[1])
            y_bottoms.append(bbox[3])
    if not y_tops:
        return None, None
    return min(y_tops) / norm_height, max(y_bottoms) / norm_height


def extract_zone_text_pdf(pdf_path, page_num, y_start_frac, y_end_frac):
    try:
        import pymupdf as fitz
        doc = fitz.open(pdf_path)
        page = doc[page_num]
        h, w = page.rect.height, page.rect.width
        rect = fitz.Rect(0, h * y_start_frac, w, h * y_end_frac)
        text = page.get_text(clip=rect).strip()
        doc.close()
        return text
    except Exception:
        return ""


def extract_pdf_text_in_band(pdf_path, page_num, y_start_frac, y_end_frac):
    try:
        import pymupdf as fitz
        doc = fitz.open(pdf_path)
        page = doc[page_num]
        h = page.rect.height
        y_lo = h * y_start_frac
        y_hi = h * y_end_frac
        words = page.get_text("words")
        doc.close()
        parts = []
        for w in words:
            if len(w) >= 5:
                y0, y1 = float(w[1]), float(w[3])
                if y0 < y_hi and y1 > y_lo:
                    parts.append(w[4])
        return " ".join(parts).strip()
    except Exception:
        return ""


def ocr_zone(image_path, y_start_frac, y_end_frac):
    zone_name = "header" if y_end_frac < 0.5 else "footer"
    try:
        from PIL import Image
        img = Image.open(image_path).convert("RGB")
        w, h = img.size
        y0 = max(0, int(h * y_start_frac))
        y1 = min(h, int(h * y_end_frac))
        if y1 <= y0:
            return ""

        crop = img.crop((0, y0, w, y1))
        cw, ch = crop.size

        if ch < MIN_CROP_HEIGHT or (cw * ch) < MIN_CROP_PIXELS:
            need_h = max(ch, MIN_CROP_HEIGHT)
            need_w = max(cw, 1)
            if (need_w * need_h) < MIN_CROP_PIXELS:
                need_w = max(need_w, (MIN_CROP_PIXELS + need_h - 1) // need_h)
            canvas = Image.new("RGB", (need_w, need_h), (255, 255, 255))
            if zone_name == "header":
                canvas.paste(crop, (0, 0))
            else:
                canvas.paste(crop, (0, need_h - ch))
            crop = canvas

        fd, path = tempfile.mkstemp(suffix=".jpg")
        os.close(fd)
        try:
            crop.save(path, "JPEG", quality=92)
            parser = get_parser()
            out = parser.parse(path)
            if not isinstance(out, list):
                out = [out]
            if out and getattr(out[0], "markdown_result", None):
                return (out[0].markdown_result or "").strip()
        finally:
            try:
                os.unlink(path)
            except Exception:
                pass
    except Exception as e:
        log.warning("[%s] ocr_zone failed: %s", zone_name, e, exc_info=True)
    return ""


def fix_account_number(hdr: str) -> str:
    if not hdr:
        return hdr
    if "Account Number:" in hdr and "Account Number: " not in hdr:
        m = re.search(r"[0-9]{5,}", hdr)
        if m:
            hdr = hdr.replace("Account Number:", "Account Number: " + m.group(0))
    acct_match = re.search(r"Account Number: ([0-9]{5,})", hdr)
    if acct_match:
        acct = acct_match.group(1)
        if hdr.startswith(acct):
            hdr = hdr[len(acct):].lstrip()
    return hdr


# --------------------------
# Table stabilization helpers
# --------------------------
def close_unclosed_html(md: str) -> str:
    if not md:
        return md
    open_tags = re.findall(r"<(table|tbody|thead|tr|td|th)\b", md, flags=re.IGNORECASE)
    close_tags = re.findall(r"</(table|tbody|thead|tr|td|th)>", md, flags=re.IGNORECASE)

    def count(tags, name):
        return sum(1 for t in tags if t.lower() == name)

    for tag in reversed(["td", "th", "tr", "thead", "tbody", "table"]):
        opened = count(open_tags, tag)
        closed = count(close_tags, tag)
        if opened > closed:
            md += ("</%s>" % tag) * (opened - closed)
    return md


def looks_like_markdown_table(block: str) -> bool:
    lines = [ln.rstrip() for ln in block.strip().splitlines() if ln.strip()]
    if len(lines) < 2:
        return False
    if "|" not in lines[0]:
        return False
    sep = lines[1].replace(" ", "")
    return ("---" in sep) and ("|" in sep)


def md_table_to_html(block: str) -> str:
    lines = [ln.strip() for ln in block.strip().splitlines() if ln.strip()]
    if len(lines) < 2:
        return block

    def split_row(row: str):
        row = row.strip()
        if row.startswith("|"):
            row = row[1:]
        if row.endswith("|"):
            row = row[:-1]
        return [p.strip() for p in row.split("|")]

    header = split_row(lines[0])
    body_lines = [ln for ln in lines[2:] if "|" in ln]

    html_rows = []
    html_rows.append("<tr>" + "".join(f"<th>{html.escape(c)}</th>" for c in header) + "</tr>")
    for ln in body_lines:
        cols = split_row(ln)
        if len(cols) < len(header):
            cols += [""] * (len(header) - len(cols))
        html_rows.append("<tr>" + "".join(f"<td>{html.escape(c)}</td>" for c in cols[: len(header)]) + "</tr>")
    return "<table>\n" + "\n".join(html_rows) + "\n</table>"


def normalize_money_glyphs(text: str) -> str:
    if not text:
        return text
    t = text.replace("−", "-").replace("–", "-").replace("—", "-")
    t = re.sub(
        r"\(\s*\$?\s*([0-9]{1,3}(?:,[0-9]{3})*|[0-9]+)(\.[0-9]{2})\s*\)",
        r"-\1\2",
        t,
    )

    def o_to_zero(m):
        token = m.group(0)
        return token.replace("O", "0").replace("o", "0")

    t = re.sub(r"\b[0-9Oo\$,.\-]{4,}\b", o_to_zero, t)
    return t


# ---- Generic HTML table normalizer ----
class TableGridParser(HTMLParser):
    """Parse a <table> into rows of (text, colspan, rowspan)."""

    def __init__(self):
        super().__init__()
        self.rows = []
        self._current_row = []
        self._cell_text = []
        self._colspan = 1
        self._rowspan = 1
        self._in_cell = False

    def handle_starttag(self, tag, attrs):
        if tag == "tr":
            self._current_row = []
        elif tag in ("td", "th"):
            attrs_d = dict(attrs)
            self._colspan = max(1, int(attrs_d.get("colspan", 1)))
            self._rowspan = max(1, int(attrs_d.get("rowspan", 1)))
            self._cell_text = []
            self._in_cell = True

    def handle_endtag(self, tag):
        if tag in ("td", "th"):
            text = "".join(self._cell_text).strip().replace("\n", " ")
            self._current_row.append((text, self._colspan, self._rowspan))
            self._in_cell = False
        elif tag == "tr":
            self.rows.append(self._current_row)

    def handle_data(self, data):
        if self._in_cell:
            self._cell_text.append(data)


def _build_grid(rows_data):
    if not rows_data:
        return []
    blocked = defaultdict(set)
    grid = []

    for r, row_cells in enumerate(rows_data):
        grid.append([])
        col = 0
        for content, C, R in row_cells:
            while col in blocked[r]:
                grid[r].append("")
                col += 1
            for k in range(C):
                grid[r].append(content if k == 0 else "")
            for k in range(1, R):
                blocked[r + k].add(col)
            col += C

    max_cols = max(len(row) for row in grid) if grid else 0
    for row in grid:
        while len(row) < max_cols:
            row.append("")
    return grid


def _grid_to_html(grid):
    if not grid:
        return ""
    lines = ["<table>"]
    for r, row in enumerate(grid):
        lines.append("<tr>")
        tag = "th" if r == 0 else "td"
        for cell in row:
            escaped = (
                (cell or "")
                .replace("&", "&amp;")
                .replace("<", "&lt;")
                .replace(">", "&gt;")
                .replace('"', "&quot;")
            )
            lines.append(f"<{tag}>{escaped}</{tag}>")
        lines.append("</tr>")
    lines.append("</table>")
    return "\n".join(lines)


def _is_amount_like(s: str) -> bool:
    if not s:
        return False
    return re.fullmatch(r"\$?-?\d{1,3}(?:,\d{3})*(?:\.\d{2})?", s.strip()) is not None


def _drop_truly_empty_columns(grid):
    """
    Drop columns that are empty in header AND almost always empty in body.
    This fixes summary tables that have an extra blank trailing column.
    """
    if not grid or len(grid) < 1:
        return grid

    header = [str(c or "").strip() for c in grid[0]]
    ncols = len(header)
    if ncols <= 1:
        return grid

    body = grid[1:]
    keep = [True] * ncols

    for i in range(ncols):
        if header[i] != "":
            continue
        total = 0
        non_empty = 0
        for row in body:
            if i >= len(row):
                continue
            total += 1
            if str(row[i] or "").strip():
                non_empty += 1
        if total > 0 and (non_empty / total) <= 0.05:
            keep[i] = False

    if all(keep):
        return grid

    new_grid = []
    for row in grid:
        new_grid.append([cell for idx, cell in enumerate(row) if idx < len(keep) and keep[idx]])
    return new_grid


def _merge_blank_header_text_columns(grid):
    """
    If the header row has blank columns, and most body rows have non-empty *text* in that blank column,
    merge that column into the nearest non-empty header to the left (usually DESCRIPTION),
    then remove the blank column.
    This fixes transaction tables where DESCRIPTION is split across two columns.
    """
    if not grid or len(grid) < 2:
        return grid

    header = [str(c or "").strip() for c in grid[0]]
    ncols = len(header)
    body = grid[1:]
    blank_cols = [i for i, h in enumerate(header) if h == ""]
    if not blank_cols:
        return grid

    keep = [True] * ncols

    for i in blank_cols:
        j = i - 1
        while j >= 0 and header[j] == "":
            j -= 1
        if j < 0:
            continue

        total = 0
        non_empty = 0
        texty = 0
        for row in body:
            if i >= len(row) or j >= len(row):
                continue
            v = str(row[i] or "").strip()
            total += 1
            if v:
                non_empty += 1
                if not _is_amount_like(v):
                    texty += 1

        if total == 0:
            continue
        non_empty_ratio = non_empty / total
        texty_ratio = (texty / non_empty) if non_empty else 0.0

        if non_empty_ratio >= 0.55 and texty_ratio >= 0.70:
            for r in range(1, len(grid)):
                row = grid[r]
                if i >= len(row) or j >= len(row):
                    continue
                left = str(row[j] or "").strip()
                right = str(row[i] or "").strip()
                if right:
                    row[j] = (left + " " + right).strip() if left else right
            keep[i] = False

    if all(keep):
        return grid

    new_grid = []
    for row in grid:
        new_grid.append([cell for idx, cell in enumerate(row) if idx < len(keep) and keep[idx]])
    return new_grid


# --------------------------
# Shared keyword definitions
# --------------------------

# Matches a standalone money value across currencies.
# Covers: $3,447.10  -1,234.56  £500.00  1000  etc.
_MONEY_RE = re.compile(
    r"^[\$£€¥]?\s*-?\s?\d{1,3}(?:,\d{3})*(?:\.\d{2})?$"
)

# All column-header keywords recognised across any bank statement format.
_HEADER_KW_PATTERNS = [
    # Date variants
    r"DATE",
    r"POSTING\s+DATE",
    r"VALUE\s+DATE",
    r"TXN\s+DATE",
    r"TRANSACTION\s+DATE",
    r"ENTRY\s+DATE",
    r"EFFECTIVE\s+DATE",

    # Transaction ID / reference
    r"TRANSACTION(?:\s+(?:ID|TYPE|NO|NUMBER))?",
    r"TXN(?:\s+(?:ID|NO|TYPE))?",
    r"REF(?:ERENCE)?(?:\s*(?:NO|NUM|NUMBER))?",
    r"CHEQUE(?:\s*(?:NO|NUMBER))?",
    r"CHQ(?:\s*(?:NO|NUMBER))?",
    r"VOUCHER(?:\s*(?:NO|NUMBER))?",
    r"SR\.?\s*NO\.?",
    r"SERIAL(?:\s*(?:NO|NUMBER))?",

    # Description variants
    r"DESCRIPTION",
    r"DETAILS?",
    r"PARTICULARS?(?:\s+OF\s+TRANSACTION)?",
    r"NARRATION",
    r"REMARKS?",
    r"NOTES?",

    # Debit variants
    r"DEBIT",
    r"DEBITS?",
    r"DR\.?",
    r"WITHDRAWALS?",
    r"PAID\s+OUT",
    r"MONEY\s+OUT",

    # Credit variants
    r"CREDIT",
    r"CREDITS?",
    r"CR\.?",
    r"DEPOSITS?",
    r"PAID\s+IN",
    r"MONEY\s+IN",

    # Amount
    r"AMOUNT",

    # Balance variants
    r"BALANCE",
    r"BAL\.?",
    r"RUNNING\s+BALANCE",
    r"RUNNING\s+BAL\.?",
    r"AVAILABLE\s+BALANCE",
    r"AVAILABLE\s+BAL\.?",
    r"AVAIL\.?\s+BAL\.?",
    r"LEDGER\s+BALANCE",
    r"LEDGER\s+BAL\.?",
    r"CLOSING\s+BALANCE",
    r"CLOSING\s+BAL\.?",
    r"OPENING\s+BALANCE",
    r"OPENING\s+BAL\.?",
]

# Pre-compiled: each pattern anchored at start, case-insensitive
_HEADER_KW_RES = [
    re.compile(r"^(" + p + r")(.*)", re.IGNORECASE | re.DOTALL)
    for p in _HEADER_KW_PATTERNS
]

# Quick exact-match check: "is this entire string a known keyword?"
_HEADER_KW_EXACT_RE = re.compile(
    r"^(?:" + r"|".join(_HEADER_KW_PATTERNS) + r")$",
    re.IGNORECASE,
)

# Minimum fraction of cells in a row that must be pure keywords for that row
# to be considered a misplaced header row.
_HEADER_ROW_KEYWORD_THRESHOLD = 0.5


# --------------------------
# Fix 1: Fused-header recovery
# --------------------------

def _split_keyword_remainder(cell_text: str):
    """
    If cell_text starts with a known header keyword followed by extra content,
    return (keyword, remainder).  Otherwise return (cell_text, "").
    """
    s = cell_text.strip()
    for pattern in _HEADER_KW_RES:
        m = pattern.match(s)
        if m:
            keyword = m.group(1).strip()
            remainder = m.group(2).strip()
            return keyword, remainder
    return s, ""


def _extract_fused_header_artifacts(header_row):
    """
    Generic detector for the OCR artifact where the first data row of a table
    gets fused into the header cells during OCR.

    The artifact pattern (BOTH signals must fire simultaneously):
      - Signal 1 — text-fused: a header cell contains KEYWORD + free descriptive text
        e.g.  "DESCRIPTIONBeginning Balance"  "NARRATIONOpening Balance"
      - Signal 2 — money-fused: a DIFFERENT header cell contains KEYWORD + money amount
        e.g.  "BALANCE$3,447.10"  "AMOUNT 5,000.00"

    Two-signal guard prevents false positives on clean PDFs from any bank.

    Returns:
        (cleaned_header : list[str], recovered_row : list[str] | None)
    """
    if not header_row:
        return list(header_row), None

    ncols = len(header_row)
    cleaned = list(header_row)
    recovered = [""] * ncols

    text_fused = []
    money_fused = []

    for idx, cell in enumerate(header_row):
        cell_s = str(cell or "").strip()
        if not cell_s:
            continue

        keyword, remainder = _split_keyword_remainder(cell_s)

        if not remainder:
            continue

        remainder_no_space = remainder.replace(" ", "")

        if _MONEY_RE.match(remainder_no_space):
            cleaned[idx] = keyword
            recovered[idx] = remainder
            money_fused.append(idx)

        elif not _HEADER_KW_EXACT_RE.match(remainder.split()[0] if remainder.split() else ""):
            cleaned[idx] = keyword
            recovered[idx] = remainder
            text_fused.append(idx)

    if text_fused and money_fused:
        return cleaned, recovered

    return list(header_row), None


def _clean_header_artifacts(grid):
    """
    Entry point for Fix 1 called from normalize_html_tables.
    Returns (grid, recovered_row | None).
    """
    if not grid or not grid[0]:
        return grid, None

    cleaned_header, recovered_row = _extract_fused_header_artifacts(grid[0])
    grid[0] = cleaned_header
    return grid, recovered_row


# --------------------------
# Fix 2: Misplaced header row promotion
# --------------------------

def _row_keyword_score(row):
    """
    Return the fraction of non-empty cells in this row that are pure header
    keywords (e.g. DATE, DESCRIPTION, DEBIT, CREDIT, BALANCE).
    Also handles cells like "DATE DESCRIPTION" where two keywords are space-joined
    into one cell — these count as a keyword cell too.
    """
    non_empty = [str(c or "").strip() for c in row if str(c or "").strip()]
    if not non_empty:
        return 0.0

    keyword_hits = 0
    for cell in non_empty:
        # Check if the whole cell is a known keyword
        if _HEADER_KW_EXACT_RE.match(cell):
            keyword_hits += 1
            continue
        # Check if the cell is two or more space-separated keywords
        # e.g. "DATE DESCRIPTION"
        parts = cell.split()
        if all(_HEADER_KW_EXACT_RE.match(p) for p in parts) and len(parts) > 1:
            keyword_hits += 1
            continue

    return keyword_hits / len(non_empty)


def _split_merged_keyword_cell(cell_text: str, target_ncols: int, col_idx: int, total_cols: int):
    """
    When a data cell contains multiple space-separated keywords (e.g. "DATE DESCRIPTION"),
    split it into individual keyword tokens.  Returns a list of strings to replace
    the single cell — padded/trimmed to fit the available column space.
    """
    parts = cell_text.strip().split()
    # Only split if every part is a keyword
    if len(parts) > 1 and all(_HEADER_KW_EXACT_RE.match(p) for p in parts):
        return parts
    return [cell_text]


def _promote_misplaced_header_row(grid):
    """
    Detects and fixes the OCR artifact where the real column headers (DATE,
    DESCRIPTION, DEBIT, CREDIT, BALANCE …) land in a <td> data row instead
    of the <th> header row.

    This happens when a bank PDF wraps the transaction section in a larger
    table that also contains section titles (e.g. "ACCOUNT ACTIVITY" and
    "Transactions by Date (continued)"), so the OCR emits:

        <th>ACCOUNT ACTIVITY</th>          ← wrong header (section title)
        <td>Transactions by Date …</td>    ← subtitle row
        <td>DATE DESCRIPTION</td> <td>DEBIT</td> <td>CREDIT</td> <td>BALANCE</td>
        <td>01/05 …</td> …                 ← real data

    The fix:
      1. Scan every data row (rows 1+) for a row whose cells are predominantly
         pure header keywords.
      2. When found, promote that row to be the new header (row 0) and discard
         all rows above it (they are section titles / subtitles).
      3. If any cell in the promoted row contains multiple space-joined keywords
         (e.g. "DATE DESCRIPTION"), expand it into separate columns and
         re-align the data rows below accordingly.

    Guard conditions (no-op if not met — safe for all other PDFs):
      - The current header row must NOT already look like real column headers
        (keyword score < threshold).
      - At least one data row must have keyword score >= threshold.
      - The candidate row must appear within the first 5 data rows (it is
        always near the top of the table, never buried mid-table).
    """
    if not grid or len(grid) < 2:
        return grid

    current_header = grid[0]
    current_header_score = _row_keyword_score(current_header)

    # If the current header already looks like real column headers, do nothing.
    if current_header_score >= _HEADER_ROW_KEYWORD_THRESHOLD:
        return grid

    # Search the first few data rows for a candidate header row.
    candidate_idx = None
    candidate_score = 0.0
    search_limit = min(len(grid), 6)  # rows 1..5

    for i in range(1, search_limit):
        score = _row_keyword_score(grid[i])
        if score >= _HEADER_ROW_KEYWORD_THRESHOLD and score > candidate_score:
            candidate_score = score
            candidate_idx = i

    if candidate_idx is None:
        return grid

    # --- Promote the candidate row ---
    candidate_row = grid[candidate_idx]

    # Expand any merged keyword cells (e.g. "DATE DESCRIPTION" → ["DATE", "DESCRIPTION"])
    expanded_header = []
    for cell in candidate_row:
        cell_s = str(cell or "").strip()
        parts = cell_s.split()
        if len(parts) > 1 and all(_HEADER_KW_EXACT_RE.match(p) for p in parts):
            expanded_header.extend(parts)
        else:
            expanded_header.append(cell_s)

    new_ncols = len(expanded_header)
    old_ncols = len(candidate_row)
    col_expansion = new_ncols - old_ncols  # how many extra columns were added

    # Re-align data rows below the candidate (expand first cell if needed)
    new_data_rows = []
    for row in grid[candidate_idx + 1:]:
        if col_expansion > 0 and row:
            # The first cell of a data row often contains DATE + start of DESCRIPTION
            # fused together (e.g. "01/05 TD ZELLE RECEIVED,...").
            # Split it into [date_part, description_part] to fill the expanded columns.
            first_cell = str(row[0] or "").strip()
            date_match = re.match(r"^(\d{1,2}/\d{1,2})\s+(.*)", first_cell, re.DOTALL)
            if date_match and col_expansion == 1:
                new_row = [date_match.group(1), date_match.group(2).strip()] + list(row[1:])
            else:
                # Can't split cleanly — pad with empty cells to maintain column count
                new_row = list(row) + [""] * col_expansion
            new_data_rows.append(new_row)
        else:
            new_data_rows.append(list(row))

    # Normalise all rows to new_ncols width
    def pad(row, n):
        r = list(row)
        while len(r) < n:
            r.append("")
        return r[:n]

    new_grid = [pad(expanded_header, new_ncols)]
    for row in new_data_rows:
        new_grid.append(pad(row, new_ncols))

    return new_grid


# --------------------------
# Normalizer entry point
# --------------------------

def normalize_html_tables(text: str) -> str:
    """
    For every <table>...</table>:
      1. Parse to grid (expand colspan/rowspan).
      2. Drop truly-empty columns.
      3. Merge blank-header text columns into the left column.
      4. Fix 1 — clean fused header artifacts and recover the lost first data row.
      5. Fix 2 — promote a misplaced header row when real column headers landed
                  in a <td> data row instead of the <th> header row.
      6. Emit normalised <table> HTML.
    """
    if not text or "<table" not in text.lower():
        return text

    pattern = re.compile(r"<table[^>]*>.*?</table>", re.DOTALL | re.IGNORECASE)
    out = []
    last = 0

    for m in pattern.finditer(text):
        out.append(text[last : m.start()])
        table_html = m.group(0)
        try:
            p = TableGridParser()
            p.feed(table_html)
            grid = _build_grid(p.rows)
            grid = _drop_truly_empty_columns(grid)
            grid = _merge_blank_header_text_columns(grid)

            # Fix 1: fused-header artifact (two-signal guard — safe on clean PDFs)
            grid, recovered_row = _clean_header_artifacts(grid)
            if recovered_row is not None:
                grid.insert(1, recovered_row)

            # Fix 2: misplaced header row promotion (keyword-score guard — safe on clean PDFs)
            grid = _promote_misplaced_header_row(grid)

            out.append(_grid_to_html(grid) if grid else table_html)
        except Exception:
            out.append(table_html)
        last = m.end()

    out.append(text[last:])
    return "".join(out)


def stabilize_tables_and_text(page_md: str) -> str:
    if not page_md:
        return page_md

    page_md = normalize_money_glyphs(page_md)

    blocks = re.split(r"\n\s*\n", page_md.strip())
    out_blocks = []
    for b in blocks:
        if looks_like_markdown_table(b):
            out_blocks.append(md_table_to_html(b))
        else:
            out_blocks.append(b)

    stabilized = "\n\n".join(out_blocks)
    stabilized = normalize_html_tables(stabilized)
    return close_unclosed_html(stabilized)


# --------------------------
# PDF rendering with padding
# --------------------------
def render_pdf_pages_to_images(pdf_path: str) -> Tuple[List[str], List[int]]:
    import pymupdf as fitz
    from PIL import Image, ImageEnhance

    doc = fitz.open(pdf_path)
    page_images: List[str] = []
    page_heights: List[int] = []

    for i in range(len(doc)):
        page = doc[i]
        pix = page.get_pixmap(matrix=fitz.Matrix(RENDER_SCALE, RENDER_SCALE), alpha=False)

        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)

        if ENABLE_CONTRAST:
            img = ImageEnhance.Contrast(img).enhance(1.12)

        w, h = img.size
        pad_l = int(w * PAD_LEFT_FRAC)
        pad_r = int(w * PAD_RIGHT_FRAC)
        pad_t = int(h * PAD_TOP_FRAC)
        pad_b = int(h * PAD_BOTTOM_FRAC)

        if any(p > 0 for p in (pad_l, pad_r, pad_t, pad_b)):
            canvas = Image.new("RGB", (w + pad_l + pad_r, h + pad_t + pad_b), (255, 255, 255))
            canvas.paste(img, (pad_l, pad_t))
            img = canvas

        img_path = os.path.join(tempfile.gettempdir(), f"glmocr_page_{os.getpid()}_{i}.png")
        img.save(img_path, "PNG", compress_level=6)
        page_images.append(img_path)
        page_heights.append(img.height)

    doc.close()
    return page_images, page_heights


# --------------------------
# GLM-OCR result extraction
# --------------------------
def get_page_md_and_regions(page_result):
    md = ""
    if hasattr(page_result, "markdown_result") and page_result.markdown_result:
        md = (page_result.markdown_result or "").strip()
    regions = []
    if hasattr(page_result, "json_result"):
        jr = page_result.json_result
        if isinstance(jr, dict) and "regions" in jr:
            regions = jr.get("regions") or []
        elif isinstance(jr, list) and len(jr) > 0:
            r = jr[0] if isinstance(jr[0], list) else jr
            if isinstance(r, list):
                regions = r
            elif isinstance(r, dict) and "regions" in r:
                regions = r.get("regions") or []
    return md, regions


# --------------------------
# Main entry
# --------------------------
def run_ocr(uploaded_file):
    if uploaded_file is None:
        return "Please upload a file."

    page_images = []
    try:
        path = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)
        is_pdf = path.lower().endswith(".pdf")
        parser = get_parser()

        page_heights = []

        if is_pdf:
            page_images, page_heights = render_pdf_pages_to_images(path)
            results = parser.parse(page_images)
        else:
            page_images = [path]
            page_heights = [1000]
            results = parser.parse(path)

        if not isinstance(results, list):
            results = [results]

        all_pages = []
        for page_num, page_result in enumerate(results):
            page_md, regions = get_page_md_and_regions(page_result)
            img_h = page_heights[page_num] if page_num < len(page_heights) else 1000
            header_end_frac, footer_start_frac = get_header_footer_zones(regions, img_h)

            he = header_end_frac if header_end_frac is not None else DEFAULT_ZONE_FRAC
            fs = footer_start_frac if footer_start_frac is not None else (1.0 - DEFAULT_ZONE_FRAC)

            he = max(0.02, min(0.25, he))
            fs = max(0.75, min(0.98, fs))

            parts = []

            # Header inclusion: PDF text -> band words -> OCR band
            hdr = ""
            if is_pdf:
                hdr = extract_zone_text_pdf(path, page_num, 0, he)
                if not (hdr and hdr.strip()):
                    hdr = extract_pdf_text_in_band(path, page_num, 0, PDF_HEADER_BAND_FRAC)
            if not (hdr and hdr.strip()) and page_num < len(page_images):
                hdr = ocr_zone(page_images[page_num], 0, he)
            if hdr and hdr.strip():
                parts.append(fix_account_number(normalize_money_glyphs(hdr.strip())))

            # Main OCR markdown, stabilized
            if page_md and page_md.strip():
                parts.append(stabilize_tables_and_text(page_md.strip()))

            # Optional footer
            if ENABLE_FOOTER_OCR and page_num < len(page_images):
                ftr = ""
                if is_pdf:
                    ftr = extract_zone_text_pdf(path, page_num, fs, 1.0)
                    if not (ftr and ftr.strip()):
                        ftr = extract_pdf_text_in_band(path, page_num, PDF_FOOTER_BAND_FRAC, 1.0)
                if not (ftr and ftr.strip()):
                    ftr = ocr_zone(page_images[page_num], fs, 1.0)
                if ftr and ftr.strip():
                    parts.append(normalize_money_glyphs(ftr.strip()))

            if parts:
                all_pages.append("\n\n".join(parts))

        return "\n\n---page-separator---\n\n".join(all_pages) if all_pages else "(No content)"

    except Exception as e:
        import traceback
        log.exception("run_ocr failed: %s", e)
        return f"Error: {e}\n\n{traceback.format_exc()}"

    finally:
        for p in page_images:
            try:
                if isinstance(p, str) and p.endswith(".png") and "glmocr_page_" in os.path.basename(p):
                    os.unlink(p)
            except Exception:
                pass


with gr.Blocks(title="GLM-OCR") as demo:
    gr.Markdown("# GLM-OCR\nUpload a PDF or image. Headers included; tables stabilized.")
    file_in = gr.File(label="Upload PDF or image", file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"])
    run_btn = gr.Button("Run OCR", variant="primary")
    out = gr.Textbox(lines=40, label="Output (markdown)")
    run_btn.click(fn=run_ocr, inputs=file_in, outputs=out)

if __name__ == "__main__":
    demo.launch()