import gradio as gr
from transformers import AutoProcessor, AutoModelForImageTextToText
import torch
from PIL import Image

# --- KONFIGURASI INTI ---
MODEL_PATH = "zai-org/GLM-OCR"

# 1. HUKUM KEKALAN HARDWARE (Auto Deteksi CPU/GPU)
if torch.cuda.is_available():
    device = "cuda"
    dtype = torch.float16
else:
    device = "cpu"
    dtype = torch.float32

print(f"🚀 ENGINE STARTED: Device={device} | Dtype={dtype}")

# 2. INISIASI MODEL (DENGAN SAFE LOAD)
try:
    print("⏳ Menyiapkan Otak GLM...")
    
    # Processor (Penyortir Data)
    processor = AutoProcessor.from_pretrained(
        MODEL_PATH, 
        trust_remote_code=True
    )
    
    # Model (Otak Utama) - Balik ke AutoModelForImageTextToText karena Library GIT udah support
    model = AutoModelForImageTextToText.from_pretrained(
        MODEL_PATH,
        torch_dtype=dtype,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
        device_map="auto"
    )
    
    # Mode Hemat Energi (Evaluasi)
    model.eval()

except Exception as e:
    print(f"⚠️ Warning Model Loading: {e}")
    # Biasanya warning doang, gas terus.
    pass

# 3. PROSES REAKSI FISIKA (Input -> Tensor -> Output)
def proses_intelijen(image):
    if image is None:
        return "⚠️ Upload dulu gambarnya Bos! Jangan scan angin."

    # Format Pesan Sesuai Standar GLM
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": "Text Recognition:"}
            ],
        }
    ]

    try:
        # --- PERBAIKAN FATAL ADA DISINI ---
        # tokenize=True adalah KUNCINYA. Biar dia jadi Tensor (Angka), bukan Teks doang.
        inputs = processor.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,      # <--- INI YG DULU KURANG!
            return_dict=True,
            return_tensors="pt" # Pastikan keluar format PyTorch
        )
        
        # Buang data sampah yang gak perlu
        inputs.pop("token_type_ids", None)
        
        # Pindahkan ke Mesin (Device) secara Manual biar aman
        inputs = {k: v.to(device) for k, v in inputs.items()}

        # GEBER MESINNYA (GENERATE)
        with torch.no_grad():
            generated_ids = model.generate(
                **inputs, 
                max_new_tokens=2048, # Naikkin dikit biar dokumen panjang kebaca
                do_sample=False
            )

        # DECODE (Angka -> Teks Manusia)
        # Kita potong input promptnya, ambil hasilnya doang
        hasil_potong = generated_ids[0][inputs["input_ids"].shape[1]:]
        teks_final = processor.decode(hasil_potong, skip_special_tokens=False) # Coba False dulu biar detail
        
        return teks_final

    except Exception as e:
        return f"🚨 SYSTEM CRITICAL FAILURE: {str(e)}\n(Kemungkinan memori tidak cukup atau input rusak)"

# 4. ANTARMUKA 2026 (UI yang Elegan tapi BODOH-PROOF)
css_style = """
.container { max-width: 1200px; margin: auto; padding-top: 20px; }
h1 { text-align: center; color: #1e3a8a; font-family: sans-serif; font-weight: 900; }
.btn-primary { background: linear-gradient(90deg, #1e3a8a 0%, #3b82f6 100%); color: white; border: none; font-size: 1.1rem; }
.btn-primary:hover { opacity: 0.9; }
"""

with gr.Blocks(css=css_style, title="GLM-OCR V1") as app:
    with gr.Column(elem_classes="container"):
        gr.Markdown("# 👁️ GLM-OCR ULTRA SCANNER")
        gr.Markdown("Ekstrak teks dari dokumen apapun. Upload, Klik, Jadi.")
        
        with gr.Row():
            # Kolom Input
            with gr.Column(scale=1):
                input_img = gr.Image(type="pil", label="Input Visual (Gambar)", sources=["upload", "clipboard"], height=450)
                scan_btn = gr.Button("⚡ AKTIFKAN PEMINDAI", variant="primary", elem_classes="btn-primary")
            
            # Kolom Output (Pake Textbox biasa biar anti error 'show_copy_button')
            with gr.Column(scale=1):
                output_txt = gr.Textbox(label="Data Terekstraksi", lines=22, show_copy_button=True, interactive=False)
    
    # Pemicu
    scan_btn.click(fn=proses_intelijen, inputs=input_img, outputs=output_txt)

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