#!/usr/bin/env python3

import logging
import os
import re
import subprocess
import tempfile
import time
from typing import Any, List

import gradio as gr
import httpx
import yaml

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

MODEL_ID = "zai-org/GLM-OCR"
OCR_API_PORT = 8080
OCR_API_BASE = f"http://127.0.0.1:{OCR_API_PORT}"
VLLM_MAX_MODEL_LEN = "8192"
VLLM_GPU_MEMORY_UTIL = "0.75"
VLLM_EXTRA_ARGS = ""
GLMOCR_MAX_OUTPUT_TOKENS = 2048

_parser = None
_vllm_proc = None


def _render_pdf_pages_to_images(pdf_path: str) -> List[str]:
    import pymupdf as fitz

    doc = fitz.open(pdf_path)
    page_images: List[str] = []
    try:
        for i in range(len(doc)):
            page = doc[i]
            pix = page.get_pixmap(matrix=fitz.Matrix(2.5, 2.5), alpha=False)
            img_path = os.path.join(
                tempfile.gettempdir(),
                f"glmocr_page_{os.getpid()}_{i}_{int(time.time()*1000)}.png",
            )
            pix.save(img_path)
            page_images.append(img_path)
    finally:
        doc.close()
    return page_images


def _start_vllm_if_needed() -> None:
    global _vllm_proc
    if _vllm_proc is not None and _vllm_proc.poll() is None:
        return

    cmd: List[str] = [
        "python", "-m", "vllm.entrypoints.openai.api_server",
        "--model", MODEL_ID,
        "--served-model-name", "glm-ocr",
        "--port", str(OCR_API_PORT),
        "--max-model-len", VLLM_MAX_MODEL_LEN,
        "--gpu-memory-utilization", VLLM_GPU_MEMORY_UTIL,
    ]
    if VLLM_EXTRA_ARGS:
        cmd.extend(VLLM_EXTRA_ARGS.split())

    log.info("Starting vLLM server: %s", " ".join(cmd))
    _vllm_proc = subprocess.Popen(cmd)

    deadline = time.time() + 420
    last_err: str = ""
    while time.time() < deadline:
        if _vllm_proc.poll() is not None:
            raise RuntimeError("vLLM server exited early. Check Space logs for details.")
        try:
            resp = httpx.get(f"{OCR_API_BASE}/v1/models", timeout=8.0)
            if resp.status_code == 200:
                log.info("vLLM is ready.")
                return
        except Exception as e:
            last_err = str(e)
        time.sleep(3)
    raise RuntimeError(f"Timed out waiting for vLLM startup. Last error: {last_err}")


def _configure_glmocr_to_keep_header_footer() -> None:
    import glmocr

    config_path = os.path.join(os.path.dirname(glmocr.__file__), "config.yaml")
    with open(config_path, "r", encoding="utf-8") as f:
        cfg = yaml.safe_load(f) or {}

    pipeline = cfg.setdefault("pipeline", {})
    maas = pipeline.setdefault("maas", {})
    maas["enabled"] = False
    ocr_api = pipeline.setdefault("ocr_api", {})
    ocr_api["api_host"] = "127.0.0.1"
    ocr_api["api_port"] = OCR_API_PORT
    ocr_api["model"] = "glm-ocr"
    ocr_api["api_path"] = "/v1/chat/completions"
    ocr_api["api_mode"] = "openai"
    ocr_api["verify_ssl"] = False
    page_loader = pipeline.setdefault("page_loader", {})
    page_loader["max_tokens"] = GLMOCR_MAX_OUTPUT_TOKENS
    layout = pipeline.setdefault("layout", {})
    label_task_mapping = layout.setdefault("label_task_mapping", {})
    text_labels = set(label_task_mapping.get("text", []) or [])
    abandon_labels = set(label_task_mapping.get("abandon", []) or [])
    skip_labels = set(label_task_mapping.get("skip", []) or [])
    text_labels.update({"header", "footer"})
    abandon_labels.discard("header")
    abandon_labels.discard("footer")
    skip_labels.update({"header_image", "footer_image"})
    abandon_labels.discard("header_image")
    abandon_labels.discard("footer_image")
    label_task_mapping["text"] = sorted(text_labels)
    label_task_mapping["abandon"] = sorted(abandon_labels)
    label_task_mapping["skip"] = sorted(skip_labels)

    with open(config_path, "w", encoding="utf-8") as f:
        yaml.safe_dump(cfg, f, sort_keys=False)


