import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")

import random
import re
import tempfile
from typing import Any

import spaces
import torch
import gradio as gr
import numpy as np
from diffusers import LTX2ConditionPipeline, LTX2VideoTransformer3DModel
from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition
from diffusers.pipelines.ltx2.utils import DEFAULT_NEGATIVE_PROMPT, DISTILLED_SIGMA_VALUES
from diffusers.utils import encode_video

SULPHUR_TRANSFORMER_ID = "CalamitousFelicitousness/LTX-2.3-Sulphur2-Distilled-Diffusers"
BASE_PIPELINE_ID = "diffusers/LTX-2.3-Distilled-Diffusers"
SULPHUR_REPO = "SulphurAI/Sulphur-2-base"

MAX_SEED = np.iinfo(np.int32).max
FPS = 24.0
DISTILLED_STEPS = 8

RESOLUTIONS = {
    "high": {"16:9": (768, 512), "9:16": (512, 768), "1:1": (768, 768)},
    "low": {"16:9": (640, 384), "9:16": (384, 640), "1:1": (512, 512)},
}

DEFAULT_PROMPT = (
    "An astronaut hatches from a fragile egg on the surface of the Moon, "
    "the shell cracking and peeling apart in gentle low-gravity motion. "
    "Fine lunar dust lifts and drifts outward with each movement, floating "
    "in slow arcs before settling back onto the ground."
)

_BLOCKED = (
    re.compile(
        r"\b(child|children|kid|minor|underage|teen(?:ager)?s?)\b.{0,80}"
        r"\b(nude|naked|sex|sexual|explicit|porn|nsfw)\b",
        re.I,
    ),
    re.compile(
        r"\b(nude|naked|sex|sexual|explicit|porn|nsfw)\b.{0,80}"
        r"\b(child|children|kid|minor|underage|teen(?:ager)?s?)\b",
        re.I,
    ),
    re.compile(r"\b(csam|child porn|rape|non[- ]consensual|revenge porn)\b", re.I),
)


def _prompt_allowed(prompt: str) -> bool:
    return not any(pattern.search(prompt) for pattern in _BLOCKED)


print("Loading Sulphur 2 distilled transformer + LTX-2.3 Distilled pipeline...")
transformer = LTX2VideoTransformer3DModel.from_pretrained(
    SULPHUR_TRANSFORMER_ID,
    subfolder="transformer",
    dtype=torch.bfloat16,
)
pipe = LTX2ConditionPipeline.from_pretrained(
    BASE_PIPELINE_ID,
    transformer=transformer,
    dtype=torch.bfloat16,
)

# Pack only the diffusion modules. Gemma stays on CPU so the ZeroGPU pack fits
# in 48GB (a full .to("cuda") packed ~70GB and forced xlarge, which queued slowly).
for name in ("transformer", "vae", "audio_vae", "vocoder", "connectors"):
    module = getattr(pipe, name, None)
    if module is not None:
        module.to("cuda")
pipe.vae.enable_tiling()

# Diffusers puts token IDs on `_execution_device` (CUDA). Gemma weights are on CPU,
# so encode on CPU and move the resulting embeddings to the GPU afterwards.
_orig_get_gemma_prompt_embeds = pipe._get_gemma_prompt_embeds


def _get_gemma_prompt_embeds_cpu(*args, **kwargs):
    kwargs["device"] = torch.device("cpu")
    prompt_embeds, prompt_attention_mask = _orig_get_gemma_prompt_embeds(*args, **kwargs)
    device = pipe._execution_device
    return prompt_embeds.to(device), prompt_attention_mask.to(device)


pipe._get_gemma_prompt_embeds = _get_gemma_prompt_embeds_cpu
print("Pipeline ready.")


def detect_aspect_ratio(image) -> str:
    """Return the closest 16:9, 9:16, or 1:1 ratio for an optional PIL image."""
    if image is None:
        return "16:9"
    if hasattr(image, "size"):
        width, height = image.size
    elif hasattr(image, "shape"):
        height, width = image.shape[:2]
    else:
        return "16:9"
    ratio = width / max(height, 1)
    candidates = {"16:9": 16 / 9, "9:16": 9 / 16, "1:1": 1.0}
    return min(candidates, key=lambda key: abs(ratio - candidates[key]))


def on_image_upload(image, high_res: bool):
    """Snap width/height to the image aspect when a first frame is uploaded."""
    aspect = detect_aspect_ratio(image)
    tier = "high" if high_res else "low"
    width, height = RESOLUTIONS[tier][aspect]
    return gr.update(value=width), gr.update(value=height)


def on_highres_toggle(image, high_res: bool):
    """Update resolution when the high-res toggle changes."""
    return on_image_upload(image, high_res)


def _gpu_duration(
    input_image,
    prompt: str,
    duration: float,
    enhance_prompt: bool,
    seed: int,
    randomize_seed: bool,
    height: int,
    width: int,
    *args,
    **kwargs,
) -> int:
    del input_image, prompt, enhance_prompt, seed, randomize_seed, args, kwargs
    extra = 15 if int(height) * int(width) >= 768 * 512 else 0
    # CPU Gemma encode is slow; keep a buffer so the lease isn't cut mid-run.
    return min(120, int(70 + float(duration) * 10 + extra))


