import os
import sys
import base64
import tempfile
import json
import traceback

# Install pinned runtime deps at startup (before importing torch/transformers)
_RUNTIME_PKGS = [
    "torch==2.10.0",
    "torchvision==0.25.0",
    "transformers==4.57.1",
    "PyMuPDF==1.26.1",
    "Pillow>=10.0.0",
]

print("Installing pinned runtime dependencies...", flush=True)
import subprocess
subprocess.run(
    [sys.executable, "-m", "pip", "install", "--quiet", "--no-cache-dir"] + _RUNTIME_PKGS,
    check=True,
)
print("Runtime deps installed.", flush=True)

import torch
from transformers import AutoModel, AutoTokenizer

MODEL_NAME = "baidu/Unlimited-OCR"

print("=== Pre-loading Baidu Unlimited-OCR model at startup ===", flush=True)
print("Loading tokenizer...", flush=True)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
print("Loading model (CPU)...", flush=True)
model = AutoModel.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
    use_safetensors=True,
    torch_dtype=torch.float32,
).eval()
print("OCR model ready (CPU).", flush=True)
print("=== OCR model pre-load complete ===", flush=True)

from mcp.server.fastmcp import FastMCP

mcp = FastMCP("unlimited-ocr")


def _collect_output(out_dir: str) -> str:
    result = ""
    for fname in sorted(os.listdir(out_dir)):
        if fname.endswith((".txt", ".md")):
            with open(os.path.join(out_dir, fname), "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()


def _run_ocr_internal(image_path: str, mode: str, prompt: str) -> str:
    out_dir = tempfile.mkdtemp(prefix="ocr_out_")

    if mode == "gundam":
        base_size, image_size, crop_mode, ngram_window = 1024, 640, True, 128
    else:
        base_size, image_size, crop_mode, ngram_window = 1024, 1024, False, 128

    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_window,
        save_results=True,
    )

    return _collect_output(out_dir)


@mcp.tool()
def run_ocr(
    image_base64: str = "",
    image_path: str = "",
    mode: str = "gundam",
    prompt: str = "document parsing.",
) -> str:
    """Run OCR on an image using Baidu Unlimited-OCR (CPU inference).

    Extracts text from images of documents, screenshots, signs, handwriting, etc.
    Supports multilingual text extraction with high accuracy.

    Args:
        image_base64: Base64-encoded image data (PNG/JPEG). Takes precedence if provided.
        image_path: File path to the image. Used if image_base64 is empty.
        mode: 'gundam' for fast mode (640px crop), 'base' for accurate mode (1024px).
        prompt: Instruction for the OCR model (e.g. "document parsing.", "read all text.").

    Returns:
        Extracted text from the image.
    """
    if image_base64:
        try:
            image_bytes = base64.b64decode(image_base64)
        except Exception as e:
            return f"Error decoding base64 image: {e}"
        tmp_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
        with open(tmp_path, "wb") as f:
            f.write(image_bytes)
    elif image_path:
        if not os.path.exists(image_path):
            return f"Error: Image file not found: {image_path}"
        tmp_path = image_path
    else:
        return "Error: Provide either image_base64 or image_path."

    try:
        result = _run_ocr_internal(tmp_path, mode, prompt)
        return result if result else "No text detected in image."
    except Exception as e:
        return f"OCR error: {e}"
    finally:
        if image_base64 and os.path.exists(tmp_path):
            os.unlink(tmp_path)


@mcp.tool()
def run_ocr_pdf(
    pdf_base64: str = "",
    pdf_path: str = "",
    mode: str = "gundam",
    prompt: str = "document parsing.",
) -> str:
    """Run OCR on a PDF document. Converts each page to an image and runs OCR.

    Args:
        pdf_base64: Base64-encoded PDF data. Takes precedence if provided.
        pdf_path: File path to the PDF. Used if pdf_base64 is empty.
        mode: 'gundam' for fast mode, 'base' for accurate mode.
        prompt: Instruction for the OCR model.

    Returns:
        Concatenated text from all pages.
    """
    import fitz

    if pdf_base64:
        try:
            pdf_bytes = base64.b64decode(pdf_base64)
        except Exception as e:
            return f"Error decoding base64 PDF: {e}"
        tmp_pdf = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name
        with open(tmp_pdf, "wb") as f:
            f.write(pdf_bytes)
    elif pdf_path:
        if not os.path.exists(pdf_path):
            return f"Error: PDF file not found: {pdf_path}"
        tmp_pdf = pdf_path
    else:
        return "Error: Provide either pdf_base64 or pdf_path."

    try:
        doc = fitz.open(tmp_pdf)
        tmp_dir = tempfile.mkdtemp(prefix="pdf_ocr_")
        mat = fitz.Matrix(200 / 72, 200 / 72)
        all_text = []
        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)
            text = _run_ocr_internal(out, mode, prompt)
            all_text.append(f"--- Page {i + 1} ---\n{text}")
            os.unlink(out)
        doc.close()
        return "\n\n".join(all_text) if all_text else "No text detected in PDF."
    except Exception as e:
        return f"PDF OCR error: {e}"
    finally:
        if pdf_base64 and os.path.exists(tmp_pdf):
            os.unlink(tmp_pdf)


if __name__ == "__main__":
    mcp.run(transport="streamable-http", host="0.0.0.0", port=7860)