import os
import json
import torch
import sys
from PIL import Image
from diffsynth.pipelines.z_image import ZImagePipeline, ModelConfig, ControlNetInput
from multiprocessing import Process, set_start_method

try:
    from controlnet_aux.open_pose import (
        draw_poses, BodyResult, PoseResult, Keypoint,
    )
except ImportError:
    from controlnet_aux.openpose import (
        draw_poses, BodyResult, PoseResult, Keypoint,
    )

def xy_to_kp(xy):
    if xy is None:
        return None
    return Keypoint(float(xy[0]), float(xy[1]), 1.0, -1)

def xy_list_to_kp_list(lst):
    if lst is None:
        return None
    return [xy_to_kp(xy) for xy in lst]

def dict_to_pose(d):
    body_kps = xy_list_to_kp_list(d.get("body"))
    if body_kps is None:
        body = None
    else:
        body = BodyResult(
            keypoints=body_kps,
            total_score=0.0,
            total_parts=0,
        )
    return PoseResult(
        body=body,
        left_hand=xy_list_to_kp_list(d.get("left_hand")),
        right_hand=xy_list_to_kp_list(d.get("right_hand")),
        face=xy_list_to_kp_list(d.get("face")),
    )

def process_partition(gpu_id, subset_tasks, lora_path, base_output_dir):
    device = f"cuda:{gpu_id}"
    print(f"[GPU {gpu_id}] Loading controlnet model on {device}...")
    
    pipe = ZImagePipeline.from_pretrained(
        torch_dtype=torch.bfloat16,
        device=device,
        model_configs=[
            ModelConfig(model_id="PAI/Z-Image-Turbo-Fun-Controlnet-Union-2.1", origin_file_pattern="Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps.safetensors"),
            ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="transformer/*.safetensors"),
            ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="text_encoder/*.safetensors"),
            ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
        ],
        tokenizer_config=ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="tokenizer/"),
    )
    
    print(f"[GPU {gpu_id}] Loading LoRA: {lora_path}")
    if os.path.exists(lora_path):
        pipe.load_lora(pipe.dit, lora_config=lora_path, alpha=1.0)
    else:
        print(f"[GPU {gpu_id}] Warning: LoRA path {lora_path} does not exist.")
    
    for task in subset_tasks:
        fname = task["fname"]
        info = task["info"]
        
        try:
            content_desc = info.get("content_description", "")
            comp_analysis = info.get("composition_analysis", "")
            prompt = f"{content_desc} {comp_analysis}".strip()
            
            control_pose = info.get("control_pose", {})
            poses = [dict_to_pose(p) for p in control_pose.get("poses", [])]
            canvas_h = int(control_pose.get("canvas_h", 1024))
            canvas_w = int(control_pose.get("canvas_w", 1024))
            
            # 使用脚本 2 的渲染逻辑
            canvas = draw_poses(
                poses, canvas_h, canvas_w,
                draw_body=True,
                draw_hand=True,
                draw_face=True,
            )
            controlnet_img = Image.fromarray(canvas)
            
            h, w = canvas_h, canvas_w
            
            # 确保宽高是 16 的倍数
            h = (h // 16) * 16
            w = (w // 16) * 16
            controlnet_img = controlnet_img.resize((w, h))
            
            gen_img = pipe(
                prompt=prompt, 
                seed=0, 
                height=h, 
                width=w, 
                num_inference_steps=40,
                controlnet_inputs=[ControlNetInput(image=controlnet_img, scale=0.7)]
            )
            
            out_path = os.path.join(base_output_dir, fname)
            gen_img.save(out_path)
            print(f"[GPU {gpu_id}] Saved {fname}")
            
        except Exception as e:
            print(f"[GPU {gpu_id}] Error processing {fname}: {e}")

def get_available_gpus():
    if torch.cuda.is_available():
        return list(range(torch.cuda.device_count()))
    return [0]

def main():
    try:
        set_start_method('spawn')
    except RuntimeError:
        pass
        
    gpus = get_available_gpus()
    print(f"Available GPUs: {gpus}")
    
    lora_path = "models/train/AI4VA-Pose-Controlnet-LoRA-DPO-8x1-3-0515/step-200.safetensors"
    base_output_dir = "output"
    input_json = "inference/track_2_test.json"
    
    os.makedirs(base_output_dir, exist_ok=True)
    
    with open(input_json, "r", encoding="utf-8") as f:
        data = json.load(f)
        
    print(f"Total entries: {len(data)}")
    
    # 构造任务列表
    tasks = [{"fname": fname, "info": info} for fname, info in data.items()]
    
    chunk_size = len(tasks) // len(gpus) + 1
    processes = []
    
    for i, gpu_id in enumerate(gpus):
        start = i * chunk_size
        end = min((i + 1) * chunk_size, len(tasks))
        subset = tasks[start:end]
        if not subset:
            continue
            
        p = Process(target=process_partition, 
                    args=(gpu_id, subset, lora_path, base_output_dir))
        p.start()
        processes.append(p)
        
    for p in processes:
        p.join()
        
    print("All tasks finished.")

if __name__ == "__main__":
    main()