@spaces.GPU(duration=_gpu_duration)
def generate_video(
    input_image: Any,
    prompt: str,
    duration: float,
    enhance_prompt: bool = False,
    seed: int = 42,
    randomize_seed: bool = True,
    height: int = 384,
    width: int = 640,
    progress=gr.Progress(track_tqdm=True),
) -> tuple[str, int]:
    """Generate a short audio-video clip from a text prompt and optional first-frame image.

    Args:
        input_image: optional PIL image used as the first frame (image-to-video).
        prompt: description of one shot: subject, motion, camera, lighting, and sound.
        duration: clip length in seconds.
        enhance_prompt: reserved for prompt enhancement (not used by this demo).
        seed: RNG seed for reproducible generation.
        randomize_seed: pick a fresh random seed when True.
        height: output height in pixels.
        width: output width in pixels.

    Returns:
        Path to the generated .mp4 file and the seed that was used.
    """
    del progress, enhance_prompt
    prompt = (prompt or "").strip()
    if len(prompt) < 8:
        raise gr.Error("Please enter a more descriptive prompt.")
    if len(prompt) > 2000:
        raise gr.Error("Please keep the prompt under 2,000 characters.")
    if not _prompt_allowed(prompt):
        raise gr.Error("This public demo cannot process that prompt.")

    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
    num_frames = int(float(duration) * FPS) + 1
    num_frames = ((num_frames - 1 + 7) // 8) * 8 + 1
    generator = torch.Generator(device="cuda").manual_seed(current_seed)

    conditions = None
    if input_image is not None:
        conditions = [LTX2VideoCondition(frames=input_image, index=0, strength=1.0)]

    print(
        f"Generating {int(width)}x{int(height)}, {num_frames} frames "
        f"({duration}s), seed={current_seed}, i2v={conditions is not None}"
    )
    try:
        video, audio = pipe(
            conditions=conditions,
            prompt=prompt,
            negative_prompt=DEFAULT_NEGATIVE_PROMPT,
            height=int(height),
            width=int(width),
            num_frames=num_frames,
            frame_rate=FPS,
            num_inference_steps=DISTILLED_STEPS,
            sigmas=DISTILLED_SIGMA_VALUES,
            guidance_scale=1.0,
            audio_guidance_scale=1.0,
            enable_prompt_enhancement=False,
            generator=generator,
            output_type="np",
            return_dict=False,
        )
        output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
        output.close()
        encode_video(
            video[0],
            fps=int(FPS),
            audio=audio[0].float().cpu(),
            audio_sample_rate=pipe.vocoder.config.output_sampling_rate,
            output_path=output.name,
        )
    except gr.Error:
        raise
    except Exception as exc:
        raise gr.Error(f"Generation failed: {exc}") from exc

    return output.name, current_seed


CSS = """
#col-container { max-width: 1180px !important; margin: 0 auto; }
.hero h1 { font-size: clamp(1.8rem, 4vw, 2.8rem); margin-bottom: 0.25rem; }
.hero p { color: var(--body-text-color-subdued); }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="Sulphur 2 Demo") as demo:
    gr.Markdown(
        f"""
<div class="hero">

# 🌋 Sulphur 2 Demo

Text-to-video and image-to-video with synchronized audio, using
[{SULPHUR_REPO}](https://huggingface.co/{SULPHUR_REPO})
on the LTX 2.3 distilled Diffusers pipeline.

Free ZeroGPU has a short queue. If a run times out waiting for a GPU, wait a few seconds and try again.

</div>
        """
    )

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="First frame (optional, enables image-to-video)", type="pil")
            prompt = gr.Textbox(
                label="Prompt",
                info="Describe one shot: subject, motion, camera, lighting, and sound.",
                value=DEFAULT_PROMPT,
                lines=5,
                placeholder="A cinematic tracking shot...",
            )
            with gr.Row():
                duration = gr.Slider(
                    label="Duration (seconds)",
                    minimum=1.0,
                    maximum=5.0,
                    value=2.0,
                    step=0.1,
                )
                with gr.Column():
                    enhance_prompt = gr.Checkbox(
                        label="Enhance prompt",
                        value=False,
                        info="Not used in this demo (LTX-2.3 needs a separate enhancer).",
                    )
                    high_res = gr.Checkbox(label="Higher resolution", value=False)

            generate_btn = gr.Button("Generate video", variant="primary", size="lg")

            with gr.Accordion("Advanced", open=False):
                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, value=42, step=1)
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                with gr.Row():
                    width = gr.Number(label="Width", value=640, precision=0)
                    height = gr.Number(label="Height", value=384, precision=0)

        with gr.Column():
            output_video = gr.Video(label="Generated video", autoplay=True)

    gr.Examples(
        examples=[
            [DEFAULT_PROMPT],
            [
                "A tiny moss-covered clockwork fox trots through a rain-soaked neon market at night. "
                "Low tracking shot, wet asphalt reflections, steam from food stalls, distant bass and rainfall."
            ],
            [
                "Macro shot of a glass terrarium at dawn. A brass hummingbird unfolds its wings, "
                "dew glints on ferns, soft mechanical clicks and distant birdsong."
            ],
        ],
        inputs=[prompt],
        cache_examples=False,
    )

    input_image.change(fn=on_image_upload, inputs=[input_image, high_res], outputs=[width, height])
    high_res.change(fn=on_highres_toggle, inputs=[input_image, high_res], outputs=[width, height])
    generate_btn.click(
        fn=generate_video,
        inputs=[
            input_image,
            prompt,
            duration,
            enhance_prompt,
            seed,
            randomize_seed,
            height,
            width,
        ],
        outputs=[output_video, seed],
        api_name="generate",
    )

demo.queue(default_concurrency_limit=1).launch(
    mcp_server=True,
    theme=gr.themes.Citrus(),
    css=CSS,
)