import gradio as gr
import pandas as pd
import torch
from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
from PIL import Image
import io
from datetime import datetime

# โหลดโมเดลและ processor จาก Hugging Face
model_name = "Qwen/Qwen2-VL-7B-Instruct"
processor = AutoProcessor.from_pretrained(model_name)
model = Qwen2VLForConditionalGeneration.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

def extract_data_from_image(images):
    results = []

    for idx, img_file in enumerate(images):
        try:
            image = Image.open(io.BytesIO(img_file.read())).convert("RGB")

            # Prompt บอกโมเดลว่าให้ทำอะไร
            prompt = """
            กรุณาสกัดข้อมูลสำคัญจากเอกสารนี้:
            - วันที่
            - ยอดรวม
            - ชื่อร้านค้า
            - เลขใบเสร็จ
            
            กรุณาตอบในรูปแบบ JSON
            """

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

            text_prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
            inputs = processor(text=text_prompt, images=image, return_tensors="pt").to(model.device).bfloat16()

            with torch.no_grad():
                generated_ids = model.generate(**inputs, max_new_tokens=512)

            generated_ids_trimmed = [out_ids[len(inputs["input_ids"][0]):] for out_ids in generated_ids]
            answer = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]

            try:
                structured = eval(answer.replace("```json", "").replace("```", ""))
            except:
                structured = {"raw_response": answer}

            results.append({
                "file_name": img_file.name,
                "data": str(structured),
                "timestamp": datetime.now().isoformat()
            })

        except Exception as e:
            results.append({
                "file_name": img_file.name,
                "data": f"เกิดข้อผิดพลาด: {str(e)}",
                "timestamp": datetime.now().isoformat()
            })

    df = pd.DataFrame(results)
    df["structured_data"] = df["data"].astype(str)

    # บันทึกเป็น Parquet
    parquet_path = "output.parquet"
    df.to_parquet(parquet_path)

    return {
        "table": df[["file_name", "structured_data"]],
        "download": parquet_path
    }

# UI Components
title = "📄 ระบบสกัดข้อมูลเอกสารอัตโนมัติ (รองรับภาษาไทย)"
description = "อัปโหลดภาพหลายไฟล์ → สกัดข้อมูล → แยกหัวข้อ → บันทึกเป็น Parquet"

interface = gr.Interface(
    fn=extract_data_from_image,
    inputs=gr.File(type="file", file_types=["image"], multiple=True),
    outputs=[
        gr.Dataframe(label="ผลลัพธ์"),
        gr.File(label="ดาวน์โหลด Parquet")
    ],
    title=title,
    description=description,
    allow_flagging="never"
)

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