import gc
import os

import tempfile
import warnings
os.environ.setdefault("GRADIO_SSR_MODE", "false")
warnings.filterwarnings(
    "ignore",
    message=".*HTTP_422_UNPROCESSABLE_ENTITY.*",
    category=DeprecationWarning,
)
#!!!!!

try:
    import spaces
except ImportError:
    class _SpacesFallback:
        @staticmethod
        def GPU(*args, **kwargs):
            def decorator(fn):
                return fn

            return decorator

    spaces = _SpacesFallback()


import gradio as gr
import torch
from diffusers import AutoencoderKLWan, WanPipeline
from diffusers import HunyuanVideo15Pipeline, HunyuanVideo15ImageToVideoPipeline
from diffusers.utils import export_to_video
from huggingface_hub import snapshot_download
from PIL import Image
from transformers import AutoTokenizer, AutoModelForCausalLM


# ---------------------------------------------------------------------------
# Model registry
# ---------------------------------------------------------------------------
MODEL_OPTIONS = {
    "HunyuanVideo-1.5-T2V": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
    "HunyuanVideo-1.5-I2V": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v",
    "Wan2.2-TI2V-5B": "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
    "Wan2.2-T2V-A14B": "Wan-AI/Wan2.2-T2V-A14B",
}

LLM_MODEL_ID = "Qwen/Qwen3.6-35B-A3B"

DEFAULT_MODEL_LABEL = os.getenv("DEFAULT_MODEL_LABEL", "HunyuanVideo-1.5-T2V")
PRELOAD_MODELS = os.getenv("PRELOAD_MODELS", "HunyuanVideo-1.5-T2V")
OUTPUT_DIR = os.getenv("OUTPUT_DIR", tempfile.gettempdir())
CACHE_ROOT = os.getenv("MODEL_CACHE_DIR", "/data" if os.path.isdir("/data") else tempfile.gettempdir())
HY_T2V_CKPT_DIR = os.getenv("HY_T2V_CKPT_DIR", os.path.join(CACHE_ROOT, "HunyuanVideo-1.5-T2V"))
HY_I2V_CKPT_DIR = os.getenv("HY_I2V_CKPT_DIR", os.path.join(CACHE_ROOT, "HunyuanVideo-1.5-I2V"))
WAN_CKPT_DIR = os.getenv("WAN_CKPT_DIR", os.path.join(CACHE_ROOT, "Wan2.2-TI2V-5B"))
WAN_14B_CKPT_DIR = os.getenv("WAN_14B_CKPT_DIR", os.path.join(CACHE_ROOT, "Wan2.2-T2V-A14B"))
LLM_CKPT_DIR = os.getenv("LLM_CKPT_DIR", os.path.join(CACHE_ROOT, "Qwen3.6-35B-A3B"))

# ---------------------------------------------------------------------------
# Global state
# ---------------------------------------------------------------------------
pipe = None
loaded_model_id = None

# LLM state (kept separate from video pipelines)
llm_model = None
llm_tokenizer = None


# ===========================================================================
# LLM (Qwen3.6-35B-A3B) — prompt enhancement
# ===========================================================================
def _unload_llm():
    """Free the Qwen LLM from GPU memory."""
    global llm_model, llm_tokenizer
    if llm_model is not None:
        del llm_model
        llm_model = None
    if llm_tokenizer is not None:
        del llm_tokenizer
        llm_tokenizer = None
    gc.collect()
    torch.cuda.empty_cache()


def _load_llm():
    """Load Qwen3.6-35B-A3B for prompt enhancement (lazy-loaded, cached)."""
    global llm_model, llm_tokenizer
    if llm_model is not None and llm_tokenizer is not None:
        return llm_model, llm_tokenizer

    model_id = LLM_CKPT_DIR
    llm_tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
    llm_model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
    )
    llm_model.eval()
    return llm_model, llm_tokenizer


