import gradio as gr
import cv2
import numpy as np
from PIL import Image
import torch

# Import โค้ดสมองที่เราเขียนไว้
from scanner import scan_document
from ocr import extract_text

def process_image(input_image):
    if input_image is None:
        return None, "กรุณาอัปโหลดรูปภาพ"
    
    # 1. บันทึกรูปชั่วคราวเพื่อให้ scanner.py อ่านได้
    temp_path = "temp_upload.jpg"
    # input_image จาก Gradio เป็น numpy array (RGB) ต้องแปลงเป็น BGR ก่อน save ด้วย cv2
    input_bgr = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)
    cv2.imwrite(temp_path, input_bgr)
    
    try:
        # 2. เรียกใช้ Scanner (ดัดภาพ)
        # scan_document คืนค่า (ภาพสี, ภาพขาวดำ) เราเอาภาพสี
        warped_color, _ = scan_document(temp_path)
        
        # 3. เรียกใช้ OCR (อ่านข้อความ)
        # ส่งภาพที่ดัดแล้วไปให้ AI
        text_result = extract_text(warped_color)
        
        # แปลงภาพกลับเป็น RGB เพื่อแสดงผลใน Gradio
        output_img = cv2.cvtColor(warped_color, cv2.COLOR_BGR2RGB)
        
        return output_img, text_result

    except Exception as e:
        return None, f"เกิดข้อผิดพลาด: {str(e)}"

# สร้างหน้าเว็บด้วย Gradio
iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(label="อัปโหลดรูปใบเสร็จ/เอกสาร"),
    outputs=[
        gr.Image(label="ภาพที่สแกนแล้ว (Crop & Warp)"),
        gr.Textbox(label="ข้อความที่ AI อ่านได้ (JSON)", lines=10)
    ],
    title="🤖 AI Document Scanner & OCR",
    description="อัปโหลดรูปถ่ายเอกสาร -> ระบบจะดัดภาพให้ตรงอัตโนมัติ และอ่านข้อความด้วย Donut Transformer",
    theme="default"
)

# สั่งรัน
if __name__ == "__main__":
    iface.launch()