import os

# ZeroGPU: torch.compile / dynamo are unsupported — disable before torch import.
os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")

import random
import tempfile

import numpy as np
import spaces
import torch
import gradio as gr
from PIL import Image, ImageFilter
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/ltx2.3-inpainting-lora"
LORA_FILE = "ltx-2.3-22b-ic-lora-inpainting.safetensors"
LORA_SCALE = 1.0
FPS = 24
NUM_STEPS = len(DISTILLED_SIGMA_VALUES)  # 8
MASK_FILL = 128  # masked region painted neutral grey in the reference the model fills
MAX_SEED = np.iinfo(np.int32).max
HF_TOKEN = os.environ.get("HF_TOKEN")

RES_PRESETS = {"Fast (768×448)": (768, 448), "Quality (960×544)": (960, 544)}
FRAME_CHOICES = [49, 73, 97, 121]

# --- Load pipeline once at module scope (ZeroGPU registers it) ---------------
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="inpaint")
pipe.set_adapters("inpaint", LORA_SCALE)


# --- Helpers ----------------------------------------------------------------
def _resample(frames, n):
    idx = np.linspace(0, len(frames) - 1, n).round().astype(int)
    return [frames[i] for i in idx]


def _pick_resolution(first_frame: Image.Image, preset: str):
    w, h = RES_PRESETS[preset]
    if first_frame.height > first_frame.width:
        w, h = h, w
    return w, h


def first_frame(video):
    """Populate the mask editor with the uploaded video's first frame."""
    if video is None:
        return None
    frames = load_video(video)
    return np.array(frames[0].convert("RGB")) if frames else None


def _mask_from_editor(editor_value, width, height):
    """Extract a binary mask (H,W) from a gr.ImageEditor value — union of painted layers."""
    if not editor_value:
        return None
    layers = editor_value.get("layers") or []
    bg = editor_value.get("background")
    if bg is None:
        return None
    H0, W0 = np.asarray(bg).shape[:2]
    acc = np.zeros((H0, W0), dtype=bool)
    for layer in layers:
        arr = np.asarray(layer)
        if arr.ndim == 3 and arr.shape[2] == 4:
            acc |= arr[..., 3] > 10
        elif arr.ndim == 3:
            acc |= arr.sum(axis=2) > 10
    if not acc.any():
        return None
    m = Image.fromarray((acc * 255).astype(np.uint8)).resize((width, height), Image.NEAREST)
    return np.array(m) > 127


def _duration(*args, **kwargs):
    preset = args[3] if len(args) > 3 else "Fast"
    num_frames = args[4] if len(args) > 4 else 73
    per_frame = 1.6 if "Quality" in str(preset) else 1.0
    return int(50 + int(num_frames) * per_frame)


# --- Inference --------------------------------------------------------------
@spaces.GPU(duration=_duration)
def inpaint(video, mask_editor, prompt, preset, num_frames, seed, randomize,
            progress=gr.Progress(track_tqdm=True)):
    if video is None:
        raise gr.Error("Please upload a video.")
    if not prompt.strip():
        raise gr.Error("Describe what should fill the masked region.")

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

    frames = load_video(video)
    if not frames:
        raise gr.Error("Could not read any frames from that video.")

    width, height = _pick_resolution(frames[0], preset)
    num_frames = int(num_frames)

    mask = _mask_from_editor(mask_editor, width, height)
    if mask is None:
        raise gr.Error("Draw a mask over the region to inpaint (use the brush on the frame).")

    orig = [np.array(f.convert("RGB").resize((width, height), Image.LANCZOS))
            for f in _resample(frames, num_frames)]

    # Reference = video with the masked region painted neutral grey; the model fills it.
    ref = []
    for fr in orig:
        m = fr.copy()
        m[mask] = MASK_FILL
        ref.append(Image.fromarray(m))

    ref_cond = LTX2ReferenceCondition(frames=ref, strength=1.0)
    video_out, _audio = pipe(
        prompt=prompt,
        negative_prompt="",
        reference_conditions=[ref_cond],
        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,
    )

    # Composite: keep original pixels outside the mask, generated pixels inside (feathered).
    gen = (np.clip(video_out[0], 0, 1) * 255).astype(np.uint8)
    soft = np.array(Image.fromarray((mask * 255).astype(np.uint8)).filter(ImageFilter.GaussianBlur(3))) / 255.0
    soft = soft[None, :, :, None]
    orig_arr = np.stack(orig).astype(np.float32)
    n = min(len(gen), len(orig_arr))
    out = (gen[:n].astype(np.float32) * soft + orig_arr[:n] * (1 - soft)).astype(np.uint8)

    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    encode_video(out, fps=FPS, output_path=out_path)
    return out_path, seed


# --- UI ---------------------------------------------------------------------
with gr.Blocks(title="LTX-2.3 Video Inpainting") as demo:
    gr.Markdown(
        "# 🪄 LTX-2.3 Video Inpainting\n"
        "Mask a region of a video and regenerate it from a prompt, keeping the rest of the frame intact. "
        "Upload a clip, **brush over the area to replace** on the first frame, describe what should appear there. "
        "The mask is applied across all frames. "
        "IC-LoRA: [`linoyts/ltx2.3-inpainting-lora`](https://huggingface.co/linoyts/ltx2.3-inpainting-lora) · "
        "base: distilled LTX-2.3."
    )
    with gr.Row():
        with gr.Column():
            video_in = gr.Video(label="Input video")
            mask_editor = gr.ImageEditor(
                label="Brush the region to inpaint (loads from the video's first frame)",
                type="numpy",
                layers=False,
                brush=gr.Brush(colors=["#ff2d55"], color_mode="fixed"),
            )
            prompt = gr.Textbox(
                label="What should fill the masked region",
                placeholder="a lush green bush with small white flowers",
                lines=2,
            )
            with gr.Accordion("Settings", open=False):
                preset = gr.Dropdown(list(RES_PRESETS), value="Fast (768×448)", label="Resolution")
                num_frames = gr.Dropdown(FRAME_CHOICES, value=73, 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("Inpaint", variant="primary")
        with gr.Column():
            video_out = gr.Video(label="Inpainted result")
            used_seed = gr.Number(label="Seed used", interactive=False)

    video_in.change(first_frame, inputs=video_in, outputs=mask_editor)
    run.click(
        inpaint,
        inputs=[video_in, mask_editor, prompt, preset, num_frames, seed, randomize],
        outputs=[video_out, used_seed],
    )

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