def enhance_prompt(prompt):
    """Use Qwen LLM to expand a short prompt into a detailed video description."""
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt to enhance.")

    # Unload video pipeline to free GPU memory for LLM
    global pipe, loaded_model_id
    if pipe is not None:
        del pipe
        pipe = None
        loaded_model_id = None
        gc.collect()
        torch.cuda.empty_cache()

    model, tokenizer = _load_llm()

    messages = [
        {
            "role": "system",
            "content": (
                "You are a professional video prompt engineer. Given a short description, "
                "expand it into a detailed, cinematic video generation prompt in English. "
                "Include camera angles, lighting, motion, atmosphere, and visual details. "
                "Output ONLY the enhanced prompt, nothing else. Keep it under 200 words."
            ),
        },
        {"role": "user", "content": prompt.strip()},
    ]

    text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(text, return_tensors="pt").to(model.device)

    with torch.inference_mode():
        outputs = model.generate(
            **inputs,
            max_new_tokens=300,
            temperature=0.7,
            do_sample=True,
        )

    # Decode only the newly generated tokens (exclude the input prompt)
    new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
    result = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()

    # Clean up possible thinking tags from Qwen
    if "<think>" in result:
        # Remove everything between <think>...</think>
        import re
        result = re.sub(r"<think>.*?</think>", "", result, flags=re.DOTALL).strip()

    return result if result else prompt.strip()


def enhance_prompt_ui(prompt):
    """UI wrapper for prompt enhancement with status feedback."""
    try:
        enhanced = enhance_prompt(prompt)
        return enhanced, f"Prompt enhanced successfully using {LLM_MODEL_ID}."
    except Exception as e:
        return prompt, f"Enhancement failed: {e}. Using original prompt."


def download_llm():
    """Download Qwen LLM weights to local cache."""
    if os.path.exists(os.path.join(LLM_CKPT_DIR, "config.json")):
        return f"LLM ({LLM_MODEL_ID}) already downloaded."

    snapshot_download(
        repo_id=LLM_MODEL_ID,
        local_dir=LLM_CKPT_DIR,
        local_dir_use_symlinks=False,
    )
    return f"Downloaded LLM: {LLM_MODEL_ID}"


# ===========================================================================
# Video pipeline loading
# ===========================================================================
def _resolve_model_id(model_label):
    """Map a model label to its checkpoint directory."""
    mapping = {
        "HunyuanVideo-1.5-T2V": HY_T2V_CKPT_DIR,
        "HunyuanVideo-1.5-I2V": HY_I2V_CKPT_DIR,
        "Wan2.2-TI2V-5B": WAN_CKPT_DIR,
        "Wan2.2-T2V-A14B": WAN_14B_CKPT_DIR,
    }
    if model_label not in mapping:
        raise gr.Error(f"Unknown model: {model_label}")
    return mapping[model_label]


def _load_pipeline(model_label):
    global loaded_model_id
    global pipe

    model_id = _resolve_model_id(model_label)

    if pipe is not None and loaded_model_id == model_id:
        return pipe

    # Unload LLM before loading video pipeline
    _unload_llm()

    if pipe is not None:
        del pipe
        pipe = None
        gc.collect()
        torch.cuda.empty_cache()

    # --- HunyuanVideo-1.5 T2V ---
    if model_label == "HunyuanVideo-1.5-T2V":
        pipe = HunyuanVideo15Pipeline.from_pretrained(
            model_id,
            torch_dtype=torch.bfloat16,
        )
        pipe.to("cuda")
        pipe.vae.enable_tiling()
        loaded_model_id = model_id
        return pipe

    # --- HunyuanVideo-1.5 I2V ---
    if model_label == "HunyuanVideo-1.5-I2V":
        pipe = HunyuanVideo15ImageToVideoPipeline.from_pretrained(
            model_id,
            torch_dtype=torch.bfloat16,
        )
        pipe.to("cuda")
        pipe.vae.enable_tiling()
        loaded_model_id = model_id
        return pipe

    # --- Wan2.2-TI2V-5B ---
    if model_label == "Wan2.2-TI2V-5B":
        vae = AutoencoderKLWan.from_pretrained(
            model_id,
            subfolder="vae",
            torch_dtype=torch.float32,
        )
        pipe = WanPipeline.from_pretrained(
            model_id,
            vae=vae,
            torch_dtype=torch.bfloat16,
        )
        pipe.to("cuda")
        loaded_model_id = model_id
        return pipe

    # --- Wan2.2-T2V-A14B (14B parameter text-to-video) ---
    if model_label == "Wan2.2-T2V-A14B":
        vae = AutoencoderKLWan.from_pretrained(
            model_id,
            subfolder="vae",
            torch_dtype=torch.float32,
        )
        pipe = WanPipeline.from_pretrained(
            model_id,
            vae=vae,
            torch_dtype=torch.bfloat16,
        )
        pipe.to("cuda")
        loaded_model_id = model_id
        return pipe

    raise gr.Error(f"Unsupported model: {model_label}")


