#!/usr/bin/env python3

import json
import logging
import os
import subprocess
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 = os.getenv("GLMOCR_MODEL_ID", "zai-org/GLM-OCR")
OCR_API_PORT = int(os.getenv("OCR_API_PORT", "8080"))
OCR_API_BASE = f"http://127.0.0.1:{OCR_API_PORT}"
VLLM_MAX_MODEL_LEN = os.getenv("VLLM_MAX_MODEL_LEN", "8192")
VLLM_GPU_MEMORY_UTIL = os.getenv("VLLM_GPU_MEMORY_UTIL", "0.92")
VLLM_EXTRA_ARGS = os.getenv("VLLM_EXTRA_ARGS", "").strip()

_parser = None
_vllm_proc = None


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

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

    # Keep image-only decorative zones out of OCR by default.
    # If you want logo OCR too, move these to text_labels.
    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()
    return _parser


def _extract_md(result: Any) -> str:
    if result is None:
        return ""
    if isinstance(result, list):
        chunks = []
        for item in result:
            md = getattr(item, "markdown_result", "")
            if md:
                chunks.append(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 str(md).strip() if md else ""


def _extract_json(result: Any) -> str:
    if result is None:
        return "{}"
    if isinstance(result, list):
        payload = [getattr(item, "json_result", None) for item in result]
        return json.dumps(payload, ensure_ascii=False, indent=2)
    return json.dumps(getattr(result, "json_result", None), ensure_ascii=False, indent=2)


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)
    try:
        parser = get_parser()
        result = parser.parse(path)
        md = _extract_md(result) or "(No content)"
        jr = _extract_json(result)
        return md, jr
    except Exception as e:
        import traceback

        log.exception("run_ocr failed: %s", e)
        return f"Error: {e}\n\n{traceback.format_exc()}", "{}"


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)
        out_json = gr.Textbox(label="JSON (layout details)", lines=20)
        run_btn.click(fn=run_ocr, inputs=file_in, outputs=[out_md, out_json])
    return demo


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