import gradio as gr
import spaces
from huggingface_hub import hf_hub_download
import requests
import json
import time
import random
import os

# ====================== DOWNLOAD MODELOV ======================
def download_models():
    print("📥 Downloading models...")
    # Qwen
    hf_hub_download("Comfy-Org/Qwen-Image-Edit_ComfyUI", "split_files/diffusion_models/qwen_image_edit_2509_fp8_e4m3fn.safetensors", local_dir="models/diffusion_models", force_download=False)
    hf_hub_download("Comfy-Org/Qwen-Image_ComfyUI", "split_files/vae/qwen_image_vae.safetensors", local_dir="models/vae", force_download=False)
    hf_hub_download("Comfy-Org/Qwen-Image_ComfyUI", "split_files/text_encoders/qwen_2.5_vl_7b_fp8_scaled.safetensors", local_dir="models/text_encoders", force_download=False)
    hf_hub_download("lightx2v/Qwen-Image-Lightning", "Qwen-Image-Edit-2509/Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors", local_dir="models/loras", force_download=False)

    # Wan 2.2
    hf_hub_download("Comfy-Org/Wan_2.2_ComfyUI_Repackaged", "split_files/vae/wan_2.1_vae.safetensors", local_dir="models/vae", force_download=False)
    hf_hub_download("NSFW-API/NSFW-Wan-UMT5-XXL", "nsfw_wan_umt5-xxl_fp8_scaled.safetensors", local_dir="models/text_encoders", force_download=False)
    hf_hub_download("FX-FeiHou/wan2.2-Remix", "NSFW/Wan2.2_Remix_NSFW_i2v_14b_low_lighting_v2.0.safetensors", local_dir="models/diffusion_models", force_download=False)
    hf_hub_download("FX-FeiHou/wan2.2-Remix", "NSFW/Wan2.2_Remix_NSFW_i2v_14b_high_lighting_v2.0.safetensors", local_dir="models/diffusion_models", force_download=False)
    hf_hub_download("Comfy-Org/Wan_2.2_ComfyUI_Repackaged", "split_files/loras/wan2.2_i2v_lightx2v_4steps_lora_v1_low_noise.safetensors", local_dir="models/loras", force_download=False)
    hf_hub_download("Comfy-Org/Wan_2.2_ComfyUI_Repackaged", "split_files/loras/wan2.2_i2v_lightx2v_4steps_lora_v1_high_noise.safetensors", local_dir="models/loras", force_download=False)

    print("✅ Models ready!")

download_models()

COMFY_URL = "http://127.0.0.1:8188"

@spaces.GPU(duration=180)
def run_workflow(prompt_image, prompt_video, bedroom_image, woman_image, steps, cfg, length, seed):
    if bedroom_image is None or woman_image is None:
        return None, None, "❌ Naloži obe sliki!"

    if seed == -1:
        seed = random.randint(0, 999999999)

    try:
        with open("workflow_api.json", "r") as f:
            workflow = json.load(f)

        # Tukaj lahko kasneje dinamično spreminjaš node vrednosti
        # Za zdaj zaženemo osnovni workflow
        response = requests.post(f"{COMFY_URL}/prompt", json={"prompt": workflow})
        prompt_id = response.json().get("prompt_id")

        status = f"🔄 Running workflow... (ID: {prompt_id})"
        
        # Čakaj na končanje
        for _ in range(120):
            time.sleep(3)
            history = requests.get(f"{COMFY_URL}/history/{prompt_id}").json()
            if prompt_id in history:
                status = "✅ Workflow finished!"
                break

        return None, None, status

    except Exception as e:
        return None, None, f"Error: {str(e)}"

with gr.Blocks(title="Wan2.2 Video Generator") as demo:
    gr.Markdown("# 🛏️ Qwen Image Edit + Wan2.2 Video")

    with gr.Row():
        with gr.Column():
            prompt_image = gr.Textbox(label="Prompt za sliko", lines=3, value="Natural integration, cozy bedroom...")
            prompt_video = gr.Textbox(label="Prompt za video", lines=3, value="smooth cinematic motion...")

            with gr.Row():
                img1 = gr.Image(label="Image 1", type="filepath")
                img2 = gr.Image(label="Image 2", type="filepath")

            with gr.Accordion("Advanced", open=False):
                steps = gr.Slider(1, 20, value=4, step=1, label="Steps")
                cfg = gr.Slider(0.5, 3.0, value=1.0, step=0.1, label="CFG")
                length = gr.Slider(16, 121, value=81, step=1, label="Video Length")
                seed = gr.Number(value=-1, label="Seed (-1 = random)")

            btn = gr.Button("🚀 Generate", variant="primary")

        with gr.Column():
            out_image = gr.Image(label="Generated Image")
            out_video = gr.Video(label="Generated Video")
            status = gr.Textbox(label="Status", interactive=False)

    btn.click(
        run_workflow,
        inputs=[prompt_image, prompt_video, img1, img2, steps, cfg, length, seed],
        outputs=[out_image, out_video, status]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, theme=gr.themes.Soft())