import gradio as gr
from PIL import Image
import os

def load_hairstyles():
    folder = "hairstyles"
    if not os.path.exists(folder):
        return []
    return [
        Image.open(os.path.join(folder, f)).convert("RGBA")
        for f in sorted(os.listdir(folder)) if f.endswith(".png")
    ]

hairstyles = load_hairstyles()

def apply_hairstyle(user_img, style_index, x_offset, y_offset, scale):
    if user_img is None or not hairstyles:
        return None

    user_img = user_img.convert("RGBA")
    base_w, base_h = user_img.size

    hairstyle = hairstyles[style_index]

    # Resize the hairstyle based on scale
    new_size = (int(base_w * scale), int(hairstyle.height * (base_w * scale / hairstyle.width)))
    hairstyle = hairstyle.resize(new_size)

    # Create a blank transparent image to position the hairstyle
    composite = Image.new("RGBA", user_img.size)
    paste_x = int((base_w - new_size[0]) / 2 + x_offset)
    paste_y = int(y_offset)
    composite.paste(hairstyle, (paste_x, paste_y), hairstyle)

    # Overlay it
    result = Image.alpha_composite(user_img, composite)
    return result.convert("RGB")

with gr.Blocks() as demo:
    gr.Markdown("## 💇 Salon Virtual Hairstyle Try-On (Adjustable)")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="📷 Upload an Image")
            style_slider = gr.Slider(0, max(len(hairstyles)-1, 0), step=1, label="🎨 Select Hairstyle")
            x_offset = gr.Slider(-200, 200, value=0, step=1, label="⬅️➡️ Move Left / Right")
            y_offset = gr.Slider(-200, 200, value=0, step=1, label="⬆️⬇️ Move Up / Down")
            scale = gr.Slider(0.3, 2.0, value=1.0, step=0.05, label="📏 Scale Hairstyle")
            apply_btn = gr.Button("✨ Apply Hairstyle")
        with gr.Column():
            result_output = gr.Image(label="🔍 Result Preview")

    apply_btn.click(
        fn=apply_hairstyle,
        inputs=[image_input, style_slider, x_offset, y_offset, scale],
        outputs=result_output
    )

demo.launch()