def _resize_image(image, width, height):
    image = image.convert("RGB")
    return image.resize((width, height), Image.Resampling.LANCZOS)


# ===========================================================================
# Weight downloading
# ===========================================================================
def _ensure_wan_weights():
    if os.path.exists(os.path.join(WAN_CKPT_DIR, "model_index.json")):
        return WAN_CKPT_DIR

    return snapshot_download(
        repo_id=MODEL_OPTIONS["Wan2.2-TI2V-5B"],
        local_dir=WAN_CKPT_DIR,
        local_dir_use_symlinks=False,
    )


def _ensure_wan_14b_weights():
    if os.path.exists(os.path.join(WAN_14B_CKPT_DIR, "model_index.json")):
        return WAN_14B_CKPT_DIR

    return snapshot_download(
        repo_id=MODEL_OPTIONS["Wan2.2-T2V-A14B"],
        local_dir=WAN_14B_CKPT_DIR,
        local_dir_use_symlinks=False,
    )


def _ensure_hunyuan_weights(model_label):
    if model_label == "HunyuanVideo-1.5-T2V":
        ckpt_dir = HY_T2V_CKPT_DIR
    else:
        ckpt_dir = HY_I2V_CKPT_DIR

    if os.path.exists(os.path.join(ckpt_dir, "model_index.json")):
        return ckpt_dir

    return snapshot_download(
        repo_id=MODEL_OPTIONS[model_label],
        local_dir=ckpt_dir,
        local_dir_use_symlinks=False,
    )


def _model_cache_ready(model_label):
    if model_label in ("HunyuanVideo-1.5-T2V", "HunyuanVideo-1.5-I2V"):
        ckpt_dir = HY_T2V_CKPT_DIR if model_label == "HunyuanVideo-1.5-T2V" else HY_I2V_CKPT_DIR
        return os.path.exists(os.path.join(ckpt_dir, "model_index.json"))

    if model_label == "Wan2.2-TI2V-5B":
        return os.path.exists(os.path.join(WAN_CKPT_DIR, "model_index.json"))

    if model_label == "Wan2.2-T2V-A14B":
        return os.path.exists(os.path.join(WAN_14B_CKPT_DIR, "model_index.json"))

    return False


def download_model_assets(model_label):
    if model_label not in MODEL_OPTIONS:
        raise gr.Error("Please choose a supported model.")

    if model_label in ("HunyuanVideo-1.5-T2V", "HunyuanVideo-1.5-I2V"):
        _ensure_hunyuan_weights(model_label)
        return f"Downloaded {model_label}: Diffusers model weights"

    if model_label == "Wan2.2-TI2V-5B":
        _ensure_wan_weights()
        return f"Downloaded {model_label}: Diffusers model weights"

    if model_label == "Wan2.2-T2V-A14B":
        _ensure_wan_14b_weights()
        return f"Downloaded {model_label}: Diffusers model weights"

    raise gr.Error(f"Unsupported model: {model_label}")


def _preload_configured_models():
    if PRELOAD_MODELS.strip().lower() in {"", "0", "false", "none", "off"}:
        return "Startup preload disabled. Use the download button before generating."

    if PRELOAD_MODELS.strip().lower() == "all":
        model_labels = list(MODEL_OPTIONS.keys())
    else:
        requested = [item.strip() for item in PRELOAD_MODELS.split(",")]
        model_labels = [item for item in requested if item in MODEL_OPTIONS]

    if not model_labels:
        return f"No valid PRELOAD_MODELS entries found: {PRELOAD_MODELS}"

    messages = []
    for model_label in model_labels:
        try:
            messages.append(download_model_assets(model_label))
        except Exception as error:
            messages.append(f"Failed to download {model_label}: {error}")

    return "\n".join(messages)


