"""
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 goal 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 so header/data columns align (expand colspan/rowspan)
- clean common header artifacts (e.g. "DESCRIPTIONBeginning Balance", "BALANCE$3,447.10")
"""

# 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)
# ============================================================

# 1) GLM-OCR MaaS API key
# IMPORTANT: Do NOT hard-code secrets in a public Space.
# If your Space is public, switch to HF Secrets instead.
GLMOCR_API_KEY = "e2b138b2005a41cb9d87dd18805838aa.lyd51L23rcDbsw0w"

# 2) Render quality (higher = better OCR for small/right-aligned digits; slower)
RENDER_SCALE = 2.2   # try 2.5 if Balance column is still missing

# 3) Add padding to protect columns near edges (Balance is usually 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

# 4) Mild contrast boost (helps faint gray text)
ENABLE_CONTRAST = True

# 5) Header/footer band heuristics
DEFAULT_ZONE_FRAC = 0.12      # OCR band for header (top 12%) when regions not available
PDF_HEADER_BAND_FRAC = 0.10   # PDF text fallback: take top 10% words if clip returns empty

# Footer: disabled by default to avoid duplicating what GLM already returns
ENABLE_FOOTER_OCR = False
PDF_FOOTER_BAND_FRAC = 0.88   # bottom 12% (if footer enabled)

# MaaS minimum image sizes for crops; we pad if needed
MIN_CROP_HEIGHT = 112
MIN_CROP_PIXELS = 112 * 112

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


# Single shared parser to avoid re-init per request
_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):
    """Infer header/footer extents from bbox regions if present."""
    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):
    """Extract text from a horizontal band using a clip rect (works if PDF has text layer)."""
    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):
    """Extract words whose bbox intersects a vertical band (robust fallback)."""
    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):
    """Run OCR on a horizontal band. Pads small crops to meet MaaS minimum size."""
    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:
    """Fix common account-number formatting issues."""
    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:
    """Close unclosed <table>/<tr>/<td> tags to prevent bleed."""
    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 clean_table_header_artifacts(text: str) -> str:
    """
    Generic cleanup for common OCR artifacts in table headers:
    - 'DESCRIPTIONBeginning Balance' -> 'DESCRIPTION'
    - 'DESCRIPTIONEnding Balance' -> 'DESCRIPTION'
    - 'BALANCE$3,447.10' -> 'BALANCE'
    - 'BALANCE 3,447.10' -> 'BALANCE'
    Works for any PDF; no bank-specific logic.
    """
    if not text:
        return text

    # DESCRIPTION + (Beginning/Ending Balance) glued
    text = re.sub(
        r"(>[^<]*\bDESCRIPTION)\s*(Beginning Balance|Ending Balance)\b([^<]*<)",
        r"\1\3",
        text,
        flags=re.IGNORECASE,
    )

    # BALANCE with a number glued or appended (keep the word BALANCE only)
    text = re.sub(
        r"(>[^<]*\bBALANCE)\s*\$?\s*\d{1,3}(?:,\d{3})*(?:\.\d{2})?\s*([^<]*<)",
        r"\1\2",
        text,
        flags=re.IGNORECASE,
    )
    text = re.sub(
        r"(>[^<]*\bBALANCE)\$",
        r"\1",
        text,
        flags=re.IGNORECASE,
    )

    return text


# ---- Colspan/rowspan expansion: align header and data rows for any PDF ----
class TableGridParser(HTMLParser):
    """Parse <table> HTML 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

    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 = []

    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))
        elif tag == "tr":
            self.rows.append(self._current_row)

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


def _build_grid(rows_data):
    """Expand colspan/rowspan into a rectangular grid."""
    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):
    """Emit a normalized table without colspan/rowspan."""
    if not grid:
        return ""
    lines = ["<table>"]
    for i, row in enumerate(grid):
        lines.append("<tr>")
        tag = "th" if i == 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 expand_colspan_rowspan_in_tables(text):
    """Normalize every <table>...</table> in text."""
    if not text or "<table" not in text.lower():
        return text

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

    for match in pattern.finditer(text):
        result.append(text[last_end : match.start()])
        table_html = match.group(0)
        try:
            parser = TableGridParser()
            parser.feed(table_html)
            if parser.rows:
                grid = _build_grid(parser.rows)
                result.append(_grid_to_html(grid))
            else:
                result.append(table_html)
        except Exception:
            result.append(table_html)
        last_end = match.end()

    result.append(text[last_end:])
    return "".join(result)


def looks_like_markdown_table(block: str) -> bool:
    """Detect simple markdown pipe tables."""
    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:
    """Convert a simple markdown pipe table to HTML table (best-effort)."""
    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:
    """Conservative normalization for OCR number quirks."""
    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


def stabilize_tables_and_text(page_md: str) -> str:
    """Convert markdown pipe tables to HTML, clean header artifacts, normalize tables, and close tags."""
    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)

    # 1) Fix common header text artifacts
    stabilized = clean_table_header_artifacts(stabilized)

    # 2) Normalize tables to a rectangular grid (expand colspan/rowspan)
    stabilized = expand_colspan_rowspan_in_tables(stabilized)

    # 3) Close any unclosed tags to prevent bleed
    return close_unclosed_html(stabilized)


# --------------------------
# PDF rendering with padding (critical for Balance column)
# --------------------------
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)

            # clamp
            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:
        # cleanup rendered images
        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()