import gradio as gr
import os
import uuid
import subprocess
from PIL import Image

def restore_image(input_image):
    # Tạo thư mục tạm
    temp_id = str(uuid.uuid4())
    input_folder = f"./temp_input/{temp_id}"
    output_folder = f"./temp_output/{temp_id}/final_output"
    os.makedirs(input_folder, exist_ok=True)

    # Lưu ảnh tạm thời
    input_path = os.path.join(input_folder, "input.png")
    input_image.save(input_path)

    # Gọi run.py để phục hồi ảnh
    command_restore = f"python run.py --input_folder {input_folder} --output_folder ./temp_output/{temp_id} --GPU -1 --HR"
    os.system(command_restore)

    # Kiểm tra ảnh đầu ra
    if not os.path.exists(output_folder):
        return None, "Không tìm thấy thư mục đầu ra."

    output_files = [f for f in os.listdir(output_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    if not output_files:
        return None, "Không tìm thấy ảnh đầu ra."

    output_path = os.path.join(output_folder, output_files[0])
    output_image = Image.open(os.path.join(output_folder, output_files[0]))

    # Gọi main.py để sinh mô tả
    try:
        result = subprocess.run(
            ["python", "main.py", output_path],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        description = result.stdout.strip()
        if not description:
            description = "Không có mô tả từ AI."
    except Exception as e:
        description = f"Lỗi khi mô tả ảnh: {e}"

    return output_image, description

# Tạo Gradio app
demo = gr.Interface(
    fn=restore_image,
    inputs=gr.Image(type="pil", label="Upload ảnh cũ bị hư"),
    outputs=[
        gr.Image(label="Ảnh sau phục hồi"),
        gr.Textbox(label="Mô tả triết học")
    ],
    title="Phục hồi ảnh & mô tả triết học",
    description="Upload ảnh cũ để phục hồi và AI sẽ phân tích nội dung triết học trong ảnh."
)
demo.launch()