import gradio as gr
from overlay import overlay_source
from detect_face import predict, NUM_CLASSES
import os
from pathlib import Path

BASE_DIR = Path(__file__).parent  # thư mục chứa app.py
FOLDER = BASE_DIR / "example_wigs"

# --- Hàm load ảnh từ folder ---
def load_images_from_folder(folder_path: str) -> list[str]:
    """
    Trả về list[str] chứa tất cả các hình (jpg, png, gif, bmp) trong folder_path.
    """
    supported = {'.jpg', '.jpeg', '.png', '.gif', '.bmp'}
    if not os.path.isdir(folder_path):
        print(f"Cảnh báo: '{folder_path}' không phải folder hợp lệ.")
        return []
    files = [
        os.path.join(folder_path, fn)
        for fn in os.listdir(folder_path)
        if os.path.splitext(fn)[1].lower() in supported
    ]
    if not files:
        print(f"Không tìm thấy hình trong: {folder_path}")
    return files

# --- Handler khi click thumbnail của Gallery ---
# def on_gallery_select(evt: gr.SelectData):
#     """
#     Xử lý khi click vào ảnh trong gallery - tối ưu và robust.
#     """
#     val = evt.value
#     if isinstance(val, dict):
#         img = val.get("image")
#         if isinstance(img, str): return img
#         if isinstance(img, dict):
#             path = img.get("path") or img.get("url")
#             if isinstance(path, str): return path
#             for v in img.values():
#                 if isinstance(v, str) and os.path.isfile(v):
#                     return v
#         for v in val.values():
#             if isinstance(v, str) and os.path.isfile(v):
#                 return v
#         raise ValueError(f"Không trích được filepath từ dict: {val}")
#     if isinstance(val, str):
#         return val
#     raise ValueError(f"Kiểu không hỗ trợ: {type(val)}")

def on_gallery_select(evt: gr.SelectData):
    """
    Khi click thumbnail: trả về
      1) filepath để nạp vào Image Source
      2) tên file (basename) để hiển thị trong Textbox
    """
    val = evt.value

    # --- logic trích filepath y như cũ ---
    if isinstance(val, dict):
        img = val.get("image")
        if isinstance(img, str):
            filepath = img
        elif isinstance(img, dict):
            filepath = img.get("path") or img.get("url")
        else:
            filepath = next(
                (v for v in val.values() if isinstance(v, str) and os.path.isfile(v)),
                None
            )
    elif isinstance(val, str):
        filepath = val
    else:
        raise ValueError(f"Kiểu không hỗ trợ: {type(val)}")

    filename = os.path.basename(filepath) if filepath else ""
    return filepath, filename

# --- Hàm xác định folder dựa trên phân lớp ---
def infer_folder(image) -> str:
    cls = predict(image)["predicted_class"]
    folder = str(FOLDER / cls)
    return folder

# --- Hàm gộp: phân loại + load ảnh ---
def handle_bg_change(image):
    """
    Khi thay đổi background:
    1. Phân loại khuôn mặt
    2. Load ảnh từ folder tương ứng
    """
    if image is None:
        return "", []
    
    try:
        folder = infer_folder(image)
        images = load_images_from_folder(folder)
        return folder, images
    except Exception as e:
        print(f"Lỗi xử lý ảnh: {e}")
        return "", []

# --- Xây dựng giao diện Gradio ---
def build_demo():
    with gr.Blocks(title="Xử lý hai hình ảnh", theme=gr.themes.Soft()) as demo:
        gr.Markdown("Upload Background & Source, click **Run** to try on wigs.")

        with gr.Row():
            bg = gr.Image(type="pil", label="Background", height=500)
            src = gr.Image(type="pil", label="Source", height=500, interactive=False)
            out = gr.Image(label="Result", height=500, interactive=False)

        folder_path_box = gr.Textbox(label="Folder path", visible=False)

        
        with gr.Row():
            src_name_box = gr.Textbox(
                        label="Wigs Name",
                        interactive=False,
                        show_copy_button=True ,     # tuỳ chọn – tiện copy đường dẫn
                        scale = 1
                        )
            gallery = gr.Gallery(
                            label="Recommend For You",
                            height=300,
                            value=[],
                            type="filepath",
                            interactive=False,
                            columns=5,
                            object_fit="cover",
                            allow_preview=True,
                            scale = 8
                            )
            btn = gr.Button("🔄 Run", variant="primary",scale = 1)  
 


        # Chạy ghép tóc
        btn.click(fn=overlay_source, inputs=[bg, src], outputs=[out])
        # Khi đổi ảnh background, tự động phân loại và load ảnh gợi ý
        bg.change(
            fn=handle_bg_change,
            inputs=[bg],
            outputs=[folder_path_box, gallery],
            show_progress=True
        )
        # Nút tải lại ảnh thủ công (backup)
        # Khi chọn ảnh trong gallery, cập nhật vào khung Source
        gallery.select(
            fn=on_gallery_select,
            outputs=[src, src_name_box]
        )

    return demo

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