import gradio as gr
import yaml
import re
import json
import os
import torch

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: Fix config.yaml ───────────────────────────────────
with open(config_path, "r") as f:
    config = yaml.safe_load(f)

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

# ── 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")

# ── STEP 3: Load model at startup ────────────────────────────
from transformers import AutoProcessor, GlmOcrForConditionalGeneration

print("Loading model at startup...")
processor = AutoProcessor.from_pretrained("zai-org/GLM-OCR")
model = GlmOcrForConditionalGeneration.from_pretrained(
    "zai-org/GLM-OCR",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
print(f"✅ Model ready on {next(model.parameters()).device}")

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

# ── STEP 4: OCR one image ─────────────────────────────────────
def ocr_image(img_path):
    messages = [
        {"role": "user", "content": [
            {"type": "image", "url": img_path},
            {"type": "text",  "text": "Document Parsing:"}
        ]}
    ]
    inputs = processor.apply_chat_template(
        messages, tokenize=True, add_generation_prompt=True,
        return_dict=True, return_tensors="pt"
    ).to(model.device)
    inputs.pop("token_type_ids", None)

    with torch.no_grad():
        output_ids = model.generate(**inputs, max_new_tokens=2048)

    raw = processor.decode(
        output_ids[0][inputs["input_ids"].shape[1]:],
        skip_special_tokens=False
    )

    for token in ["<|assistant|>", "<|user|>", "<|system|>",
                  "<|endoftext|>", "</s>", "<s>"]:
        raw = raw.replace(token, "")
    raw = raw.strip()

    json_match = re.search(r'\[.*\]', raw, re.DOTALL)
    if json_match:
        try:
            regions = json.loads(json_match.group())
            return regions, raw, "json"
        except:
            pass

    return [], raw, "raw"

# ── STEP 5: Main OCR function ─────────────────────────────────
def run_ocr(uploaded_file):
    if uploaded_file is None:
        return "Please upload a file.", "No regions detected."

    try:
        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/page_{i}.png"
                pix.save(img_path)
                page_images.append(img_path)
            doc.close()
        else:
            page_images = [path]

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

        for page_num, img_path in enumerate(page_images):
            regions, raw, output_type = ocr_image(img_path)

            page_md      = []
            page_summary = [f"── PAGE {page_num + 1} of {len(page_images)} ──"]

            if output_type == "raw":
                page_md.append(raw)
                page_summary.append("[raw HTML/text output]")
                page_summary.append(f"Content length: {len(raw)} chars")
            else:
                for region in regions:
                    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(page_images)}\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)
        return markdown, summary

    except Exception as e:
        import traceback
        return f"Error: {str(e)}\n\n{traceback.format_exc()}", "Failed."

# ── STEP 6: Gradio UI ─────────────────────────────────────────
with gr.Blocks(title="GLM-OCR — Header & Footer Fix") as demo:
    gr.Markdown("""
    # 🔍 GLM-OCR — Header & Footer Fix
    Upload PDF or image. Headers 🔵 and Footers 🟢 now appear in output.
    Multi-page PDFs fully supported.
    > First request takes ~2 min to load model. After that it is fast.
    """)

    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()