# ===========================================================================
# Video generation
# ===========================================================================
def _duration(model_label, prompt, image, negative_prompt, width, height, frames, steps, guidance_scale, seed, enhance):
    if "HunyuanVideo" in model_label:
        base_seconds = 120
    elif "A14B" in model_label:
        base_seconds = 150  # 14B model needs more time
    else:
        base_seconds = 90
    extra = 60 if enhance else 0  # LLM enhancement overhead
    # ZeroGPU xlarge max is 300s
    return min(300, max(60, int(base_seconds + steps * 4 + frames * 1.0 + extra)))


def _call_pipeline(pipeline, prompt, image, negative_prompt, width, height, frames, steps, guidance_scale, generator):
    kwargs = {
        "prompt": prompt,
        "width": width,
        "height": height,
        "num_frames": frames,
        "num_inference_steps": steps,
        "guidance_scale": guidance_scale,
        "generator": generator,
    }

    if image is not None:
        kwargs["image"] = image

    if negative_prompt:
        kwargs["negative_prompt"] = negative_prompt

    try:
        return pipeline(**kwargs)
    except TypeError as error:
        raise gr.Error(f"The selected pipeline rejected these inputs: {error}") from error


@spaces.GPU(size="xlarge", duration=_duration)
def generate_video(
    model_label,
    prompt,
    image,
    negative_prompt,
    width,
    height,
    frames,
    steps,
    guidance_scale,
    seed,
    enhance,
):
    if model_label not in MODEL_OPTIONS:
        raise gr.Error("Please choose a supported model.")

    # HunyuanVideo I2V requires an input image
    if image is None and model_label == "HunyuanVideo-1.5-I2V":
        raise gr.Error("HunyuanVideo-1.5-I2V requires an input image for image-to-video generation.")

    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    # --- Optional: enhance prompt with Qwen LLM ---
    if enhance:
        try:
            prompt = enhance_prompt(prompt)
        except Exception as e:
            raise gr.Error(f"Prompt enhancement failed: {e}")

    width = int(width)
    height = int(height)
    frames = int(frames)
    steps = int(steps)
    seed = int(seed)

    # Force model-native resolutions before validation
    if "HunyuanVideo" in model_label:
        width, height = (848, 480) if width >= height else (480, 848)
    elif model_label in ("Wan2.2-TI2V-5B", "Wan2.2-T2V-A14B"):
        width, height = (1280, 704) if width >= height else (704, 1280)

    # HunyuanVideo VAE compresses 16x, Wan VAE compresses 32x
    divisor = 16 if "HunyuanVideo" in model_label else 32
    if width % divisor != 0 or height % divisor != 0:
        raise gr.Error(f"Width and height must be divisible by {divisor}.")

    # Frame constraint: 4n+1 for both HunyuanVideo and Wan
    if (frames - 1) % 4 != 0:
        raise gr.Error("Frame count must be 4n + 1, for example 25, 29, 33, 37, 41, 45, 49, 53, 57, 61, 81, 97, or 121.")

    if not _model_cache_ready(model_label):
        raise gr.Error(f"{model_label} is not downloaded yet. Click Download selected model first.")

    resized_image = _resize_image(image, width, height) if image is not None else None
    generator = None

    if seed >= 0:
        generator = torch.Generator(device="cuda").manual_seed(seed)

    pipeline = _load_pipeline(model_label)

    with torch.inference_mode():
        output = _call_pipeline(
            pipeline=pipeline,
            prompt=prompt.strip(),
            image=resized_image,
            negative_prompt=negative_prompt.strip(),
            width=width,
            height=height,
            frames=frames,
            steps=steps,
            guidance_scale=float(guidance_scale),
            generator=generator,
        )

    video_file = tempfile.NamedTemporaryFile(
        prefix=f"{model_label.lower().replace('.', '_').replace('-', '_')}_",
        suffix=".mp4",
        dir=OUTPUT_DIR,
        delete=False,
    )
    video_file.close()
    video_path = video_file.name
    export_to_video(output.frames[0], video_path, fps=24)

    gc.collect()
    torch.cuda.empty_cache()
    return video_path


