import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces
import sys
import os
import time
import tempfile
import torch
import gradio as gr
import numpy as np
from PIL import Image
from huggingface_hub import hf_hub_download, snapshot_download


# ---------------------------------------------------------------------------
# Kill the per-call cuDNN / kernel-autotune warmup tax.
#
# On ZeroGPU every @spaces.GPU call runs in a *fresh* forked worker process, so
# any per-call (lazy, first-use) kernel autotuning is paid on every single
# invocation instead of amortizing once. The video pipeline's Conv3d patch
# embedding + SDPA attention would otherwise re-autotune cuDNN/kernels on the
# first denoising step of each call (the ~114s -> ~7s decaying step-time curve).
#
# Disabling cudnn.benchmark stops cuDNN from re-benchmarking algorithms per
# call; the tensor shapes fed into the DiT blocks / attention are already static
# across the denoising loop (fixed height/width/num_frames), so a single fixed
# algorithm choice is correct and avoids the re-selection warmup.
# ---------------------------------------------------------------------------

def _configure_static_kernels():
    """Pin cuDNN / SDPA kernel selection so nothing re-autotunes per call."""
    # cuDNN: do NOT benchmark/autotune conv algorithms (this is the main tax).
    torch.backends.cudnn.benchmark = False
    # Belt-and-suspenders: make sure nothing flipped it on.
    if hasattr(torch.backends.cudnn, "benchmark_limit"):
        torch.backends.cudnn.benchmark_limit = 0
    # Keep TF32 matmul/conv enabled for steady-state speed (shape-independent).
    try:
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
    except Exception:
        pass


# Apply at import time (main process); inherited by forked workers.
_configure_static_kernels()

# Ensure kairos package is importable (it's uploaded alongside app.py)
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from mmengine import Config

# ---------------------------------------------------------------------------
# Download model weights at module scope so ZeroGPU can pack them
# ---------------------------------------------------------------------------

MODELS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")

def _ensure_models():
    """Download all required model weights to local directories."""
    os.makedirs(MODELS_DIR, exist_ok=True)

    # Kairos DiT checkpoint
    kairos_dir = os.path.join(MODELS_DIR, "Kairos3.1-4B-robot-480P")
    if not os.path.exists(os.path.join(kairos_dir, "kairos-4B-robot-3.1-480P.safetensors")):
        print("Downloading Kairos DiT weights...")
        snapshot_download(
            "kairos-agi/Kairos3.1-4B-robot-480P",
            local_dir=kairos_dir,
        )
        print("Kairos DiT weights downloaded.")

    # Qwen3.5-2B text encoder
    qwen_dir = os.path.join(MODELS_DIR, "Qwen3.5-2B")
    if not os.path.exists(os.path.join(qwen_dir, "model.safetensors-00001-of-00001.safetensors")):
        print("Downloading Qwen3.5-2B text encoder...")
        snapshot_download(
            "Qwen/Qwen3.5-2B",
            local_dir=qwen_dir,
        )
        print("Qwen3.5-2B text encoder downloaded.")

    # Wan2.1 VAE
    vae_dir = os.path.join(MODELS_DIR, "Wan2.1-T2V-1.3B")
    if not os.path.exists(os.path.join(vae_dir, "Wan2.1_VAE.pth")):
        print("Downloading Wan2.1 VAE...")
        snapshot_download(
            "Wan-AI/Wan2.1-T2V-1.3B",
            local_dir=vae_dir,
            allow_patterns=["Wan2.1_VAE.pth"],
        )
        print("Wan2.1 VAE downloaded.")


_ensure_models()

# ---------------------------------------------------------------------------
# Build the Kairos pipeline config
# ---------------------------------------------------------------------------

KAIROS_MODEL_DIR = MODELS_DIR
pretrained_dit = os.path.join(KAIROS_MODEL_DIR, "Kairos3.1-4B-robot-480P", "kairos-4B-robot-3.1-480P.safetensors")
vae_path = os.path.join(KAIROS_MODEL_DIR, "Wan2.1-T2V-1.3B", "Wan2.1_VAE.pth")
qwen_text_encoder_path = os.path.join(KAIROS_MODEL_DIR, "Qwen3.5-2B")

