"""
🔱 Imperial OCR Space — GLM-OCR Edition
نموذج: zai-org/GLM-OCR
يدعم: Text Recognition / Formula / Table
API متوافق مع البوت: /predict(b64_img, lang) → text
"""

import os
import gc
import base64
import tempfile
from io import BytesIO
from threading import Thread

import gradio as gr
import torch
from PIL import Image, ImageOps
from transformers import (
    AutoProcessor,
    AutoModelForImageTextToText,
    TextIteratorStreamer,
)

# ─── إعدادات النموذج ──────────────────────────────────────────
MODEL_PATH  = "zai-org/GLM-OCR"
device      = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"🔱 OCR Space — جهاز: {device}", flush=True)

processor = AutoProcessor.from_pretrained(MODEL_PATH, trust_remote_code=True)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_PATH,
    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
    device_map="auto",
    trust_remote_code=True,
).eval()

# ─── خريطة المهام ─────────────────────────────────────────────
TASK_MAP = {
    "text":    "Text Recognition:",
    "formula": "Formula Recognition:",
    "table":   "Table Recognition:",
    # اللغات → نفس مهمة Text (GLM-OCR multilingual)
    "korean":  "Text Recognition:",
    "japanese":"Text Recognition:",
    "chinese": "Text Recognition:",
    "english": "Text Recognition:",
    "arabic":  "Text Recognition:",
    "spanish": "Text Recognition:",
}

MAX_NEW_TOKENS = 4096


# ─── الدالة الرئيسية (متوافقة مع البوت القديم) ───────────────
def extract_text(b64_img: str, lang: str = "korean") -> str:
    """
    الإدخال : base64 string للصورة (مع أو بدون data URI header)
    الخروج  : النص المستخرج كـ string
    
    lang يقبل: korean/japanese/chinese/english/arabic/spanish/text/formula/table
    """
    tmp_path = None
    try:
        # ── تنظيف الـ base64 ──────────────────────────────────
        if "," in b64_img:
            b64_img = b64_img.split(",", 1)[1]
        
        img_bytes = base64.b64decode(b64_img)
        img = Image.open(BytesIO(img_bytes)).convert("RGB")
        img = ImageOps.exif_transpose(img)

        # ── حفظ مؤقت (GLM-OCR يحتاج path) ────────────────────
        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
            img.save(tmp.name, "PNG")
            tmp_path = tmp.name

        # ── اختيار الـ prompt ─────────────────────────────────
        prompt = TASK_MAP.get(lang.lower().strip(), "Text Recognition:")

        messages = [{
            "role": "user",
            "content": [
                {"type": "image", "url": tmp_path},
                {"type": "text",  "text": prompt},
            ],
        }]

        inputs = processor.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_dict=True,
            return_tensors="pt",
        )
        inputs.pop("token_type_ids", None)
        inputs = {k: v.to(model.device) if hasattr(v, "to") else v
                  for k, v in inputs.items()}

        # ── توليد مع Streamer لتجنب الـ timeout ──────────────
        streamer = TextIteratorStreamer(
            processor.tokenizer if hasattr(processor, "tokenizer") else processor,
            skip_prompt=True,
            skip_special_tokens=True,
        )

        gen_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": MAX_NEW_TOKENS}
        gen_error  = {"e": None}

        def _run():
            try:
                model.generate(**gen_kwargs)
            except Exception as e:
                gen_error["e"] = e
                try: streamer.end()
                except: pass

        t = Thread(target=_run, daemon=True)
        t.start()

        output = ""
        for chunk in streamer:
            output += chunk
        t.join(timeout=2.0)

        if gen_error["e"]:
            return f"Error: {gen_error['e']}"

        return output.strip() if output.strip() else ""

    except Exception as e:
        return f"Error: {e}"
    finally:
        if tmp_path and os.path.exists(tmp_path):
            try: os.unlink(tmp_path)
            except: pass
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


# ─── واجهة Gradio (للاختبار اليدوي + API) ────────────────────
with gr.Blocks(title="🔱 Imperial OCR — GLM") as demo:
    gr.Markdown("## 🔱 Imperial OCR Space\n**GLM-OCR** — يدعم: كوري / ياباني / صيني / إنجليزي")

    with gr.Row():
        with gr.Column():
            inp_img  = gr.Image(type="pil",    label="الصورة")
            inp_lang = gr.Dropdown(
                choices=["korean","japanese","chinese","english","arabic","spanish","formula","table"],
                value="korean",
                label="اللغة / المهمة"
            )
            btn = gr.Button("استخراج النص 🔍")
        with gr.Column():
            out_text = gr.Textbox(label="النص المستخرج", lines=20)

    # زر للواجهة المرئية فقط
    def run_ui(img: Image.Image, lang: str) -> str:
        if img is None:
            return "❌ ارفع صورة أولاً"
        buf = BytesIO()
        img.save(buf, format="PNG")
        b64 = base64.b64encode(buf.getvalue()).decode()
        return extract_text(b64, lang)

    btn.click(fn=run_ui, inputs=[inp_img, inp_lang], outputs=out_text)

    # ── API متوافق مع البوت (نفس الـ endpoint القديم) ────────
    gr.Interface(
        fn=extract_text,
        inputs=[
            gr.Textbox(label="b64_img"),
            gr.Textbox(label="lang", value="korean"),
        ],
        outputs=gr.Textbox(label="text"),
        api_name="predict",
        title="API",
        description="متوافق مع البوت — يقبل base64 ويرجع نص",
    )

if __name__ == "__main__":
    demo.queue(max_size=10).launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
    )