import gradio as gr
import cv2
import numpy as np
from gradio_client import Client
from io import BytesIO
from PIL import Image

# -----------------------------
# Hugging Face API Key
# -----------------------------
API_KEY = "HF_Token"  # <-- Replace with your Hugging Face token

# -----------------------------
# Local "try-on" simulation
# -----------------------------
def tryon_local(person_img, garment_img, seed, randomize_seed):
    if person_img is None or garment_img is None:
        return None, None, "Empty image"

    h_person, w_person = person_img.shape[:2]
    h_garment, w_garment = garment_img.shape[:2]

    scale = w_person / w_garment
    new_w = int(w_garment * scale)
    new_h = int(h_garment * scale)
    garment_resized = cv2.resize(garment_img, (new_w, new_h))

    overlay = person_img.copy()
    y_offset = 0
    x_offset = max(0, (w_person - new_w) // 2)

    if garment_resized.shape[2] == 4:
        alpha = garment_resized[:, :, 3] / 255.0
        for c in range(3):
            overlay[y_offset:y_offset+new_h, x_offset:x_offset+new_w, c] = \
                alpha * garment_resized[:, :, c] + (1 - alpha) * overlay[y_offset:y_offset+new_h, x_offset:x_offset+new_w, c]
    else:
        overlay[y_offset:y_offset+new_h, x_offset:x_offset+new_w] = garment_resized

    info = "Success (local simulation)"
    return overlay, seed, info

# -----------------------------
# Initialize Hunyuan3D API client
# -----------------------------
hunyuan_client = Client("tencent/Hunyuan3D-2.1", hf_token=API_KEY)

# -----------------------------
# Convert NumPy image to bytes for API
# -----------------------------
def np_to_bytes(img: np.ndarray):
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    pil_img = Image.fromarray(img_rgb)
    buf = BytesIO()
    pil_img.save(buf, format="PNG")
    buf.seek(0)
    return buf.read()

# -----------------------------
# Try-on using Hunyuan3D
# -----------------------------
def tryon_to_3d(person_img, garment_img, seed, randomize_seed):
    tryon_img, seed_used, info = tryon_local(person_img, garment_img, seed, randomize_seed)
    if tryon_img is None:
        return None, "Local try-on failed"

    try:
        person_bytes = np_to_bytes(person_img)
        garment_bytes = np_to_bytes(garment_img)

        # Use the correct endpoint without keyword args
        result_bytes = hunyuan_client.predict(person_bytes, garment_bytes)
        result_img = np.array(Image.open(BytesIO(result_bytes)))
        return result_img, info + " + Hunyuan3D processed"
    except Exception as e:
        return None, f"Hunyuan3D API error: {e}"

# -----------------------------
# Gradio UI
# -----------------------------
css = """
#col-left, #col-mid, #col-right {
    margin: 0 auto;
    max-width: 430px;
}
#col-showcase {
    margin: 0 auto;
    max-width: 1100px;
}
#button { color: blue; }
"""

with gr.Blocks(css=css) as app:
    with gr.Row():
        with gr.Column(elem_id="col-left"):
            person_input = gr.Image(label="Person Image", type="numpy")
        with gr.Column(elem_id="col-mid"):
            garment_input = gr.Image(label="Garment Image", type="numpy")
        with gr.Column(elem_id="col-right"):
            output_img = gr.Image(label="3D Result")
            result_info = gr.Text(label="Info")
            seed = gr.Slider(0, 999999, value=0, step=1, label="Seed")
            randomize_seed = gr.Checkbox(label="Random seed", value=True)
            run_btn = gr.Button("Run")

    run_btn.click(
        fn=tryon_to_3d,
        inputs=[person_input, garment_input, seed, randomize_seed],
        outputs=[output_img, result_info]
    )

app.launch()