import gradio as gr
import yaml
import re
import os

# Paths from installed glmocr package
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")

# ── STEP 1: Load and fix config (MaaS + 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)

print("✅ config.yaml fixed (MaaS enabled, header/footer kept)")

# ── STEP 2: Fix result_formatter.py ────────────────────────────────────────
with open(formatter_path, "r") as f:
    source = f.read()

labels_to_remove = [
    '"header"', "'header'", '"footer"', "'footer'",
    '"doc_header"', "'doc_header'", '"doc_footer"', "'doc_footer'"
]
for label in labels_to_remove:
    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)

print("✅ result_formatter.py fixed")

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

# ── STEP 3: OCR via Zhipu MaaS (no local model) ─────────────────────────────
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: ZHIPU_API_KEY is not set. Add it in Space Settings → Variables and secrets.",
            "Failed."
        )

    try:
        from glmocr import parse

        path = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)

        # PDF → list of image paths (SDK expects images per page)
        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()
            # parse() with list = one document, multiple pages
            result = parse(page_images)
        else:
            result = parse(path)

        # result.json_result: list of pages, each page = list of regions
        # Region: {"label": "header"|"footer"|"text"|..., "content": "...", ...}
        pages_data = result.json_result if hasattr(result, "json_result") else []
        if not isinstance(pages_data, list):
            pages_data = [pages_data] if pages_data else []

        total_headers = 0
        total_footers = 0
        all_markdown = []
        all_summary = []

        for page_num, page_regions in enumerate(pages_data):
            if not isinstance(page_regions, list):
                page_regions = [page_regions] if page_regions else []
            page_md = []
            page_summary = [f"── PAGE {page_num + 1} of {len(pages_data)} ──"]

            for region in page_regions:
                if not isinstance(region, dict):
                    continue
                label = region.get("label", "text")
                content = str(region.get("content", ""))

                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."

# ── STEP 4: Gradio UI ──────────────────────────────────────────────────────
with gr.Blocks(title="GLM-OCR — MaaS (Header & Footer)") as demo:
    gr.Markdown("""
    # 🔍 GLM-OCR — Zhipu MaaS
    Upload PDF or image. Headers 🔵 and Footers 🟢 from the **full pipeline** (layout + OCR).
    Uses your Zhipu API key (set in Space secrets).
    """)
    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()