def get_parser():
    global _parser
    if _parser is None:
        _start_vllm_if_needed()
        _configure_glmocr_to_keep_header_footer()
        from glmocr import GlmOcr
        _parser = GlmOcr(mode="selfhosted")
    return _parser


def _extract_md(result: Any) -> str:
    if result is None:
        return ""

    def _close_unbalanced_tables(chunk: str) -> str:
        opens = len(re.findall(r"<table\b", chunk, flags=re.IGNORECASE))
        closes = len(re.findall(r"</table>", chunk, flags=re.IGNORECASE))
        missing = max(0, opens - closes)
        if missing:
            chunk = chunk.rstrip() + ("\n</table>" * missing)
        return chunk

    if isinstance(result, list):
        chunks = []
        for item in result:
            md = getattr(item, "markdown_result", "")
            if md:
                chunks.append(_close_unbalanced_tables(str(md).strip()))
        return "\n\n---page-separator---\n\n".join([c for c in chunks if c]).strip()
    md = getattr(result, "markdown_result", "")
    return _close_unbalanced_tables(str(md).strip()) if md else ""


def _get_cell_text(cell) -> str:
    """Extract clean text from a BeautifulSoup table cell."""
    # Replace <br> tags with space before getting text
    for br in cell.find_all("br"):
        br.replace_with(" ")
    text = cell.get_text(separator=" ", strip=True)
    # Collapse multiple spaces
    text = re.sub(r"\s+", " ", text).strip()
    return text


def _html_table_to_markdown(table_html: str) -> str:
    """
    Convert a single HTML table to a clean markdown table.
    - Removes empty columns that only contain whitespace
    - Handles colspan/rowspan by ignoring them (takes first value)
    - Works for both regular tables and bordered tables
    """
    try:
        from bs4 import BeautifulSoup
        soup = BeautifulSoup(table_html, "html.parser")
        table = soup.find("table")
        if not table:
            return table_html
    except ImportError:
        # Fallback: regex-based if bs4 not available
        return _html_table_to_markdown_regex(table_html)

    # Collect all rows from thead + tbody + direct tr
    all_rows = []
    for tr in table.find_all("tr"):
        cells = []
        for cell in tr.find_all(["td", "th"]):
            cells.append(_get_cell_text(cell))
        if cells:
            all_rows.append(cells)

    if not all_rows:
        return table_html

    # Normalize all rows to same width
    max_cols = max(len(r) for r in all_rows)
    all_rows = [r + [""] * (max_cols - len(r)) for r in all_rows]

    # Remove columns that are empty in ALL rows
    non_empty_cols = []
    for col_idx in range(max_cols):
        if any(all_rows[row_idx][col_idx] for row_idx in range(len(all_rows))):
            non_empty_cols.append(col_idx)

    all_rows = [[row[i] for i in non_empty_cols] for row in all_rows]

    if not all_rows:
        return table_html

    num_cols = len(all_rows[0])
    md_lines = []

    # First row as header
    md_lines.append("| " + " | ".join(all_rows[0]) + " |")
    md_lines.append("| " + " | ".join(["---"] * num_cols) + " |")

    # Remaining rows
    for row in all_rows[1:]:
        # Pad or trim to match header width
        padded = (row + [""] * num_cols)[:num_cols]
        md_lines.append("| " + " | ".join(padded) + " |")

    return "\n".join(md_lines)


