import os

os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")

import math
import random
import tempfile

import numpy as np
import spaces
import torch
import gradio as gr
from PIL import Image, ImageOps
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

from diffusers import LTX2InContextPipeline
from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition
from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES
from diffusers.utils import encode_video

# --- Config -----------------------------------------------------------------
# FAST distilled variant of the ingredients (reference-sheet) IC-LoRA: 8-step schedule, CFG off.
BASE_MODEL = "diffusers/LTX-2.3-Distilled-Diffusers"
LORA_REPO = "linoyts/LTX-2.3-loras"
LORA_FILE = "ltx-2.3-22b-ic-lora-ingredients-0.9"  # no .safetensors extension in the repo
LORA_SCALE = 1.4
FPS = 24
WIDTH, HEIGHT = 768, 448
NUM_FRAMES = 121
NUM_STEPS = len(DISTILLED_SIGMA_VALUES)  # 8
MAX_SEED = np.iinfo(np.int32).max
HF_TOKEN = os.environ.get("HF_TOKEN")

pipe = LTX2InContextPipeline.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16)
pipe.to("cuda")
pipe.vae.enable_tiling()
_lora_path = hf_hub_download(LORA_REPO, LORA_FILE, token=HF_TOKEN)
pipe.load_lora_weights(load_file(_lora_path), adapter_name="ingredients")
pipe.set_adapters("ingredients", LORA_SCALE)


def _gallery_paths(gallery):
    paths = []
    for item in gallery or []:
        if isinstance(item, (list, tuple)):
            item = item[0]
        if isinstance(item, dict):
            item = item.get("path") or item.get("name") or item.get("image")
        if isinstance(item, str):
            paths.append(item)
    return paths


def compose_sheet(gallery):
    paths = _gallery_paths(gallery)
    if not paths:
        raise gr.Error("Upload at least one subject image to build a sheet.")
    imgs = [Image.open(p).convert("RGB") for p in paths]
    if len(imgs) == 1:
        return imgs[0]
    CW, CH = 1536, 896
    canvas = Image.new("RGB", (CW, CH), (0, 0, 0))
    n = len(imgs)
    cols = math.ceil(math.sqrt(n))
    rows = math.ceil(n / cols)
    g = 16
    cw = (CW - g * (cols + 1)) // cols
    ch = (CH - g * (rows + 1)) // rows
    for i, im in enumerate(imgs):
        r, c = divmod(i, cols)
        canvas.paste(ImageOps.fit(im, (cw, ch), Image.LANCZOS), (g + c * (cw + g), g + r * (ch + g)))
    return canvas


def _build_prompt(sheet, action):
    return f"Reference sheet: {sheet.strip()}\n\nGenerated video: {action.strip()}"


def _export(video_np, audio, path):
    kw = {}
    if audio is not None:
        kw = dict(audio=audio[0].float().cpu(), audio_sample_rate=pipe.vocoder.config.output_sampling_rate)
    encode_video(video_np, fps=FPS, output_path=path, **kw)


def _duration(*args, **kwargs):
    return 200


@spaces.GPU(duration=_duration)
def generate(sheet_image, sheet, action, lora_scale, seed, randomize, progress=gr.Progress()):
    if sheet_image is None:
        raise gr.Error("Add a reference sheet (upload one, or build one from subject images in the other tab).")
    if not sheet.strip():
        raise gr.Error("Describe the elements in the reference sheet (characters, props, location).")
    if not action.strip():
        raise gr.Error("Describe the action / shot you want generated.")

    if randomize:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)

    sheet_img = sheet_image.convert("RGB").resize((WIDTH, HEIGHT), Image.LANCZOS)
    ref = [sheet_img] * NUM_FRAMES
    pipe.set_adapters("ingredients", float(lora_scale))
    prompt = _build_prompt(sheet, action)

    def _cb(p, i, t, kw):
        progress((i + 1) / NUM_STEPS, desc=f"Generating — step {i + 1}/{NUM_STEPS}")
        return {}

    video_out, audio_out = pipe(
        prompt=prompt, negative_prompt="",
        reference_conditions=[LTX2ReferenceCondition(frames=ref, strength=1.0)],
        reference_downscale_factor=1,
        width=WIDTH, height=HEIGHT, num_frames=NUM_FRAMES, frame_rate=FPS,
        num_inference_steps=NUM_STEPS, sigmas=DISTILLED_SIGMA_VALUES,
        guidance_scale=1.0, stg_scale=0.0, audio_guidance_scale=1.0, audio_stg_scale=0.0,
        generator=torch.Generator(device="cuda").manual_seed(seed),
        output_type="np", return_dict=False, callback_on_step_end=_cb,
    )
    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    _export(video_out[0], audio_out, out_path)
    return out_path, seed


