import gradio as gr
from pathlib import Path
from ui import build_ui
from overlay import apply_hairstyle
from segmentation import segment_image, estimate_landmarks

# Folder containing your PNG hairstyle overlays
ASSETS_DIR = Path("assets/hairstyles")

def get_hairstyle_list():
    """Return sorted list of PNG files in the assets/hairstyles folder."""
    if not ASSETS_DIR.exists():
        print("⚠️ hairstyles folder not found:", ASSETS_DIR)
        return []
    styles = sorted([f.name for f in ASSETS_DIR.glob("*.png") if f.is_file()])
    print("✅ Found hairstyles:", styles)
    return styles

def try_on(user_image, style_name):
    """Main try-on function: segments hair, aligns style, blends result."""
    if user_image is None:
        return None
    if not style_name:
        print("⚠️ No hairstyle selected")
        return user_image
    mask = segment_image(user_image)
    landmarks = estimate_landmarks(user_image)  # optional
    style_path = str(ASSETS_DIR / style_name)
    return apply_hairstyle(user_image, style_path, mask, landmarks)

def launch():
    styles = get_hairstyle_list()
    if not styles:
        print("⚠️ No PNG hairstyles found in", ASSETS_DIR)
    demo = build_ui(try_on, styles)
    demo.launch()

if __name__ == "__main__":
    launch()