import gradio as gr
import yaml
import re
import os

import glmocr
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")

# ── Config: MaaS on, keep header/footer ────────────────────────────────────
with open(config_path, "r") as f:
    config = yaml.safe_load(f)
config["pipeline"]["maas"]["enabled"] = True
config["pipeline"]["maas"]["api_key"] = "4570c28bdea5493c9efae9dae68edc66.sGbA9DLlcX1GlvqV"
config["pipeline"]["result_formatter"]["abandon"] = [
    "number", "footnote", "aside_text", "reference", "footer_image", "header_image",
]
config["pipeline"]["enable_layout"] = True
with open(config_path, "w") as f:
    yaml.dump(config, f, default_flow_style=False, sort_keys=False)

# ── Formatter: don't strip header/footer ───────────────────────────────────
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)

ABANDON = set(config["pipeline"]["result_formatter"]["abandon"])

def get_regions(page_item):
    if page_item is None:
        return []
    if isinstance(page_item, list):
        return page_item
    if isinstance(page_item, dict):
        for key in ("regions", "blocks", "items", "content_list"):
            if key in page_item and isinstance(page_item[key], list):
                return page_item[key]
        if "content" in page_item or "text" in page_item:
            return [page_item]
    if isinstance(page_item, str):
        return [{"label": "text", "content": page_item}]
    return []

def get_label(r):
    if not isinstance(r, dict):
        return "text"
    return (r.get("label") or r.get("type") or r.get("block_type") or "text")

def get_content(r):
    if isinstance(r, str):
        return r
    if not isinstance(r, dict):
        return str(r)
    return str(r.get("content") or r.get("text") or r.get("raw") or "")

def run_ocr(uploaded_file):
    if uploaded_file is None:
        return "Please upload a file.", "No regions detected."
    api_key = "4570c28bdea5493c9efae9dae68edc66.sGbA9DLlcX1GlvqV"
    if not api_key:
        return "Error: API key not set.", "Failed."
    try:
        from glmocr import parse
        path = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)
        if path.lower().endswith(".pdf"):
            import fitz
            doc = fitz.open(path)
            page_images = []
            for i in range(len(doc)):
                pix = doc[i].get_pixmap(matrix=fitz.Matrix(1.5, 1.5), alpha=False)
                img_path = f"/tmp/maas_page_{i}.png"
                pix.save(img_path)
                page_images.append(img_path)
            doc.close()
            result = parse(page_images)
        else:
            result = parse(path)
        pages_data = result if isinstance(result, list) else []
        total_headers = 0
        total_footers = 0
        all_markdown = []
        all_summary = []
        for page_num, page_item in enumerate(pages_data):
            page_regions = get_regions(page_item)
            if not isinstance(page_regions, list):
                page_regions = [page_regions]
            page_md = []
            page_summary = [f"── PAGE {page_num + 1} of {len(pages_data)} ──"]
            for region in page_regions:
                label = get_label(region)
                content = get_content(region)
                if not content.strip():
                    continue
                if label in ABANDON:
                    continue
                if label == "header":
                    total_headers += 1
                    page_summary.append(f"🔵 HEADER: {content[:100]}")
                    page_md.append(f"<!-- HEADER -->\n{content}")
                elif label == "footer":
                    total_footers += 1
                    page_summary.append(f"🟢 FOOTER: {content[:100]}")
                    page_md.append(f"<!-- FOOTER -->\n{content}")
                else:
                    page_summary.append(f"[{label}]: {content[:100]}")
                    page_md.append(content)
            all_markdown.append("\n\n".join(page_md))
            all_summary.extend(page_summary)
            all_summary.append("")
        summary = (
            f"Total pages   : {len(pages_data)}\n"
            f"Headers found : {total_headers}\n"
            f"Footers found : {total_footers}\n"
            f"{'─'*40}\n" + "\n".join(all_summary)
        )
        markdown = "\n\n---\n\n".join(all_markdown) if all_markdown else "(No content)"
        return markdown, summary
    except Exception as e:
        import traceback
        return f"Error: {str(e)}\n\n{traceback.format_exc()}", "Failed."

with gr.Blocks(title="GLM-OCR — MaaS") as demo:
    gr.Markdown("# 🔍 GLM-OCR — Zhipu MaaS\nUpload PDF or image. Headers 🔵 and Footers 🟢 when returned by the API.")
    file_input = gr.File(label="Upload PDF or Image", file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"])
    run_btn = gr.Button("▶ Run OCR", variant="primary", size="lg")
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 📄 Markdown Output")
            markdown_out = gr.Textbox(lines=30, label="")
        with gr.Column():
            gr.Markdown("### 🗂️ Detected Regions")
            regions_out = gr.Textbox(lines=30, label="")
    run_btn.click(fn=run_ocr, inputs=file_input, outputs=[markdown_out, regions_out])
demo.launch()