import os

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

import random
import tempfile
import threading
import time

import numpy as np
import imageio.v3 as iio
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 load_video, encode_video

# --- Config -----------------------------------------------------------------
BASE_MODEL = "diffusers/LTX-2.3-Distilled-Diffusers"
LORA_REPO = "linoyts/LTX-2.3-loras"
LORA_FILES = {2: "ltx-2.3-22b-ic-lora-spatial-upscaler-x2s-0.9.safetensors",
              4: "ltx-2.3-22b-ic-lora-spatial-upscaler-x4s-0.9.safetensors"}
LORA_SCALE = 1.0
FPS = 24
NUM_STEPS = len(DISTILLED_SIGMA_VALUES)
MAX_SEED = np.iinfo(np.int32).max
HF_TOKEN = os.environ.get("HF_TOKEN")

# Output resolution per factor (ref = output // factor; both divisible by 32). Capped to keep 4x fast.
OUT_DIMS = {2: (768, 448), 4: (896, 512)}
FRAME_CHOICES = [25, 49, 73]

pipe = LTX2InContextPipeline.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16)
pipe.to("cuda")
pipe.vae.enable_tiling()
_LORA_PATHS = {f: hf_hub_download(LORA_REPO, fn, token=HF_TOKEN) for f, fn in LORA_FILES.items()}
# AOTI: fuse the default factor's LoRA + load precompiled blocks at ROOT level (fused &
# unloaded so weight names match the compiled constants). Switching factor swaps a whole
# LoRA file, so undo the current one (fuse -scale) and fuse the new one (+scale), unloaded
# and additive throughout, then re-snapshot.
_FUSED_FACTOR = sorted(_LORA_PATHS)[0]
pipe.load_lora_weights(load_file(_LORA_PATHS[_FUSED_FACTOR]), adapter_name="up")
pipe.fuse_lora(lora_scale=LORA_SCALE)
pipe.unload_lora_weights()
spaces.aoti_load(module=pipe.transformer, repo_id="ltx-community/LTX-2.3-Transformer-GroupA-sm120-cu130-r9e")
def _prepare(factor):
    global _FUSED_FACTOR
    if factor == _FUSED_FACTOR:
        return
    pipe.load_lora_weights(load_file(_LORA_PATHS[_FUSED_FACTOR]), adapter_name="up")
    pipe.fuse_lora(lora_scale=-LORA_SCALE)
    pipe.unload_lora_weights()
    pipe.load_lora_weights(load_file(_LORA_PATHS[factor]), adapter_name="up")
    pipe.fuse_lora(lora_scale=LORA_SCALE)
    pipe.unload_lora_weights()
    spaces.aoti_load(module=pipe.transformer, repo_id="ltx-community/LTX-2.3-Transformer-GroupA-sm120-cu130-r9e")
    _FUSED_FACTOR = factor


def _src_fps(path, default=FPS):
    try:
        return float(iio.immeta(path, plugin="pyav").get("fps", default)) or default
    except Exception:
        return default


def _load_frames(path, num_frames, width, height):
    frames = load_video(path)
    if not frames:
        return []
    fps = _src_fps(path)
    out = []
    for i in range(num_frames):
        idx = min(int(round(i / FPS * fps)), len(frames) - 1)
        out.append(ImageOps.fit(frames[idx].convert("RGB"), (width, height), Image.LANCZOS))
    return out


def _out_dims(first_frame, factor):
    w, h = OUT_DIMS[factor]
    if first_frame.height > first_frame.width:
        w, h = h, w
    return w, h


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):
    factor = next((int(str(a)[0]) for a in args if str(a) in ("2", "4")), 2)
    num_frames = next((a for a in args if a in FRAME_CHOICES), 49)
    return int(70 + int(num_frames) * (2.2 if factor == 4 else 1.4))


@spaces.GPU(duration=_duration)
def upscale(video, factor, scene, num_frames, seed, randomize,
            progress=gr.Progress(track_tqdm=True)):
    if video is None:
        raise gr.Error("Please upload a low-resolution video.")
    factor = int(factor)
    if randomize:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)
    num_frames = int(num_frames)

    probe = load_video(video)
    if not probe:
        raise gr.Error("Could not read any frames from that video.")
    out_w, out_h = _out_dims(probe[0], factor)
    ref_w, ref_h = out_w // factor, out_h // factor

    ref = _load_frames(video, num_frames, ref_w, ref_h)
    _prepare(factor)
    desc = scene.strip() or "the scene"
    prompt = (f"{desc}, high resolution, sharp fine detail, crisp clean edges, "
              f"restored texture; identical content, composition and motion to the reference, only resolution and detail increase.")

    def _cb(p, i, t, kw):
        progress((i + 1) / NUM_STEPS, desc=f"Upscaling {factor}× — 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=factor,
        width=out_w, height=out_h, 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,
        )
    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 Spatial Upscaler") as demo:
    gr.Markdown(
        "# 🔬 LTX-2.3 Spatial Upscaler\n"
        "Upscale a low-resolution video 2× or 4×, generating sharp detail while keeping content, composition "
        "and motion. Using [LTX 2.3 Distilled](https://huggingface.co/diffusers/LTX-2.3-Distilled-Diffusers) "
        "with the [Spatial Upscaler IC-LoRA](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Pixel-Spatial-Upscaler), via diffusers 🧨."
    )
    with gr.Row():
        with gr.Column():
            video_in = gr.Video(label="Low-resolution video")
            factor = gr.Radio(["2", "4"], value="2", label="Upscale factor (4× is heavier)")
            scene = gr.Textbox(label="Prompt — describe the scene and any sounds (optional)", lines=2,
                               placeholder="a close-up of a butterfly on a flower in a sunny garden; soft garden ambience and a gentle breeze")
            with gr.Accordion("Settings", open=False):
                num_frames = gr.Dropdown(FRAME_CHOICES, value=49, label="Frames (24fps)")
                randomize = gr.Checkbox(True, label="Randomize seed")
                seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
            run = gr.Button("Upscale", variant="primary")
        with gr.Column():
            video_out = gr.Video(label="Upscaled result")

    run.click(upscale, inputs=[video_in, factor, scene, num_frames, seed, randomize],
              outputs=[video_out, seed])

    gr.Examples(
        examples=[
            ["examples/lowres_cat.mp4", "2", "a tabby cat in crisp high resolution — individual strands of brown-and-grey striped fur, fine whiskers, bright detailed eyes and a moist little nose, clean sharp edges throughout; soft contented purring", 49, 42, False],
            ["examples/lowres_landscape.mp4", "4", "a misty mountain landscape in crisp high resolution — fine detail in the pine trees and rocky slopes, delicate layers of drifting fog, sharp rippling reflections on the water and a clean crisp sky; a gentle wind and distant birdsong", 49, 42, False],
        ],
        inputs=[video_in, factor, scene, num_frames, seed, randomize],
        outputs=[video_out, seed], fn=upscale, cache_examples=True, cache_mode="lazy",
    )

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