import os
import re
import tempfile
import time

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

import spaces
import torch
import gradio as gr
from diffusers import LTX2Pipeline
from diffusers.pipelines.ltx2.export_utils import encode_video
from diffusers.pipelines.ltx2.utils import DEFAULT_NEGATIVE_PROMPT


MODEL_ID = "SulphurAI/Sulphur-2-base"
BASE_MODEL_ID = "diffusers/LTX-2.3-Diffusers"
LORA_FILE = "sulphur_lora_rank_768.safetensors"

WIDTH = 512
HEIGHT = 320
NUM_FRAMES = 49
FPS = 24.0

# A deliberately conservative public-demo guard. The upstream checkpoint is
# described as uncensored, but a public demo should not generate abuse content.
BLOCKED_PATTERNS = (
    r"\b(child|children|kid|minor|underage|teen(?:ager)?)\b.{0,50}"
    r"\b(nude|naked|sex|sexual|explicit|porn)\b",
    r"\b(nude|naked|sex|sexual|explicit|porn)\b.{0,50}"
    r"\b(child|children|kid|minor|underage|teen(?:ager)?)\b",
    r"\b(rape|sexual assault|non[- ]consensual|revenge porn|csam)\b",
    r"\b(gore|dismemberment|beheading|graphic violence)\b",
)


def _allowed(prompt: str) -> bool:
    text = prompt.casefold()
    return not any(re.search(pattern, text) for pattern in BLOCKED_PATTERNS)


pipe = LTX2Pipeline.from_pretrained(
    BASE_MODEL_ID,
    torch_dtype=torch.bfloat16,
)
pipe.load_lora_weights(MODEL_ID, weight_name=LORA_FILE)
pipe.to("cuda")
pipe.vae.enable_tiling()


def _duration(prompt: str, seed: int, steps: int, *args, **kwargs) -> int:
    del prompt, seed, args, kwargs
    return min(300, 90 + int(steps) * 9)


@spaces.GPU(duration=_duration, size="xlarge")
def generate(prompt: str, seed: int, steps: int) -> tuple[str, str]:
    """Generate a short 512×320 video with synchronized audio from a text prompt."""
    prompt = (prompt or "").strip()
    if len(prompt) < 8:
        raise gr.Error("Please enter a more descriptive prompt.")
    if len(prompt) > 1_500:
        raise gr.Error("Please keep the prompt under 1,500 characters.")
    if not _allowed(prompt):
        raise gr.Error(
            "This public demo cannot process sexual, exploitative, or graphic-violence prompts."
        )

    started = time.perf_counter()
    generator = torch.Generator(device="cuda").manual_seed(int(seed))
    video, audio = pipe(
        prompt=prompt,
        negative_prompt=DEFAULT_NEGATIVE_PROMPT,
        width=WIDTH,
        height=HEIGHT,
        num_frames=NUM_FRAMES,
        frame_rate=FPS,
        num_inference_steps=int(steps),
        guidance_scale=3.0,
        generator=generator,
        output_type="np",
        return_dict=False,
    )

    output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    output.close()
    encode_video(
        video[0],
        fps=FPS,
        audio=audio[0].float().cpu(),
        audio_sample_rate=pipe.vocoder.config.output_sampling_rate,
        output_path=output.name,
    )
    elapsed = time.perf_counter() - started
    return output.name, f"Finished in {elapsed:.1f}s · seed {int(seed)}"


CSS = """
.gradio-container { max-width: 1120px !important; }
.hero { text-align: center; margin: 1.5rem auto 1rem; }
.hero h1 { font-size: clamp(2rem, 5vw, 4rem); margin-bottom: .2rem; }
.hero p { color: #a1a1aa; font-size: 1.05rem; }
"""

with gr.Blocks(css=CSS, theme=gr.themes.Soft(primary_hue="purple")) as demo:
    gr.HTML(
        """
        <div class="hero">
          <h1>🎬 Sulphur 2 Base</h1>
          <p>Text-to-video with synchronized audio, powered by LTX 2.3.</p>
        </div>
        """
    )
    with gr.Row():
        with gr.Column(scale=5):
            prompt = gr.Textbox(
                label="Describe your shot",
                placeholder=(
                    "A cinematic tracking shot of a tiny moss-covered robot "
                    "walking through a rain-soaked neon market..."
                ),
                lines=7,
                max_lines=12,
            )
            with gr.Row():
                seed = gr.Number(label="Seed", value=42, precision=0)
                steps = gr.Slider(
                    label="Inference steps", minimum=12, maximum=30, value=20, step=1
                )
            run = gr.Button("Generate video", variant="primary", size="lg")
            gr.Markdown(
                "Public demo guardrails apply. Avoid sexual, exploitative, "
                "graphic, or deceptive content."
            )
        with gr.Column(scale=7):
            video = gr.Video(label="Generated clip", autoplay=True)
            status = gr.Markdown()

    gr.Examples(
        examples=[
            [
                "A macro cinematic shot of a glass terrarium at dawn. A tiny "
                "clockwork hummingbird unfolds its brass wings, dew glints on "
                "fern leaves, soft mechanical clicks and distant birdsong."
            ],
            [
                "A wide aerial shot over black volcanic sand at blue hour. "
                "Bioluminescent waves roll ashore under a star-filled sky, with "
                "wind and gentle surf in the soundtrack."
            ],
            [
                "Stop-motion style: a paper astronaut plants a small sunflower "
                "on a handcrafted moon, warm studio lighting, subtle paper "
                "rustling and a whimsical music-box melody."
            ],
        ],
        inputs=[prompt],
        cache_examples=False,
    )
    run.click(
        fn=generate,
        inputs=[prompt, seed, steps],
        outputs=[video, status],
        api_name="generate",
    )

demo.queue(default_concurrency_limit=1).launch(mcp_server=True)