import os
import cv2
import gradio as gr
import numpy as np
import random
import base64
import time
from gradio_client import Client

# -----------------------------
# Connect to Hunyuan3D Space
# -----------------------------
mesh_client = Client("tencent/Hunyuan3D-2.1")  # has real API
MESH_API_NAME = "/process"

# -----------------------------
# Kolors VTON code (internal)
# -----------------------------
MAX_SEED = 999999

def tryon(person_img, garment_img, seed, randomize_seed):
    if person_img is None or garment_img is None:
        return None, seed, "Empty image"
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    # Encode images
    encoded_person_img = cv2.imencode('.jpg', cv2.cvtColor(person_img, cv2.COLOR_RGB2BGR))[1].tobytes()
    encoded_person_img = base64.b64encode(encoded_person_img).decode('utf-8')
    encoded_garment_img = cv2.imencode('.jpg', cv2.cvtColor(garment_img, cv2.COLOR_RGB2BGR))[1].tobytes()
    encoded_garment_img = base64.b64encode(encoded_garment_img).decode('utf-8')

    # Send request to internal tryon server (Kolors VTON)
    tryon_url = "http://" + os.environ['tryon_url'] + "Submit"
    token = os.environ['token']
    cookie = os.environ['Cookie']
    referer = os.environ['referer']
    headers = {'Content-Type': 'application/json', 'token': token, 'Cookie': cookie, 'referer': referer}
    data = {"clothImage": encoded_garment_img, "humanImage": encoded_person_img, "seed": seed}

    # Post request
    import requests
    try:
        response = requests.post(tryon_url, headers=headers, data=json.dumps(data), timeout=50)
        if response.status_code != 200:
            return None, seed, f"Tryon POST failed: {response.status_code}"
        result_json = response.json()['result']
        if result_json['status'] != "success":
            return None, seed, f"Tryon error: {result_json['status']}"
        task_uuid = result_json['result']
    except Exception as e:
        return None, seed, f"Tryon exception: {str(e)}"

    # Poll for result
    result_img = None
    for _ in range(12):  # retry 12 times
        try:
            query_url = "http://" + os.environ['tryon_url'] + f"Query?taskId={task_uuid}"
            response = requests.get(query_url, headers=headers, timeout=20)
            if response.status_code == 200:
                r = response.json()['result']
                if r['status'] == "success":
                    img_bytes = base64.b64decode(r['result'])
                    result_np = np.frombuffer(img_bytes, np.uint8)
                    result_img = cv2.imdecode(result_np, cv2.IMREAD_COLOR)
                    result_img = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
                    break
            time.sleep(1)
        except:
            time.sleep(1)

    if result_img is None:
        return None, seed, "Tryon failed or timed out"

    return result_img, seed, "Success"

# -----------------------------
# Merge: Tryon + Hunyuan3D
# -----------------------------
def tryon_to_3d(person_img, garment_img, seed, randomize_seed):
    # Run 2D try-on
    tryon_img, seed_used, tryon_info = tryon(person_img, garment_img, seed, randomize_seed)
    if tryon_img is None:
        return None, None, seed_used, tryon_info

    # Run 3D mesh generation (Hunyuan3D)
    try:
        mesh_file = mesh_client.predict(tryon_img, api_name=MESH_API_NAME)
    except Exception as e:
        return tryon_img, None, seed_used, f"Hunyuan3D error: {str(e)}"

    return tryon_img, mesh_file, seed_used, "Success"

# -----------------------------
# Gradio Interface
# -----------------------------
with gr.Blocks() as demo:
    gr.Markdown("# Virtual Try-On → 3D Mesh Generator")

    with gr.Row():
        person_img = gr.Image(label="Person Image", type="numpy")
        garment_img = gr.Image(label="Garment Image", type="numpy")

    with gr.Row():
        seed = gr.Slider(0, MAX_SEED, step=1, value=0, label="Seed")
        randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)

    with gr.Row():
        tryon_out = gr.Image(label="2D Try-On Result")
        mesh_out = gr.Model3D(label="3D Mesh Output")

    seed_used_out = gr.Number(label="Seed Used")
    status_out = gr.Textbox(label="Status")

    run_btn = gr.Button("Run")
    run_btn.click(fn=tryon_to_3d,
                  inputs=[person_img, garment_img, seed, randomize_seed],
                  outputs=[tryon_out, mesh_out, seed_used_out, status_out])

if __name__ == "__main__":
    demo.launch()