# import cv2
# import gradio as gr
# import numpy as np
# import os
# from PIL import Image
# from try_on_diffusion_client import TryOnDiffusionClient   # <‑ the file you just pulled

# # --- 1.  Initialise the API client -----------------------------------------
# API_URL = "https://try-on-diffusion.p.rapidapi.com"   # RapidAPI endpoint
# API_KEY = os.getenv("RAPIDAPI_KEY")                  # put your key in an env var
# client  = TryOnDiffusionClient(base_url=API_URL, api_key=API_KEY)

# # --- 2.  The Gradio callback ------------------------------------------------

# def try_on(user_img, outfit_img, height, chest, waist, sleeve):
#     # 1️⃣  Force 3‑channel RGB and uint8 dtype
#     user_rgb   = user_img.convert("RGB")
#     outfit_rgb = outfit_img.convert("RGB")

#     avatar_bgr   = cv2.cvtColor(np.array(user_rgb,   dtype=np.uint8), cv2.COLOR_RGB2BGR)
#     clothing_bgr = cv2.cvtColor(np.array(outfit_rgb, dtype=np.uint8), cv2.COLOR_RGB2BGR)

#     # 2️⃣  Call the API (unchanged)
#     resp = client.try_on_file(
#         clothing_image = clothing_bgr,
#         avatar_image   = avatar_bgr,
#         seed           = -1
#     )

#     if resp.status_code != 200 or resp.image is None:
#         return None, f"API error {resp.status_code}: {resp.error_details}"

#     result_rgb = cv2.cvtColor(resp.image, cv2.COLOR_BGR2RGB)

#     # 3️⃣  Simple fit message
#     fit_msg = "✅ Fit looks OK"
#     if chest < 80:
#         fit_msg = "⚠️ Chest measurement suggests a tight fit."

#     return Image.fromarray(result_rgb), fit_msg


# # --- 3. Create the Gradio Interface ------------------------------------------
# demo = gr.Interface(
#     fn=try_on,
#     inputs=[
#         gr.Image(type="pil", label="Upload your image"),
#         gr.Image(type="pil", label="Upload outfit image"),
#         gr.Number(label="Height (cm)"),
#         gr.Number(label="Chest (cm)"),
#         gr.Number(label="Waist (cm)"),
#         gr.Number(label="Sleeve length (cm)")
#     ],
#     outputs=[
#         gr.Image(type="pil", label="Try-On Result"),
#         gr.Text(label="Fit Advice")
#     ],
#     title="👕 Virtual Try-On",
#     description="Upload your image and an outfit to see how it might look on you!"
# )

# # --- 4. Launch the app -------------------------------------------------------
# if __name__ == "__main__":
#     print("Launching Virtual Try-On App...")
#     demo.launch()


import cv2
import gradio as gr
import numpy as np
import os
from PIL import Image
from try_on_diffusion_client import TryOnDiffusionClient

# --- 1. Initialize the API client -----------------------------------------
API_URL = "https://try-on-diffusion.p.rapidapi.com"
API_KEY = os.getenv("RAPIDAPI_KEY")  # This will be set as a secret in HF Spaces

if not API_KEY:
    print("⚠️ RAPIDAPI_KEY not found. Please set it in Hugging Face Spaces secrets.")

client = TryOnDiffusionClient(base_url=API_URL, api_key=API_KEY)

# --- 2. The Gradio callback ------------------------------------------------
def try_on(user_img, outfit_img, height, chest, waist, sleeve):
    """
    Process the virtual try-on request
    """
    # Check if API key is available
    if not API_KEY:
        return None, "❌ API key not configured. Please contact the app administrator."
    
    # Validate inputs
    if user_img is None or outfit_img is None:
        return None, "❌ Please upload both user image and outfit image."
    
    try:
        # Convert images to RGB and ensure uint8 dtype
        user_rgb = user_img.convert("RGB")
        outfit_rgb = outfit_img.convert("RGB")

        avatar_bgr = cv2.cvtColor(np.array(user_rgb, dtype=np.uint8), cv2.COLOR_RGB2BGR)
        clothing_bgr = cv2.cvtColor(np.array(outfit_rgb, dtype=np.uint8), cv2.COLOR_RGB2BGR)

        # Call the API
        resp = client.try_on_file(
            clothing_image=clothing_bgr,
            avatar_image=avatar_bgr,
            seed=-1
        )

        if resp.status_code != 200 or resp.image is None:
            error_msg = resp.error_details if resp.error_details else f"API error {resp.status_code}"
            return None, f"❌ {error_msg}"

        # Convert result back to RGB
        result_rgb = cv2.cvtColor(resp.image, cv2.COLOR_BGR2RGB)

        # Generate fit advice
        fit_msg = generate_fit_advice(height, chest, waist, sleeve)

        return Image.fromarray(result_rgb), fit_msg

    except Exception as e:
        return None, f"❌ Error processing request: {str(e)}"


def generate_fit_advice(height, chest, waist, sleeve):
    """
    Generate basic fit advice based on measurements
    """
    advice = ["✅ Fit Analysis:"]
    
    if chest and chest < 80:
        advice.append("⚠️ Chest measurement suggests the item might be tight.")
    elif chest and chest > 120:
        advice.append("ℹ️ Chest measurement suggests a loose fit.")
    else:
        advice.append("✅ Chest measurement looks good.")
    
    if waist and chest and abs(chest - waist) < 10:
        advice.append("ℹ️ Similar chest and waist measurements - consider fitted styles.")
    
    if height and height < 160:
        advice.append("ℹ️ For shorter stature, consider checking garment length.")
    elif height and height > 185:
        advice.append("ℹ️ For taller stature, ensure adequate garment length.")
    
    return "\n".join(advice)


# --- 3. Create the Gradio Interface ------------------------------------------
with gr.Blocks(theme=gr.themes.Soft(), title="Virtual Try-On") as demo:
    gr.Markdown(
        """
        # 👕 Virtual Try-On App
        
        Upload your photo and a clothing item to see how it would look on you!
        
        **Tips for best results:**
        - Use clear, well-lit photos
        - Person should be facing forward
        - Clothing item should be clearly visible
        """
    )
    
    with gr.Row():
        with gr.Column():
            user_img = gr.Image(type="pil", label="Your Photo", height=300)
            outfit_img = gr.Image(type="pil", label="Clothing Item", height=300)
            
        with gr.Column():
            gr.Markdown("### 📏 Your Measurements (Optional)")
            height = gr.Number(label="Height (cm)", value=170, minimum=140, maximum=220)
            chest = gr.Number(label="Chest (cm)", value=90, minimum=70, maximum=140)
            waist = gr.Number(label="Waist (cm)", value=80, minimum=60, maximum=130)
            sleeve = gr.Number(label="Sleeve Length (cm)", value=60, minimum=50, maximum=80)
    
    submit_btn = gr.Button("✨ Try On Outfit", variant="primary", size="lg")
    
    with gr.Row():
        result_img = gr.Image(type="pil", label="Try-On Result", height=400)
        fit_advice = gr.Textbox(label="Fit Advice", lines=6, max_lines=10)
    
    # Examples section
    gr.Markdown("### 📸 Example Images")
    gr.Markdown("*Upload your own images or try with sample photos*")
    
    submit_btn.click(
        fn=try_on,
        inputs=[user_img, outfit_img, height, chest, waist, sleeve],
        outputs=[result_img, fit_advice]
    )

# --- 4. Launch the app -------------------------------------------------------
if __name__ == "__main__":
    print("🚀 Launching Virtual Try-On App...")
    demo.launch()