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

# ── STEP 1: Fix config ────────────────────────────────────────
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")

# ── 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: Helper functions ──────────────────────────────────
def get_top_gap(regions, img_height):
    """
    Find the top gap the API missed using bbox_2d coordinates.
    Only extracts header (top gap) — never bottom — to avoid
    duplicating body content the API already captured correctly.
    """
    y_tops = []
    for r in regions:
        bbox = r.get("bbox_2d")
        if bbox and len(bbox) == 4:
            y_tops.append(bbox[1])
    if not y_tops:
        return None
    first_y = min(y_tops)
    return (first_y / img_height) if first_y > img_height * 0.08 else None


def extract_zone_text(pdf_path, page_num, y_start_frac, y_end_frac):
    """Extract text from a vertical slice of a PDF page."""
    import pymupdf as fitz
    doc  = fitz.open(pdf_path)
    page = doc[page_num]
    h, w = page.rect.height, page.rect.width
    text = page.get_text(clip=fitz.Rect(0, h * y_start_frac, w, h * y_end_frac)).strip()
    doc.close()
    return text


def get_page_data(page_result):
    """Extract markdown and raw regions from one PipelineResult."""
    regions = []
    if hasattr(page_result, "json_result"):
        jr = page_result.json_result
        if isinstance(jr, list) and len(jr) > 0:
            r = jr[0] if isinstance(jr[0], list) else jr
            if isinstance(r, list):
                regions = r
    md = ""
    if hasattr(page_result, "markdown_result") and page_result.markdown_result:
        md = page_result.markdown_result.strip()
    return md, regions


# ── STEP 4: Main OCR function ─────────────────────────────────
def run_ocr(uploaded_file):
    if uploaded_file is None:
        return "Please upload a file."

    try:
        from glmocr import parse
        import pymupdf as fitz

        path   = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)
        is_pdf = path.lower().endswith(".pdf")

        if is_pdf:
            doc          = fitz.open(path)
            page_images  = []
            page_heights = []
            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)
                # Image height at 1.5x matches bbox_2d coordinate space
                page_heights.append(doc[i].rect.height * 1.5)
            doc.close()
            results = parse(page_images)
            if not isinstance(results, list):
                results = [results]
        else:
            page_heights = []
            results = parse(path)
            if not isinstance(results, list):
                results = [results]

        all_pages_md = []

        for page_num, page_result in enumerate(results):
            page_md, regions = get_page_data(page_result)
            parts = []

            if is_pdf and page_num < len(page_heights):
                top_gap = get_top_gap(regions, page_heights[page_num])

                # HEADER at TOP — from PDF text layer (API misses this zone)
                if top_gap is not None:
                    hdr = extract_zone_text(path, page_num, 0, top_gap)
                    if hdr:
                        parts.append(hdr)

                # BODY — from API (includes all content + footer text correctly)
                if page_md:
                    parts.append(page_md)

                # No bottom zone extraction — API already captures footer content.
                # Extracting bottom zone causes duplication of body content.
            else:
                if page_md:
                    parts.append(page_md)

            if parts:
                all_pages_md.append("\n\n".join(parts))

        # Exactly like original SDK output format
        final = "\n\n---\n\n".join(all_pages_md) if all_pages_md else "(No content)"
        return final

    except Exception as e:
        import traceback
        return "Error: " + str(e) + "\n\n" + traceback.format_exc()


# ── STEP 5: Gradio UI ─────────────────────────────────────────
with gr.Blocks(title="GLM-OCR") as demo:
    gr.Markdown("""
    # 🔍 GLM-OCR
    Upload a PDF or image to extract text.
    Headers included at top of each page.
    """)

    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")
    output_box = gr.Textbox(lines=40, label="Output")

    run_btn.click(fn=run_ocr, inputs=file_input, outputs=output_box)

demo.launch()