pipeline_cfg = dict(
    type='KairosEmbodiedAPI',
    use_cfg_parallel=False,
    pipeline_type='KairosEmbodiedWAMPipeline',
    pipeline_args=dict(
        load_dit_fn='strict_load',
        vae_path=vae_path,
        text_encoder_config=dict(
            type='Qwen3_5_TextEncoder',
            from_pretrained=qwen_text_encoder_path,
        ),
        dit_config=dict(
            dit_type='Simple_Multi_DIT_Wrapper',
            video_dit={
                "dit_type": 'KairosDiTV2',
                "pretrained_path": pretrained_dit,
                "has_image_input": False,
                "patch_size": [1, 2, 2],
                "in_dim": 16,
                "dim": 2560,
                "ffn_dim": 10240,
                "freq_dim": 256,
                "text_dim": 2048,
                "out_dim": 16,
                "num_heads": 20,
                "num_layers": 32,
                "layers_settings": [
                    'SWA', 'SWA', 'DSWA', 'GATED',
                    'SWA', 'SWA', 'DSWA', 'GATED',
                    'SWA', 'SWA', 'DSWA', 'GATED',
                    'SWA', 'SWA', 'DSWA', 'GATED',
                    'SWA', 'SWA', 'DSWA', 'GATED',
                    'SWA', 'SWA', 'DSWA', 'GATED',
                    'SWA', 'SWA', 'DSWA', 'GATED',
                    'SWA', 'SWA', 'DSWA', 'GATED',
                ],
                "eps": 1e-6,
                "seperated_timestep": True,
                "require_clip_embedding": False,
                "require_vae_embedding": False,
                "fuse_vae_embedding_in_latents": True,
                "dilated_lengths": [4],
                "use_first_frame_cond": False,
                "use_seq_parallel": False,
                "use_tp_in_getaeddeltanet": False,
                "use_tp_in_self_attn": False,
                "attend_k0": False,
            },
        ),
    ),
)

# Initialize distributed (required by mmengine even for single-GPU)
from mmengine.dist import init_dist, get_dist_info
import torch.distributed as dist

# For ZeroGPU, we need to init distributed in a way that works without torchrun
os.environ.setdefault("RANK", "0")
os.environ.setdefault("WORLD_SIZE", "1")
os.environ.setdefault("MASTER_ADDR", "localhost")
os.environ.setdefault("MASTER_PORT", "29556")

# We won't call init_dist since we're not using torchrun;
# instead we mock the dist info the pipeline expects
if not dist.is_initialized():
    try:
        init_dist(launcher='pytorch')
    except Exception:
        pass

# Build the pipeline
print("Building Kairos pipeline...")
from kairos.apis.builder import build_model_pipeline
cfg = Config()
cfg.pipeline = pipeline_cfg

pipeline = build_model_pipeline(cfg.pipeline)
pipeline.eval()
print("Kairos pipeline built successfully.")


# ---------------------------------------------------------------------------
# Inference function
# ---------------------------------------------------------------------------

DEFAULT_NEGATIVE_PROMPT = (
    "bright tones, overexposed, static, blurred details, subtitles, style, works, "
    "paintings, images, static, overall gray, worst quality, low quality, JPEG compression "
    "residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, "
    "disfigured, deformed limbs, fused fingers, still picture, messy background, three legs, "
    "many people in the background, walking backwards, contorted human joints, objects floating "
    "against natural forces, abrupt shot changes"
)


def _save_video_frames(frames, fps=16):
    """Save a list of PIL frames as an mp4 video using imageio."""
    import imageio
    tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    tmp.close()
    frames_np = [np.array(f) for f in frames]
    writer = imageio.get_writer(tmp.name, fps=fps, codec="libx264", quality=5)
    for frame in frames_np:
        writer.append_data(frame)
    writer.close()
    return tmp.name


