import os
import random
import tempfile
import uuid

import numpy as np
import spaces
import torch
import gradio as gr
from diffusers import LTXConditionPipeline
from diffusers.utils import export_to_video

# ---------------------------------------------------------------------------
# Fix a real bug in gradio_client 1.5.2 (bundled with ZeroGPU's forced
# gradio==5.9.1): its JSON-schema-to-python-type parser assumes every schema
# node is a dict, but valid JSON Schema allows a bare boolean (e.g.
# additionalProperties: true) — that crashes EVERY API call to this Space
# with "TypeError: argument of type 'bool' is not iterable". Patched here,
# not worked around with an old pydantic pin (which broke other packages'
# serialization elsewhere in this account's other Spaces) — this is the
# actual bug, proven across studio-chat and studio-media this session.
import gradio_client.utils as _gc_utils

_orig_json_schema_to_python_type = _gc_utils._json_schema_to_python_type


def _patched_json_schema_to_python_type(schema, defs):
    if not isinstance(schema, dict):
        return "bool" if isinstance(schema, bool) else "Any"
    return _orig_json_schema_to_python_type(schema, defs)


_gc_utils._json_schema_to_python_type = _patched_json_schema_to_python_type

# ---------------------------------------------------------------------------
# Model — loaded once at import time. Under ZeroGPU's patched torch,
# .to("cuda") here is deferred (no real GPU touched yet); the ZeroGPU
# "tensors packing" step prepares it for fast real materialization inside
# the @spaces.GPU-decorated call.
# ---------------------------------------------------------------------------

MODEL_ID = "Lightricks/LTX-Video-0.9.7-distilled"
RESOLUTIONS = {
    "704x480 (landscape, fastest)": (704, 480),
    "480x704 (portrait, fastest)": (480, 704),
    "832x480 (landscape)": (832, 480),
    "480x832 (portrait)": (480, 832),
}

pipe = LTXConditionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
pipe.to("cuda")
pipe.vae.enable_tiling()


def _video_duration(prompt, negative_prompt, seed, random_seed, num_frames, steps, resolution):
    # UNMEASURED estimate ahead of the first real call — the distilled
    # checkpoint's own model card recommends ~7 steps for fast generation,
    # so this starts conservative and gets recalibrated against real
    # wall-clock time once deployed, the same discipline used for
    # studio-media's FLUX.1-dev steps slider.
    return min(120, 25 + int(num_frames) // 8 * int(steps) * 1.5)


@spaces.GPU(duration=_video_duration)
def generate_video(prompt, negative_prompt, seed, random_seed, num_frames, steps, resolution):
    if not prompt or not prompt.strip():
        raise gr.Error("Enter a prompt first.")

    width, height = RESOLUTIONS.get(resolution, (704, 480))
    num_frames = int(num_frames)
    steps = int(steps)

    s = random.randint(0, 2**31 - 1) if random_seed else int(seed)
    generator = torch.Generator("cuda").manual_seed(s)

    use_negative = bool(negative_prompt and negative_prompt.strip())
    result = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt if use_negative else None,
        width=width,
        height=height,
        num_frames=num_frames,
        num_inference_steps=steps,
        generator=generator,
        output_type="pil",
    )
    frames = result.frames[0]

    out_path = os.path.join(tempfile.gettempdir(), f"ltxvideo_{uuid.uuid4().hex}.mp4")
    export_to_video(frames, out_path, fps=24)
    return out_path, str(s)


with gr.Blocks(title="Beexly Studio Video (ZeroGPU)") as demo:
    gr.Markdown(
        "# Beexly Studio Video (ZeroGPU)\n\n"
        "Personal video-generation tool. **Not connected to any Galaxy "
        "Sports Edge product, data, or prediction engine.**\n\n"
        "`LTX-Video-0.9.7-distilled` (Lightricks) under the LTX-Video Open "
        "Weights License — free for personal/non-commercial use; commercial "
        "use requires a separate paid agreement above a revenue threshold. "
        "This stays a personal tool, same posture as the non-commercial "
        "FLUX.1-dev model on studio-media: **never for anything shipped on "
        "galaxysportsedge.com.**"
    )

    with gr.Tab("Text to Video"):
        v_prompt = gr.Textbox(label="Prompt", lines=2)
        v_negative = gr.Textbox(label="Negative prompt (optional)", value="")
        v_resolution = gr.Dropdown(list(RESOLUTIONS.keys()), value="704x480 (landscape, fastest)", label="Resolution")
        with gr.Row():
            v_seed = gr.Textbox(label="Seed", value="0")
            v_random_seed = gr.Checkbox(label="Random seed", value=True)
        v_frames = gr.Slider(41, 121, value=65, step=8, label="Frames (num_frames must be 8n+1; ~24fps, so 65 frames ≈ 2.7s)")
        v_steps = gr.Slider(4, 12, value=7, step=1, label="Steps (7 is the distilled checkpoint's recommended default)")
        v_btn = gr.Button("Generate", variant="primary")
        v_out = gr.Video(label="Generated video")
        v_seed_used = gr.Textbox(label="Seed used", interactive=False)
        v_btn.click(
            generate_video,
            inputs=[v_prompt, v_negative, v_seed, v_random_seed, v_frames, v_steps, v_resolution],
            outputs=[v_out, v_seed_used],
            api_name="generate_video",
        )

if __name__ == "__main__":
    demo.launch()