import gradio as gr
from overlay import overlay_source
from detect_face import predict, NUM_CLASSES
from swapface import swap_face_now
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

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.splitext(os.path.basename(filepath))[0] 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 "", []

# --- Hàm swap face ---
def swap_face_wrapper(background_img, result_img):
    """
    Wrapper function cho swap face giữa background và result image
    """
    if background_img is None or result_img is None:
        return None
    
    try:
        # Swap face từ background vào result image
        swapped = swap_face_now(background_img, result_img, do_enhance=True)
        return swapped
    except Exception as e:
        print(f"Lỗi swap face: {e}")
        return result_img  # Trả về ảnh gốc nếu có lỗi

# --- Hàm gộp overlay + swap face ---
def combined_hair_and_face(background_img, source_img):
    """
    Hàm gộp: chạy overlay trước, sau đó swap face
    """
    if background_img is None or source_img is None:
        return None
    
    try:
        # Bước 1: Chạy overlay (ghép tóc)
        hair_result = overlay_source(background_img, source_img)
        
        # Bước 2: Chạy swap face từ background lên kết quả overlay
        final_result = swap_face_wrapper(background_img, hair_result)
        
        return final_result
    except Exception as e:
        print(f"Lỗi trong quá trình gộp hair + face: {e}")
        return None

# --- Xây dựng giao diện Gradio ---
def build_demo():
    with gr.Blocks(title="Hair Try-On & Face Swap", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🎯 Hair Try-On & Face Swap Application
        """)
        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
                            )
            combined_btn = gr.Button("✨ Try on now", variant="primary")

        # Chạy gộp hair + face swap
        combined_btn.click(fn=combined_hair_and_face, 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
        )
        # 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()