@spaces.GPU(duration=300)
def generate_video(
    input_image: Image.Image,
    prompt: str,
    negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
    seed: int = 0,
    num_frames: int = 81,
    height: int = 480,
    width: int = 832,
    cfg_scale: float = 5.0,
    num_inference_steps: int = 50,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a future video from an initial scene image and a task instruction.

    Args:
        input_image: The initial scene image (e.g., a robot arm in a workspace).
        prompt: A text description of the task the robot should perform.
        negative_prompt: Negative prompt to guide generation away from undesired features.
        seed: Random seed for reproducibility.
        num_frames: Number of video frames to generate.
        height: Output video height (must be divisible by 16).
        width: Output video width (must be divisible by 16).
        cfg_scale: Classifier-free guidance scale.
        num_inference_steps: Number of denoising steps.
    """
    # Re-apply inside the fresh ZeroGPU worker process: cuDNN benchmark state
    # must be set before the first conv/attention op runs in this fork, otherwise
    # cuDNN re-autotunes on the first denoising step of every call.
    _configure_static_kernels()

    if input_image is None:
        return None, "Please provide an input image."
    if not prompt.strip():
        return None, "Please provide a task instruction."

    # Add prompt prefix for TI2V mode
    prompt_prefix = "high-quality video, realistic motion, single continuous shot, no jump cuts, smooth motion. "
    full_prompt = prompt_prefix + prompt

    # Run inference
    start = time.perf_counter()

    result = pipeline(
        prompt=full_prompt,
        negative_prompt=negative_prompt,
        input_image=[input_image],
        seed=seed,
        tiled=True,
        height=height,
        width=width,
        num_frames=num_frames,
        cfg_scale=cfg_scale,
        num_inference_steps=num_inference_steps,
    )

    elapsed = time.perf_counter() - start
    print(f"Inference took {elapsed:.2f}s")

    # Extract video frames
    if isinstance(result, dict):
        video_frames = result.get("video", None)
        action = result.get("action", None)
    else:
        video_frames = result
        action = None

    if video_frames is None:
        return None, "Generation failed - no output produced."

    # video_frames is a list of PIL Images (from vae_output_to_video)
    if isinstance(video_frames, list):
        frames = video_frames
    elif isinstance(video_frames, torch.Tensor):
        # Convert tensor [T, H, W, C] or [C, T, H, W] to list of PIL Images
        frames = []
        for i in range(video_frames.shape[0]):
            frame = video_frames[i]
            if isinstance(frame, torch.Tensor):
                frame = frame.cpu().numpy()
            frame = ((frame + 1) / 2 * 255).clip(0, 255).astype(np.uint8)
            if frame.ndim == 3 and frame.shape[0] in (3,):
                frame = frame.transpose(1, 2, 0)
            frames.append(Image.fromarray(frame))
    else:
        return None, f"Unexpected output type: {type(video_frames)}"

    if len(frames) == 1:
        # Single frame - return as image
        return frames[0], f"Generation completed in {elapsed:.2f}s (1 frame)"
    else:
        video_path = _save_video_frames(frames, fps=16)
        action_text = ""
        if action is not None:
            action_text = f"\nAction tensor shape: {action.shape}"
        return video_path, f"Generation completed in {elapsed:.2f}s ({len(frames)} frames){action_text}"


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# 🤖 Kairos: Native World-Action Model")
        gr.Markdown(
            "A 4B-parameter world-action model for physical AI. Given an initial scene "
            "image and a task instruction, Kairos predicts the future video showing the "
            "robot performing the described action.\n\n"
            "📄 [Paper](https://arxiv.org/abs/2606.16533) · "
            "💻 [GitHub](https://github.com/kairos-agi/kairos-sensenova) · "
            "🤗 [Model](https://huggingface.co/kairos-agi/Kairos3.1-4B-robot-480P)"
        )

        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(
                    label="Initial Scene Image",
                    type="pil",
                    height=300,
                )
                prompt = gr.Textbox(
                    label="Task Instruction",
                    placeholder="Describe what the robot should do...",
                    lines=2,
                )
                run_btn = gr.Button("Generate Video", variant="primary")

            with gr.Column(scale=1):
                output_video = gr.Video(label="Generated Future Video")
                status = gr.Textbox(label="Status", interactive=False)

        with gr.Accordion("Advanced Settings", open=False):
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                value=DEFAULT_NEGATIVE_PROMPT,
                lines=3,
            )
            seed = gr.Number(label="Seed", value=0, precision=0)
            num_frames = gr.Slider(label="Number of Frames", minimum=5, maximum=81, value=81, step=4)
            height = gr.Number(label="Height", value=480, precision=0)
            width = gr.Number(label="Width", value=832, precision=0)
            cfg_scale = gr.Slider(label="CFG Scale", minimum=1.0, maximum=15.0, value=5.0, step=0.5)
            num_inference_steps = gr.Slider(label="Inference Steps", minimum=10, maximum=100, value=50, step=5)

        gr.Examples(
            examples=[
                ["robot_1.jpg", "The robotic arm sweeps a green cloth across the wooden surface."],
                ["robot_2.jpg", "Use the right hand to pick up orange from pink plate to pale turquoise plate."],
            ],
            inputs=[input_image, prompt],
            outputs=[output_video, status],
            fn=generate_video,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        fn=generate_video,
        inputs=[
            input_image, prompt, negative_prompt, seed,
            num_frames, height, width, cfg_scale, num_inference_steps,
        ],
        outputs=[output_video, status],
    )


demo.launch(mcp_server=True)