import gradio as gr
from gradio_client import Client, handle_file
from PIL import Image
import tempfile, os, time, random

# ══════════════════════════════════════════════════════════════
# More VTON spaces to try — updated 2026 list
# ══════════════════════════════════════════════════════════════
SPACES = [
    ("yisol/IDM-VTON",           "idmvton"),
    ("yisol/IDM-VTON-DC",        "idmvton-dc"),
    ("Kwai-Kolors/Kolors-Virtual-Try-On", "kolors"),
    ("Nymbo/Virtual-Try-On",     "nymbo"),
    ("multimodalart/idm-vton",   "multimodal"),
]

def save_tmp(img):
    p = tempfile.mktemp(suffix=".jpg")
    img.save(p, quality=95)
    return p

def cleanup(*paths):
    for p in paths:
        try:
            if p and os.path.exists(p): os.remove(p)
        except: pass

def try_kolors(client, p_path, c_path, steps, seed):
    """Kwai Kolors has different API signature."""
    result = client.predict(
        human_img   = handle_file(p_path),
        garm_img    = handle_file(c_path),
        garment_des = "clothing item",
        api_name    = "/tryon"
    )
    return result[0]

def try_idmvton(client, p_path, c_path, steps, seed, is_dress):
    """Standard IDM-VTON API signature."""
    result = client.predict(
        dict            = {"background": handle_file(p_path), "layers": [], "composite": None},
        garm_img        = handle_file(c_path),
        garment_des     = "dress" if is_dress else "shirt",
        is_checked      = True,
        is_checked_crop = is_dress,
        denoise_steps   = steps,
        seed            = seed,
        api_name        = "/tryon"
    )
    return result[0]

def virtual_tryon(person_img, cloth_img, quality):
    if person_img is None:
        return None, "⚠️ Please upload your photo!"
    if cloth_img is None:
        return None, "⚠️ Please upload a clothing image!"

    steps    = {"Fast ⚡": 20, "Balanced ⚖️": 30, "Best Quality 💎": 40}.get(quality, 30)
    seed     = random.randint(0, 9999)
    t0       = time.time()
    is_dress = (cloth_img.height / max(cloth_img.width, 1)) > 1.3

    p_path = save_tmp(person_img)
    c_path = save_tmp(cloth_img)

    errors = []
    for space_id, space_type in SPACES:
        try:
            print(f"🔌 Trying {space_id}...")
            client = Client(space_id, verbose=False)

            if space_type == "kolors":
                result_path = try_kolors(client, p_path, c_path, steps, seed)
            else:
                result_path = try_idmvton(client, p_path, c_path, steps, seed, is_dress)

            result_img = Image.open(result_path).convert("RGB")
            cleanup(p_path, c_path)
            elapsed = time.time() - t0
            name = space_id.split("/")[-1]
            return result_img, f"✅ Done in {elapsed:.1f}s via {name}!"

        except Exception as e:
            err = str(e)
            errors.append(f"{space_id}: {err[:60]}")
            print(f"❌ {space_id} → {err[:80]}")
            time.sleep(2)
            continue

    cleanup(p_path, c_path)

    # Show which spaces were tried
    print("All spaces failed:\n" + "\n".join(errors))
    return None, (
        "⏳ All servers are currently busy.\n\n"
        "Please try again in 2-3 minutes.\n"
        "💡 Best time to use: early morning or late night (India time)"
    )


# ── UI ────────────────────────────────────────────────────────
with gr.Blocks(
    title="AI Virtual Try-On",
    theme=gr.themes.Soft(primary_hue="violet", secondary_hue="pink"),
    css="""
    footer { display:none !important }
    .tip { background:#f5f0ff; border-radius:12px;
           padding:10px 14px; font-size:.85em;
           color:#555; margin-top:8px;
           border:1px solid #e0d0ff }
    """
) as demo:

    gr.HTML("""
    <div style='text-align:center;padding:20px 0 10px'>
      <h1 style='font-size:2.3em;font-weight:900;
        background:linear-gradient(90deg,#7c3aed,#db2777);
        -webkit-background-clip:text;-webkit-text-fill-color:transparent'>
        AI Virtual Try-On
      </h1>
      <p style='color:#888;margin-top:4px'>
        Upload your photo + clothing — see how it looks on you instantly!
      </p>
    </div>
    """)

    with gr.Row(equal_height=True):
        with gr.Column():
            gr.Markdown("### Step 1 — Your Photo")
            p_in = gr.Image(label="Upload full-body photo", type="pil", height=380)
            gr.HTML('<div class="tip">Stand straight · Full body · Good lighting · Plain background</div>')

        with gr.Column():
            gr.Markdown("### Step 2 — Clothing Item")
            c_in = gr.Image(label="Upload clothing product photo", type="pil", height=380)
            gr.HTML('<div class="tip">White background product photo · AI detects garment type automatically</div>')

        with gr.Column():
            gr.Markdown("### Step 3 — Result")
            r_out  = gr.Image(label="Try-On Result", type="pil", height=380)
            status = gr.Textbox(
                label="Status",
                interactive=False,
                placeholder="Your result will appear here after clicking Generate…"
            )

    with gr.Row():
        quality = gr.Radio(
            ["Fast ⚡", "Balanced ⚖️", "Best Quality 💎"],
            value="Balanced ⚖️",
            label="Quality",
            info="Fast = quicker · Best Quality = slower but sharper",
            scale=2
        )
        gr.Button(
            "Generate Try-On!", variant="primary", size="lg", scale=1
        ).click(
            fn=virtual_tryon,
            inputs=[p_in, c_in, quality],
            outputs=[r_out, status]
        )

    gr.HTML("""
    <div style='text-align:center;margin-top:16px;padding:12px 16px;
      background:linear-gradient(135deg,#f5f0ff,#fff0f8);
      border-radius:12px;color:#666;font-size:.83em;
      border:1px solid #e8d8ff'>
      Tries 5 different AI servers automatically ·
      Auto-detects dress vs top · 100% Free
    </div>
    """)

demo.launch()