def _html_table_to_markdown_regex(table_html: str) -> str:
    """Regex-based fallback for HTML table conversion when bs4 is unavailable."""
    row_re = re.compile(r"<tr[^>]*>(.*?)</tr>", re.IGNORECASE | re.DOTALL)
    cell_re = re.compile(r"<t[dh][^>]*>(.*?)</t[dh]>", re.IGNORECASE | re.DOTALL)
    tag_re = re.compile(r"<[^>]+>")

    def clean_cell(c: str) -> str:
        c = re.sub(r"<br\s*/?>", " ", c, flags=re.IGNORECASE)
        c = tag_re.sub("", c)
        c = c.replace("&nbsp;", " ").replace("&amp;", "&").replace("&gt;", ">").replace("&lt;", "<")
        return re.sub(r"\s+", " ", c).strip()

    rows = []
    for tr in row_re.findall(table_html):
        cells = [clean_cell(c) for c in cell_re.findall(tr)]
        if cells:
            rows.append(cells)

    if not rows:
        return table_html

    max_cols = max(len(r) for r in rows)
    rows = [r + [""] * (max_cols - len(r)) for r in rows]

    # Remove empty columns
    non_empty_cols = [i for i in range(max_cols) if any(rows[j][i] for j in range(len(rows)))]
    rows = [[r[i] for i in non_empty_cols] for r in rows]

    if not rows:
        return table_html

    num_cols = len(rows[0])
    md = ["| " + " | ".join(rows[0]) + " |", "| " + " | ".join(["---"] * num_cols) + " |"]
    for row in rows[1:]:
        padded = (row + [""] * num_cols)[:num_cols]
        md.append("| " + " | ".join(padded) + " |")

    return "\n".join(md)


def _convert_all_html_tables(md: str) -> str:
    """Find and replace all HTML tables in the markdown with clean markdown tables."""
    table_re = re.compile(r"<table[\s\S]*?</table>", re.IGNORECASE)

    def replace_table(match: re.Match) -> str:
        return _html_table_to_markdown(match.group(0))

    return table_re.sub(replace_table, md)


def _clean_markdown(md: str) -> str:
    if not md:
        return md

    # Convert all HTML tables to clean markdown tables first
    md = _convert_all_html_tables(md)

    # Drop frequently repeated bank footer/header lines
    drop_line_patterns = [
        r"Call 1-800-937-2000 for 24-hour Bank-by-Phone services.*",
        r"Bank Deposits FDIC Insured \| TD Bank, N\.A\. \| Equal Housing Lender.*",
        r"Go paperless\.",
        r"Scan the QR code to opt in to paperless statements\.",
    ]
    drop_line_regexes = [re.compile(p, re.IGNORECASE) for p in drop_line_patterns]

    cleaned_lines: List[str] = []
    for line in md.splitlines():
        s = line.strip()
        if any(rx.fullmatch(s) for rx in drop_line_regexes):
            continue
        cleaned_lines.append(line)
    md = "\n".join(cleaned_lines)

    # Remove long legal notice sections
    legal_start_markers = [
        "## FOR CONSUMER ACCOUNTS ONLY",
        "## FOR CONSUMER LOAN ACCOUNTS ONLY",
        "## In case of Errors or Questions About Your Bill:",
    ]
    lines = md.splitlines()
    kept: List[str] = []
    skipping = False
    for line in lines:
        s = line.strip()
        if any(s.startswith(m) for m in legal_start_markers):
            skipping = True
            continue
        if skipping and s == "---page-separator---":
            skipping = False
            kept.append(line)
            continue
        if not skipping:
            kept.append(line)
    md = "\n".join(kept)

    # Remove duplicate separators and excessive blank lines
    md = re.sub(r"(?:\n\s*){3,}", "\n\n", md)
    md = re.sub(
        r"(?:\n\s*---page-separator---\s*\n){2,}",
        "\n\n---page-separator---\n\n",
        md,
    )
    return md.strip()


def run_ocr(file_obj):
    if file_obj is None:
        return "Please upload a file."
    path = file_obj.name if hasattr(file_obj, "name") else str(file_obj)
    page_images: List[str] = []
    try:
        parser = get_parser()
        is_pdf = path.lower().endswith(".pdf")
        if is_pdf:
            page_images = _render_pdf_pages_to_images(path)
            result = parser.parse(page_images)
        else:
            result = parser.parse(path)
        md = _clean_markdown(_extract_md(result)) or "(No content)"
        return md
    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


def build_demo():
    with gr.Blocks(title="GLM-OCR Self-hosted (Header/Footer enabled)") as demo:
        gr.Markdown("# GLM-OCR Self-hosted (Header/Footer enabled)")
        gr.Markdown(
            "Runs a local vLLM server inside the Space and configures GLM-OCR "
            "to include `header` and `footer` regions in OCR output."
        )
        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_md = gr.Textbox(label="Markdown", lines=30)
        run_btn.click(fn=run_ocr, inputs=file_in, outputs=out_md)
    return demo


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