#!/usr/bin/env python3
"""
Unlimited OCR — ZeroGPU Space with REST API.
Baidu's baidu/Unlimited-OCR on free GPU (ZeroGPU, no credits).
Streaming OCR output, PDF support, callable as REST API.
"""

import subprocess, sys, os, tempfile, shutil
from threading import Thread
from typing import Iterator
import queue, threading

# ── Runtime install ──────────────────────────────────────────────────────
_RUNTIME_PKGS = [
    "torch==2.10.0",
    "torchvision==0.25.0",
    "transformers==4.57.1",
    "Pillow>=12.0",
    "einops>=0.8",
    "addict>=2.4",
    "easydict>=1.13",
    "pymupdf>=1.27",
    "psutil>=7.2",
]
print("[setup] Installing pinned dependencies...")
subprocess.run(
    [sys.executable, "-m", "pip", "install", "--quiet", "--no-cache-dir"]
    + _RUNTIME_PKGS,
    check=True,
)
print("[setup] Done ✓")

# ── Imports ──────────────────────────────────────────────────────────────
import torch
from transformers import AutoModel, AutoTokenizer, TextIteratorStreamer
import gradio as gr
import spaces

# ── Model loading (ZeroGPU provides GPU at startup in emulation mode) ────
MODEL_NAME = "baidu/Unlimited-OCR"

print("[model] Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
print("[model] Loading model...")
model = AutoModel.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
    use_safetensors=True,
    torch_dtype=torch.bfloat16,
).eval().cuda()
print("[model] Ready ✓")


# ── Helpers ──────────────────────────────────────────────────────────────
def pdf_to_images(pdf_path: str, dpi: int = 200) -> list[str]:
    """Convert PDF pages to images."""
    import fitz
    doc = fitz.open(pdf_path)
    tmp_dir = tempfile.mkdtemp(prefix="pdf_ocr_")
    mat = fitz.Matrix(dpi / 72, dpi / 72)
    paths = []
    for i, page in enumerate(doc):
        out = os.path.join(tmp_dir, f"page_{i + 1:04d}.png")
        page.get_pixmap(matrix=mat).save(out)
        paths.append(out)
    doc.close()
    return paths


def _collect_output(out_dir: str) -> str:
    """Read text files written by model.infer()."""
    result = ""
    for fname in sorted(os.listdir(out_dir)):
        fpath = os.path.join(out_dir, fname)
        if fname.endswith((".txt", ".md")):
            with open(fpath, "r", encoding="utf-8") as f:
                result += f.read() + "\n"
    if not result:
        for fname in sorted(os.listdir(out_dir)):
            fpath = os.path.join(out_dir, fname)
            if os.path.isfile(fpath):
                try:
                    with open(fpath, "r", encoding="utf-8") as f:
                        result += f.read() + "\n"
                except Exception:
                    pass
    return result.strip()


class _ThreadStdout:
    """Redirect stdout from inference thread to a queue for streaming."""
    def __init__(self, target_thread, q, original):
        self.target = target_thread
        self.q = q
        self.original = original

    def write(self, data):
        self.original.write(data)
        self.original.flush()
        if threading.current_thread() == self.target and data:
            low = data.lower()
            if "tps:" in low or "tokens/s" in low:
                return len(data)
            self.q.put(data)
        return len(data)

    def flush(self):
        self.original.flush()

    def __getattr__(self, name):
        return getattr(self.original, name)


# ── OCR with streaming ──────────────────────────────────────────────────
@spaces.GPU(duration=120)
def ocr_single_stream(image_path: str, mode: str, prompt: str) -> Iterator[dict]:
    """Stream OCR output for a single image. Yields {text, done}."""
    out_dir = tempfile.mkdtemp(prefix="ocr_out_")
    try:
        if mode == "gundam":
            base_size, image_size, crop_mode, ngram_win = 1024, 640, True, 128
        else:
            base_size, image_size, crop_mode, ngram_win = 1024, 1024, False, 128

        q = queue.Queue()
        errors = []

        def _run():
            try:
                model.infer(
                    tokenizer,
                    prompt=f"<image>{prompt}",
                    image_file=image_path,
                    output_path=out_dir,
                    base_size=base_size,
                    image_size=image_size,
                    crop_mode=crop_mode,
                    max_length=8192,
                    no_repeat_ngram_size=35,
                    ngram_window=ngram_win,
                    save_results=True,
                )
            except Exception as e:
                errors.append(str(e))

        t = Thread(target=_run, daemon=True)
        orig_stdout = sys.stdout
        sys.stdout = _ThreadStdout(t, q, orig_stdout)

        acc = ""
        try:
            t.start()
            while t.is_alive() or not q.empty():
                try:
                    chunk = q.get(timeout=0.05)
                    acc += chunk
                    yield {"text": acc, "done": False}
                except queue.Empty:
                    continue
        finally:
            sys.stdout = orig_stdout
            t.join()

        full = _collect_output(out_dir)
        text = full if full else acc
        yield {"text": text, "done": True}
    finally:
        shutil.rmtree(out_dir, ignore_errors=True)


@spaces.GPU(duration=180)
def ocr_multi(image_paths: list[str], prompt: str) -> str:
    """Multi-image OCR (no streaming for simplicity)."""
    out_dir = tempfile.mkdtemp(prefix="ocr_multi_")
    try:
        model.infer_multi(
            tokenizer,
            prompt=f"<image>{prompt}",
            image_files=image_paths,
            output_path=out_dir,
            image_size=1024,
            max_length=32768,
            no_repeat_ngram_size=35,
            ngram_window=1024,
            save_results=True,
        )
        return _collect_output(out_dir)
    finally:
        shutil.rmtree(out_dir, ignore_errors=True)


