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 _compose_sheet(paths):
    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_sheet_preview(files):
    if not files:
        return None
    paths = [f if isinstance(f, str) else f.get("path", f.get("name")) for f in files]
    return _compose_sheet(paths)


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(files, sheet, action, lora_scale, seed, randomize, progress=gr.Progress(track_tqdm=True)):
    if not files:
        raise gr.Error("Upload a reference sheet image, or several subject images to build one.")
    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)

    paths = [f if isinstance(f, str) else f.get("path", f.get("name")) for f in files]
    sheet_img = _compose_sheet(paths).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 on the **distilled** checkpoint (8-step, fast). **Upload a ready reference "
        "sheet, or several subject images and we'll tile them into one.** Describe the sheet and the action "
        "(with speech/sounds for audio). For maximum fidelity use the non-distilled demo. "
        "IC-LoRA: [`linoyts/LTX-2.3-loras`](https://huggingface.co/linoyts/LTX-2.3-loras)."
    )
    with gr.Row():
        with gr.Column():
            files = gr.File(label="Reference sheet (1 image) or subject images (several)",
                            file_count="multiple", file_types=["image"], type="filepath")
            sheet_preview = gr.Image(label="Reference sheet used", type="pil", interactive=False)
            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 — the action / shot, plus any speech & sounds", lines=3,
                                placeholder="the woman walks down the alley and checks the pocket watch, slow dolly-in; footsteps on cobblestone, a soft voice saying 'almost time', 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)

    files.change(build_sheet_preview, inputs=files, outputs=sheet_preview)
    run.click(generate, inputs=[files, 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 and the rabbit explore the Greenfield Home & Garden store among the plants and garden tools, the rabbit holding a green spray bottle, warm bright store lighting, playful slow camera; cheerful ambient store sounds 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 in front of snowy Himalayan peaks and a monastery, gentle handheld camera, soft daylight; wind, distant prayer bells and the yak's low grunt",
             1.4, 42, False],
            [["examples/subj_woman.jpg", "examples/subj_horse.jpg", "examples/subj_landscape.jpg"],
             "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 and gently strokes its neck, soft daylight; gentle wind, a soft horse nicker and distant birdsong",
             1.4, 42, False],
        ],
        inputs=[files, 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)