# ===========================================================================
# Gradio UI
# ===========================================================================
example_prompt = (
    "A cinematic close-up of the subject turning toward the camera, soft natural "
    "light, detailed motion, realistic texture, synchronized ambient audio."
)

with gr.Blocks(title="Video Model ZeroGPU Test") as demo:
    gr.Markdown("# Video Model ZeroGPU Test")

    with gr.Row():
        with gr.Column(scale=1):
            model_input = gr.Dropdown(
                label="Model",
                choices=list(MODEL_OPTIONS.keys()),
                value=DEFAULT_MODEL_LABEL if DEFAULT_MODEL_LABEL in MODEL_OPTIONS else "HunyuanVideo-1.5-T2V",
            )
            download_button = gr.Button("Download selected model")
            download_status = gr.Textbox(
                label="Download status",
                value="Startup preload has not started yet.",
                lines=3,
                interactive=False,
            )
            image_input = gr.Image(
                label="Start image (required for I2V models)",
                type="pil",
                sources=["upload", "clipboard"],
                height=320,
            )
            prompt_input = gr.Textbox(
                label="Prompt",
                value=example_prompt,
                lines=5,
            )

            # --- LLM Prompt Enhancement ---
            with gr.Row():
                enhance_checkbox = gr.Checkbox(
                    label=f"Enhance prompt with LLM ({LLM_MODEL_ID})",
                    value=False,
                )
                enhance_button = gr.Button("Enhance now", size="sm")
            enhance_status = gr.Textbox(
                label="Enhancement status",
                visible=False,
                interactive=False,
            )

            negative_prompt_input = gr.Textbox(
                label="Negative prompt",
                value="low quality, blurry, distorted, flickering, artifacts",
                lines=2,
            )

            with gr.Row():
                width_input = gr.Dropdown(
                    label="Width",
                    choices=[120, 180, 256, 320, 384, 448, 480, 512, 640],
                    value=640,
                )
                height_input = gr.Dropdown(
                    label="Height",
                    choices=[120, 180, 256, 320, 384, 448, 480, 512, 640],
                    value=480,
                )
                frames_input = gr.Dropdown(
                    label="Frames",
                    choices=[17, 21, 25, 29, 33, 37, 41, 45, 49, 53, 57, 61, 65, 69, 73, 77, 81, 85, 89, 93, 97, 101, 105, 109, 113, 117, 121],
                    value=81,
                )

            with gr.Row():
                steps_input = gr.Slider(
                    label="Inference steps",
                    minimum=4,
                    maximum=50,
                    step=1,
                    value=30,
                )
                guidance_input = gr.Slider(
                    label="CFG scale",
                    minimum=1.0,
                    maximum=8.0,
                    step=0.1,
                    value=6.0,
                )
                seed_input = gr.Number(
                    label="Seed (-1 for random)",
                    value=0,
                    precision=0,
                )

            generate_button = gr.Button("Generate", variant="primary")

        with gr.Column(scale=1):
            video_output = gr.Video(label="Generated video", format="mp4")

    # --- Event bindings ---
    demo.load(
        fn=_preload_configured_models,
        outputs=download_status,
    )

    download_button.click(
        fn=download_model_assets,
        inputs=model_input,
        outputs=download_status,
    )

    # Prompt enhancement button: rewrites the prompt in-place
    enhance_button.click(
        fn=enhance_prompt_ui,
        inputs=prompt_input,
        outputs=[prompt_input, enhance_status],
    )

    generate_button.click(
        fn=generate_video,
        inputs=[
            model_input,
            prompt_input,
            image_input,
            negative_prompt_input,
            width_input,
            height_input,
            frames_input,
            steps_input,
            guidance_input,
            seed_input,
            enhance_checkbox,
        ],
        outputs=video_output,
    )

if __name__ == "__main__":
    try:
        demo.queue(max_size=20).launch(ssr_mode=False)
    except TypeError:
        demo.queue(max_size=20).launch()

# End of file