import gradio as gr
import pandas as pd
from datetime import datetime
import json
import io
from PIL import Image
import pytesseract
from huggingface_hub import hf_hub_download # Thư viện mới
from llama_cpp import Llama # Thư viện mới

# --- CẤU HÌNH VÀ TẢI MÔ HÌNH (SỬ DỤNG LLAMA.CPP) ---
print("Ứng dụng đang khởi động...")

# OCR không thay đổi
print("Sử dụng Tesseract OCR (siêu nhẹ).")

# THAY ĐỔI LỚN: DÙNG LLAMA-CPP-PYTHON
print("Đang tải mô hình LLM (Llama-3-8B GGUF for CPU)...")
model_repo = "bartowski/Meta-Llama-3-8B-Instruct-GGUF"
model_file = "Meta-Llama-3-8B-Instruct-Q4_K_M.gguf"

# Tải file mô hình về cache của Space
model_path = hf_hub_download(repo_id=model_repo, filename=model_file)
print(f"Đã tải xong file mô hình tại: {model_path}")

# Khởi tạo mô hình từ file đã tải
llm = Llama(
  model_path=model_path,
  n_ctx=4096,     # Context length
  n_gpu_layers=0, # Chạy hoàn toàn trên CPU
  verbose=True,   # In ra thông tin để debug
)
print("Tải xong và khởi tạo thành công mô hình LLM.")


# --- CÁC HÀM XỬ LÝ ---

def run_ocr(image: Image.Image) -> str:
    # (Giữ nguyên)
    try:
        text = pytesseract.image_to_string(image)
        return text
    except Exception as e:
        print(f"Lỗi Tesseract: {e}")
        return "Lỗi khi đọc chữ từ ảnh."

def extract_order_from_text(text: str) -> dict:
    # Cập nhật prompt và cách gọi cho Llama.cpp
    prompt = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>
    You are an expert assistant that only outputs valid JSON. Extract order information from the text. The JSON object must contain "ten_khach_hang" (string, null if not found) and "danh_sach_hang" (an array of items). Each item must have "ten_hang" (string), "so_luong" (number), "don_vi" (string), "ma_hang" (string, null if not found), and "ghi_chu" (string, null if not found).<|eot_id|><|start_header_id|>user<|end_header_id|>
    Text Content:
    {text}<|eot_id|><|start_header_id|>assistant<|end_header_id|>
    """
    
    output = llm(
        prompt,
        max_tokens=1024,
        stop=["<|eot_id|>"],
        temperature=0.1,
        echo=False # Không in lại prompt trong kết quả
    )
    
    response_text = output['choices'][0]['text']
    
    try:
        json_str = response_text.strip()
        start = json_str.find('{')
        end = json_str.rfind('}') + 1
        if start != -1 and end != 0:
            json_str = json_str[start:end]
        return json.loads(json_str)
    except json.JSONDecodeError:
        return {"error": "AI trả về định dạng không hợp lệ", "raw_response": response_text}

def create_excel_file(order_data: dict):
    # (Giữ nguyên)
    if not order_data or "danh_sach_hang" not in order_data or not order_data["danh_sach_hang"]: return None
    flat_data = []
    customer = order_data.get('ten_khach_hang', 'N/A')
    for item in order_data['danh_sach_hang']:
        flat_data.append({
            'Khách hàng': customer, 'Mã hàng': item.get('ma_hang'),
            'Tên hàng': item.get('ten_hang'), 'Số lượng': item.get('so_luong'),
            'Đơn vị': item.get('don_vi'), 'Ghi chú': item.get('ghi_chu')
        })
    df = pd.DataFrame(flat_data)
    output = io.BytesIO()
    with pd.ExcelWriter(output, engine='openpyxl') as writer: df.to_excel(writer, index=False, sheet_name='DonHang')
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"don_hang_{timestamp}.xlsx"
    return (filename, output.getvalue())

def process_image_and_extract(image):
    # (Giữ nguyên)
    try:
        if image is None: return "Vui lòng dán ảnh vào.", None, None
        extracted_text = run_ocr(image)
        if not extracted_text.strip(): return "Không đọc được chữ từ hình ảnh.", None, None
        order_data = extract_order_from_text(extracted_text)
        if "error" in order_data: return extracted_text, f"Lỗi từ AI: {order_data['error']}\nPhản hồi gốc: {order_data['raw_response']}", None
        excel_info = create_excel_file(order_data)
        df_display = pd.DataFrame(order_data.get('danh_sach_hang', []))
        if excel_info:
            filename, filebytes = excel_info
            with open(filename, "wb") as f: f.write(filebytes)
            return extracted_text, df_display, filename
        else: return extracted_text, df_display, None
    except Exception as e:
        import traceback
        error_str = str(e)
        traceback_str = traceback.format_exc()
        print(traceback_str)
        return f"Lỗi nghiêm trọng: {error_str}", None, None

# --- XÂY DỰNG GIAO DIỆN GRADIO (GIỮ NGUYÊN) ---
with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("# Ứng dụng Trích xuất Đơn hàng từ Ảnh chụp màn hình")
    gr.Markdown("Chụp màn hình email/tin nhắn đặt hàng, sau đó dán (Ctrl+V) vào ô bên dưới và nhấn 'Xử lý'.")
    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(label="Dán ảnh chụp màn hình vào đây", type="pil", sources=["clipboard", "upload"])
            process_btn = gr.Button("Xử lý", variant="primary")
        with gr.Column(scale=2):
            gr.Markdown("### Kết quả trích xuất")
            output_table = gr.DataFrame(label="Chi tiết đơn hàng")
            output_excel = gr.File(label="Tải file Excel")
            gr.Markdown("### Văn bản đọc được từ ảnh (OCR)")
            output_text = gr.Textbox(label="Text from Image", lines=10, interactive=False)
    process_btn.click(fn=process_image_and_extract, inputs=image_input, outputs=[output_text, output_table, output_excel])

app.launch(debug=True)