import json, subprocess, tempfile, uuid, os, shutil
import numpy as np, torch, spaces, gradio as gr
from PIL import Image
from diffusers import QwenImageEditPipeline

# ---------- GPU models: load at module level, ZeroGPU moves them on demand ----------
edit_pipe = QwenImageEditPipeline.from_pretrained(
    "Qwen/Qwen-Image-Edit-2509", torch_dtype=torch.bfloat16).to("cuda")
# 4-8 step Lightning LoRA keeps the call inside the time cap (check the repo for the exact weight file)
edit_pipe.load_lora_weights("lightx2v/Qwen-Image-Lightning", weight_name="<edit lightning .safetensors>")

from hy3dgen.shapegen import Hunyuan3DDiTFlowMatchingPipeline  # or your TRELLIS wrapper
shape_pipe = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained("tencent/Hunyuan3D-2mini").to("cuda")

# ---------- CPU pose model ----------
from rtmlib import Body
pose_model = Body(mode="balanced", backend="onnxruntime", device="cpu")  # COCO-17 keypoints

PROMPT = ("Same character, same outfit and colors, full body, T-pose with arms straight out "
          "to the sides, facing the camera, plain white background")

@spaces.GPU(duration=90)
def make_tpose(img):
    return edit_pipe(image=img, prompt=PROMPT, num_inference_steps=8,
                     true_cfg_scale=1.0).images[0]

@spaces.GPU(duration=120)
def make_mesh(tpose_img, job):
    mesh = shape_pipe(image=tpose_img)[0]          # trimesh
    path = f"{job}/raw.glb"; mesh.export(path)
    return path

def keypoints(img):
    kp, sc = pose_model(np.array(img.convert("RGB"))[:, :, ::-1])
    return kp[0].tolist(), sc[0].tolist()          # first person only

def fg_bbox(img, thr=240):
    a = np.array(img.convert("L")); ys, xs = np.where(a < thr)
    return [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())]

def blender(job, cmd, **kw):
    kw.update(cmd=cmd, job=job)
    subprocess.run(["python", "blender_worker.py", json.dumps(kw)], check=True)

# ---------- Gradio event handlers ----------
def step_generate(img):
    job = tempfile.mkdtemp(); img = img.convert("RGB")
    tpose = make_tpose(img); tpose.save(f"{job}/tpose.png"); img.save(f"{job}/orig.png")
    raw = make_mesh(tpose, job)
    kp_t, sc_t = keypoints(tpose)
    json.dump(dict(kp=kp_t, sc=sc_t, bbox=fg_bbox(tpose), size=tpose.size), open(f"{job}/tpose_kp.json", "w"))
    blender(job, "rig")                              # auto-rig from T-pose joints
    return job, tpose, f"{job}/rigged.glb"

def step_pose(job):
    orig = Image.open(f"{job}/orig.png")
    kp, sc = keypoints(orig)
    json.dump(dict(kp=kp, sc=sc, size=orig.size), open(f"{job}/target_kp.json", "w"))
    blender(job, "pose")                             # writes posed.glb/.obj/.fbx
    return f"{job}/posed.glb", f"{job}/posed.glb", f"{job}/posed.obj", f"{job}/posed.fbx"

def step_uv(job):
    blender(job, "uv"); return f"{job}/uv_layout.png", f"{job}/texture.png"

def step_texture(job, tex):
    shutil.copy(tex, f"{job}/new_texture.png"); blender(job, "retexture")
    return f"{job}/posed.glb", f"{job}/posed.glb", f"{job}/posed.obj", f"{job}/posed.fbx"

with gr.Blocks() as demo:
    job = gr.State()
    with gr.Row():
        inp = gr.Image(type="pil", label="Upload")
        tp = gr.Image(label="T-pose")
        view = gr.Model3D(label="3D")
    go = gr.Button("1. Generate T-pose + 3D + rig")
    pose = gr.Button("2. Pose like the uploaded image")
    with gr.Row():
        d_glb, d_obj, d_fbx = (gr.DownloadButton(f"Download {x}") for x in ("GLB", "OBJ", "FBX"))
    with gr.Accordion("UV / texture", open=False):
        uv_btn = gr.Button("Export UV layout + texture")
        uv_img, tex_img = gr.Image(label="UV layout"), gr.Image(label="Texture")
        new_tex = gr.File(label="Upload new texture (PNG, same UV layout)", type="filepath")
        apply_btn = gr.Button("Apply texture")
    go.click(step_generate, inp, [job, tp, view])
    pose.click(step_pose, job, [view, d_glb, d_obj, d_fbx])
    uv_btn.click(step_uv, job, [uv_img, tex_img])
    apply_btn.click(step_texture, [job, new_tex], [view, d_glb, d_obj, d_fbx])
demo.queue().launch()