# ── Gradio handlers ─────────────────────────────────────────────────────
def handle_single(image, mode, prompt):
    if image is None:
        return "⚠️ Upload an image."
    path = image if isinstance(image, str) else getattr(image, "name", str(image))
    try:
        return list(ocr_single_stream(path, mode, prompt))[-1]["text"]
    except Exception as e:
        return f"❌ {type(e).__name__}: {e}"


def handle_single_stream(image, mode, prompt):
    """Streaming generator for Gradio."""
    if image is None:
        yield {"text": "⚠️ Upload an image.", "done": True}
        return
    path = image if isinstance(image, str) else getattr(image, "name", str(image))
    try:
        yield from ocr_single_stream(path, mode, prompt)
    except Exception as e:
        yield {"text": f"❌ {type(e).__name__}: {e}", "done": True}


def handle_multi(images, prompt):
    if not images:
        return "⚠️ Upload images."
    paths = [img if isinstance(img, str) else getattr(img, "name", str(img)) for img in images]
    try:
        return ocr_multi(paths, prompt)
    except Exception as e:
        return f"❌ {type(e).__name__}: {e}"


def handle_pdf(pdf_file, prompt):
    if pdf_file is None:
        return "⚠️ Upload a PDF."
    pdf_path = pdf_file if isinstance(pdf_file, str) else getattr(pdf_file, "name", str(pdf_file))
    page_imgs = []
    try:
        page_imgs = pdf_to_images(pdf_path)
        if not page_imgs:
            return "⚠️ No pages found."
        return ocr_multi(page_imgs, prompt)
    except Exception as e:
        return f"❌ {type(e).__name__}: {e}"
    finally:
        for p in page_imgs:
            try: os.remove(p)
            except: pass


# ── Gradio UI ───────────────────────────────────────────────────────────
with gr.Blocks(title="Unlimited OCR API", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
# 🔍 Unlimited OCR API
[Baidu Unlimited-OCR](https://huggingface.co/baidu/Unlimited-OCR) — one-shot long-horizon document parsing.
Free via **ZeroGPU** — no inference credits consumed.
📡 **API:** click "Use via API" in the footer for REST endpoints.
""")

    with gr.Tabs():
        with gr.TabItem("📄 Single Image"):
            with gr.Row():
                with gr.Column():
                    si = gr.Image(label="Image", type="filepath")
                    sm = gr.Radio(["gundam", "base"], value="gundam", label="Mode",
                        info="gundam=fast(640px) | base=accurate(1024px)")
                    sp = gr.Textbox(value="document parsing.", label="Prompt")
                    sb = gr.Button("🚀 OCR", variant="primary")
                with gr.Column():
                    so = gr.Textbox(label="Result", lines=20, show_copy_button=True)
            sb.click(handle_single, inputs=[si, sm, sp], outputs=so)

        with gr.TabItem("📚 Multi Image"):
            with gr.Row():
                with gr.Column():
                    mi = gr.File(label="Images", file_count="multiple", file_types=["image"])
                    mp = gr.Textbox(value="Multi page parsing.", label="Prompt")
                    mb = gr.Button("🚀 Multi-OCR", variant="primary")
                with gr.Column():
                    mo = gr.Textbox(label="Result", lines=20, show_copy_button=True)
            mb.click(handle_multi, inputs=[mi, mp], outputs=mo)

        with gr.TabItem("📑 PDF"):
            with gr.Row():
                with gr.Column():
                    pi = gr.File(label="PDF", file_types=[".pdf"])
                    pp = gr.Textbox(value="Multi page parsing.", label="Prompt")
                    pb = gr.Button("🚀 PDF OCR", variant="primary")
                with gr.Column():
                    po = gr.Textbox(label="Result", lines=20, show_copy_button=True)
            pb.click(handle_pdf, inputs=[pi, pp], outputs=po)

    gr.Markdown("""
---
### 📡 API Usage
```python
from gradio_client import Client
client = Client("Harry00/unlimited-ocr-api")

# Single image
result = client.predict("doc.png", "gundam", "document parsing.", api_name="/handle_single")

# PDF
result = client.predict("scan.pdf", "Multi page parsing.", api_name="/handle_pdf")
```
""")


# ── REST API endpoints (separate from Gradio UI) ───────────────────────
from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import JSONResponse

app_fastapi = FastAPI(title="Unlimited OCR API")

@app_fastapi.post("/ocr")
async def api_ocr(
    image: UploadFile = File(...),
    mode: str = Form("gundam"),
    prompt: str = Form("document parsing."),
):
    """OCR a single image. Returns extracted text."""
    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    try:
        tmp.write(await image.read())
        tmp.close()
        result = handle_single(tmp.name, mode, prompt)
        return {"text": result, "model": MODEL_NAME, "mode": mode}
    finally:
        os.unlink(tmp.name)


@app_fastapi.post("/ocr/pdf")
async def api_ocr_pdf(
    pdf: UploadFile = File(...),
    prompt: str = Form("Multi page parsing."),
):
    """OCR a PDF. Returns extracted text."""
    tmp = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False)
    try:
        tmp.write(await pdf.read())
        tmp.close()
        result = handle_pdf(tmp.name, prompt)
        return {"text": result, "model": MODEL_NAME}
    finally:
        os.unlink(tmp.name)


@app_fastapi.get("/health")
async def health():
    return {"status": "ok", "model": MODEL_NAME}


# ── Mount FastAPI under Gradio ──────────────────────────────────────────
app_fastapi = gr.mount_gradio_app(app_fastapi, demo, path="/")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app_fastapi, host="0.0.0.0", port=7860)