import os
import tempfile
from pathlib import Path

# Alguns workers ZeroGPU não expõem NVML completamente ao allocator do PyTorch.
# Desativar o cache evita o assert interno sem alterar o resultado do modelo.
os.environ.setdefault("PYTORCH_NO_CUDA_MEMORY_CACHING", "1")

import gradio as gr
import spaces
import torch
from diffusers import CogVideoXImageToVideoPipeline
from diffusers.utils import export_to_video
from PIL import Image

MODEL_ID = "THUDM/CogVideoX1.5-5B-I2V"
PIPE = None


def get_pipeline():
    global PIPE
    if PIPE is None:
        print(f"Loading model: {MODEL_ID}", flush=True)
        PIPE = CogVideoXImageToVideoPipeline.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.bfloat16,
        )
        # O offload evita mover o pipeline inteiro de uma só vez para a GPU.
        PIPE.enable_model_cpu_offload()
        print("Model loaded on CUDA", flush=True)
    return PIPE


def resize_image(image: Image.Image, profile: str):
    image = image.convert("RGB")
    if profile == "1080p (upscale final)":
        target_w, target_h = 1360, 768
    else:
        target_w, target_h = 1024, 576

    target_ratio = target_w / target_h
    source_ratio = image.width / image.height
    if source_ratio > target_ratio:
        crop_h = image.height
        crop_w = round(crop_h * target_ratio)
    else:
        crop_w = image.width
        crop_h = round(crop_w / target_ratio)

    left = max(0, (image.width - crop_w) // 2)
    top = max(0, (image.height - crop_h) // 2)
    image = image.crop((left, top, left + crop_w, top + crop_h))
    return image.resize((target_w, target_h), Image.Resampling.LANCZOS), target_w, target_h


@spaces.GPU(duration=120)
def generate_video(image, prompt, profile, seed):
    if image is None:
        raise gr.Error("Envie uma imagem antes de gerar o vídeo.")
    if not prompt or not prompt.strip():
        raise gr.Error("Escreva o movimento que você quer animar.")

    try:
        pipeline = get_pipeline()
        image, width, height = resize_image(image, profile)
        generator = torch.Generator(device="cuda").manual_seed(int(seed or 42))
        full_prompt = (
            "Animate the input image while preserving the exact identity, face, clothing, "
            "objects, composition and environment. Focus on natural physically plausible motion. "
            + prompt.strip()
            + ". Smooth cinematic movement, consistent character, no morphing, no extra limbs, "
            "no scene change, no camera cuts."
        )

        result = pipeline(
            prompt=full_prompt,
            image=image,
            height=height,
            width=width,
            num_frames=33,
            guidance_scale=6.0,
            num_inference_steps=12,
            generator=generator,
        )
        frames = result.frames[0]
        output_path = Path(tempfile.mkdtemp(prefix="motion-ai-")) / "motion.mp4"
        export_to_video(frames, str(output_path), fps=8)
        return str(output_path)
    except Exception as exc:
        print(f"Generation error: {type(exc).__name__}: {exc}", flush=True)
        raise gr.Error(f"Não foi possível gerar o vídeo: {type(exc).__name__}: {exc}") from exc


with gr.Blocks(title="Motion AI Studio") as demo:
    gr.Markdown(
        "# Motion AI Studio\n"
        "Anime uma imagem com um comando de movimento. A GPU é solicitada somente ao gerar."
    )
    with gr.Row():
        with gr.Column():
            image = gr.Image(type="pil", label="Imagem de referência")
            prompt = gr.Textbox(
                label="Descreva o movimento",
                placeholder="Ex.: o homem começa a correr pela rua enquanto a câmera o acompanha lateralmente",
                lines=4,
            )
            profile = gr.Radio(
                ["720p", "1080p (upscale final)"],
                value="720p",
                label="Perfil",
            )
            seed = gr.Number(value=42, precision=0, label="Seed")
            generate = gr.Button("Gerar animação", variant="primary")
        with gr.Column():
            result = gr.Video(label="Resultado", format="mp4")
            gr.Markdown(
                "**Recomendação:** comece com 720p. O perfil 1080p gera a animação em resolução menor."
            )

    generate.click(generate_video, [image, prompt, profile, seed], result)

if __name__ == "__main__":
    demo.queue(max_size=8).launch()