with gr.Blocks(title="LTX-2.3 Ingredients (Fast)") as demo:
    gr.Markdown(
        "# ⚡ LTX-2.3 Ingredients — Fast (Distilled)\n"
        "Reference-sheet control, fast. Upload a ready sheet, or build one from individual subject images. Using "
        "[LTX 2.3 Distilled](https://huggingface.co/diffusers/LTX-2.3-Distilled-Diffusers) with the "
        "[Ingredients IC-LoRA](https://huggingface.co/ltx-community/LTX-2.3-loras), via diffusers 🧨. "
        "(For maximum fidelity, see the Dev demo.)"
    )
    with gr.Row():
        with gr.Column():
            with gr.Tabs():
                with gr.Tab("Reference sheet"):
                    sheet_image = gr.Image(type="pil", label="Reference sheet (composite of characters / props / location)")
                with gr.Tab("Build from subject images"):
                    gallery = gr.Gallery(label="Upload subject images (characters, props, location)",
                                         type="filepath", interactive=True, columns=4, height=240)
                    build_btn = gr.Button("Build reference sheet ➜")
                    gr.Markdown("*Tiles your images into one sheet and loads it into the **Reference sheet** tab.*")
            sheet = gr.Textbox(label="Reference sheet description", lines=3,
                               placeholder="a young woman with red hair in a green jacket (face close-up + turnaround); a brass pocket watch; a cobblestone alley at night")
            action = gr.Textbox(label="Generated video — action / shot, plus any speech & sounds", lines=3,
                                placeholder="the woman walks down the alley, checks the pocket watch and whispers 'almost time'; footsteps on cobblestone, distant city hum")
            with gr.Accordion("Settings", open=False):
                lora_scale = gr.Slider(0.8, 1.8, value=1.4, step=0.05, label="LoRA strength")
                randomize = gr.Checkbox(True, label="Randomize seed")
                seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
            run = gr.Button("Generate (fast)", variant="primary")
        with gr.Column():
            video_out = gr.Video(label="Generated video")
            used_seed = gr.Number(label="Seed used", interactive=False)

    build_btn.click(compose_sheet, inputs=gallery, outputs=sheet_image)
    run.click(generate, inputs=[sheet_image, sheet, action, lora_scale, seed, randomize], outputs=[video_out, used_seed])

    gr.Examples(
        examples=[
            ["examples/sheet_garden.png",
             "a cartoon hedgehog (face close-up and body turnaround) and a cartoon rabbit (turnaround); a green coiled garden hose reel and green spray bottles; the interior of a 'Greenfield Home & Garden' store with shelves of plants",
             "the hedgehog waddles up to the camera in the Greenfield Home & Garden store and says cheerfully 'Welcome to Greenfield!', while the rabbit hops past holding a green spray bottle; upbeat store music, the hedgehog's friendly voice and soft footsteps",
             1.4, 42, False],
            ["examples/sheet_hiker.png",
             "a young woman hiker in a green shirt and khaki shorts (face close-up and body turnaround); a large blue hiking backpack; a wooden walking stick; a shaggy yak with a colorful woven saddle blanket; a Himalayan stone village with prayer flags and snowy mountains",
             "the woman loads the blue backpack onto the yak, pats its neck and says warmly 'almost at the summit, buddy', snowy peaks and a monastery behind her; wind, distant prayer bells, her voice and the yak's low grunt",
             1.4, 42, False],
            ["examples/subj_composite.png",
             "a smiling young woman with curly dark hair; a dappled grey horse; a green misty mountain meadow",
             "the woman walks up to the grey horse in the misty meadow, gently strokes its neck and says softly 'good boy, easy now'; a soft horse nicker, her gentle voice, light wind and birdsong",
             1.4, 42, False],
        ],
        inputs=[sheet_image, sheet, action, lora_scale, seed, randomize],
        outputs=[video_out, used_seed], fn=generate, cache_examples=True, cache_mode="lazy",
    )

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