import os
import subprocess
import sys

# ZeroGPU: torch.compile / dynamo unsupported — disable before any torch import.
os.environ["TORCH_COMPILE_DISABLE"] = "1"
os.environ["TORCHDYNAMO_DISABLE"] = "1"
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

# --- clone + install the NATIVE LTX-2 codebase at a pinned commit ---
LTX_REPO_URL = "https://github.com/Lightricks/LTX-2.git"
LTX_REPO_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "LTX-2")
LTX_COMMIT = "780984275fd47128b02bef9b5c085404276866ee"
if not os.path.exists(LTX_REPO_DIR):
    subprocess.run(["git", "clone", LTX_REPO_URL, LTX_REPO_DIR], check=True)
    subprocess.run(["git", "-C", LTX_REPO_DIR, "checkout", LTX_COMMIT], check=True)
subprocess.run([sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps",
                "-e", os.path.join(LTX_REPO_DIR, "packages", "ltx-core"),
                "-e", os.path.join(LTX_REPO_DIR, "packages", "ltx-pipelines")], check=True)
sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-pipelines", "src"))
sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-core", "src"))

import logging
import random
import tempfile
import traceback

import torch
import imageio_ffmpeg

torch._dynamo.config.suppress_errors = True
torch._dynamo.config.disable = True

import spaces
import gradio as gr
from huggingface_hub import hf_hub_download, snapshot_download

# Import LTX modules — order matters (model modules first, then quantization/loader).
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number  # noqa: F401
from ltx_core.loader import LoraPathStrengthAndSDOps, LTXV_LORA_COMFY_RENAMING_MAP
from ltx_pipelines.retake import RetakePipeline
from ltx_pipelines.utils.constants import detect_params
from ltx_core.components.guiders import MultiModalGuiderParams
from ltx_pipelines.utils.media_io import encode_video, get_videostream_metadata

# --- ZeroGPU loader patch ---
# The native loader opens safetensors directly on CUDA (safe_open(path, device="cuda")),
# bypassing torch.Tensor.to (the call ZeroGPU patches). Patch to open on CPU then
# move via torch.Tensor.to (ZeroGPU-virtualisable).
import safetensors as _safetensors
import ltx_core.loader.sft_loader as _sft
from ltx_core.loader.primitives import StateDict as _StateDict


def _zerogpu_safe_load(self, path, sd_ops, device=None):
    device = device or torch.device("cpu")
    sd, size, dtype = {}, 0, set()
    model_paths = path if isinstance(path, list) else [path]
    for shard_path in model_paths:
        with _safetensors.safe_open(shard_path, framework="pt", device="cpu") as f:
            for name in f.keys():
                expected = name if sd_ops is None else sd_ops.apply_to_key(name)
                if expected is None:
                    continue
                value = f.get_tensor(name).to(device=device)
                kvs = ((expected, value),)
                if sd_ops is not None:
                    kvs = sd_ops.apply_to_key_value(expected, value)
                for k, v in kvs:
                    size += v.nbytes
                    dtype.add(v.dtype)
                    sd[k] = v
    return _StateDict(sd=sd, device=device, size=size, dtype=dtype)


_sft.SafetensorsStateDictLoader.load = _zerogpu_safe_load
print("[PATCH] safetensors loader -> CPU-open + torch.to (ZeroGPU-virtualisable)")

# --- attention backend patch (FA3 crashes on Blackwell ZeroGPU; use SDPA) ---
import torch.nn.functional as F
from ltx_core.model.transformer import attention as _attn_mod


def _sdpa_as_mea(query, key, value, attn_bias=None, scale=None, **kwargs):
    q, k, v = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
    return F.scaled_dot_product_attention(q, k, v, scale=scale).transpose(1, 2)


_attn_mod.memory_efficient_attention = _sdpa_as_mea
print("[ATTN] SDPA (patched at module scope, no CUDA query)")

logging.getLogger().setLevel(logging.INFO)
print("[VERSION] ltx-2.3-foley-lora v1", flush=True)

# =========================== CONFIG ===========================
TITLE = "LTX-2.3 Foley LoRA"
LTX_MODEL_REPO = "Lightricks/LTX-2.3"
BASE_CKPT = "ltx-2.3-22b-dev.safetensors"
GEMMA_REPO = "google/gemma-3-12b-it-qat-q4_0-unquantized"
LORA_REPO = "FuzzPuppy/LTX-2.3-Foley-LoRA"
LORA_FILE = "ltx-2.3-foley-400-steps.safetensors"
LORA_SCALE = 1.5  # card says 1-3; start at 1.5 for strong Foley
NUM_STEPS = 30  # LTX-2.3 non-distilled default
MAX_SEED = 2**32 - 1
HF_TOKEN = os.environ.get("HF_TOKEN")

# Recommended prompts from the model card
DEFAULT_PROMPT_SUFFIX = "No speech is present. No music is present."
DEFAULT_NEGATIVE = (
    "music, melody, song, singing, vocals, score, soundtrack, beat, "
    "rhythm bed, instrumental backing, tinny, thin, harsh, clipped, "
    "distorted, low bitrate"
)


def _duration(*args, **kwargs):
    return 600  # V2A: load full 22B transformer + Gemma + audio/video encoders/decoders


# --- Build the RetakePipeline once at module scope ---
# Download checkpoints to a local_dir (real files) inside the app dir so the
# in-worker transformer build can read them (ZeroGPU workers can't resolve
# HF cache symlinks for the 46GB checkpoint).
MODELS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
os.makedirs(MODELS_DIR, exist_ok=True)

print("Downloading checkpoints to local_dir...", flush=True)
ckpt_path = hf_hub_download(LTX_MODEL_REPO, BASE_CKPT, token=HF_TOKEN, local_dir=MODELS_DIR)
lora_path = hf_hub_download(LORA_REPO, LORA_FILE, token=HF_TOKEN, local_dir=MODELS_DIR)
# Gemma is pinned at module scope (not read in-worker) -> normal cache is fine.
gemma_root = snapshot_download(GEMMA_REPO, token=HF_TOKEN)

print("Building RetakePipeline (distilled=False)...", flush=True)
pipeline = RetakePipeline(
    checkpoint_path=ckpt_path,
    gemma_root=gemma_root,
    loras=[LoraPathStrengthAndSDOps(lora_path, LORA_SCALE, LTXV_LORA_COMFY_RENAMING_MAP)],
    distilled=False,  # Use full non-distilled model for better quality
)

# Detect params from checkpoint (LTX-2.3 uses stg_blocks=[28])
_params = detect_params(ckpt_path)
print(f"[PARAMS] detected: steps={_params.num_inference_steps}, "
      f"v_stg_blocks={_params.video_guider_params.stg_blocks}, "
      f"a_stg_blocks={_params.audio_guider_params.stg_blocks}", flush=True)

print("Pipeline ready.", flush=True)

# =========================== AoTI acceleration ===========================
# Swap the 47 non-STG transformer blocks for AoTI-compiled graphs (~1.3x faster
# generation). Precompiled graphs live in a dataset; weights stay runtime inputs.
# Fully guarded — any failure falls back to eager so the demo always works.
AOTI_REPO = "multimodalart/ltx-2.3-foley-aoti"
try:
    from dataclasses import fields as _dcf
    from torch.utils._pytree import register_pytree_node, GetAttrKey
    from ltx_core.model.transformer.transformer_args import TransformerArgs as _TA
    from ltx_core.model.transformer.model import LTXModel as _LTXModel
    from ltx_core.loader.module_ops import ModuleOps as _ModuleOps
    from spaces.zero.torch.aoti import ZeroGPUCompiledModel as _ZGCM, ZeroGPUWeights as _ZGW

    _AT = ["x", "context", "context_mask", "timesteps", "embedded_timestep",
           "positional_embeddings", "cross_positional_embeddings", "cross_scale_shift_timestep",
           "cross_gate_timestep", "prompt_timestep", "self_attention_mask",
           "self_attn_perturbation_mask", "cross_attn_perturbation_mask"]
    _AS = ["enabled", "self_attn_all_perturbed", "cross_attn_skip_all"]
    assert {f.name for f in _dcf(_TA)} == set(_AT) | set(_AS)
    register_pytree_node(
        _TA,
        lambda a: ([getattr(a, f) for f in _AT], tuple(getattr(a, f) for f in _AS)),
        lambda ch, ctx: _TA(**dict(zip(_AT, ch)), **dict(zip(_AS, ctx))),
        flatten_with_keys_fn=lambda a: ([(GetAttrKey(f), getattr(a, f)) for f in _AT],
                                        tuple(getattr(a, f) for f in _AS)),
        serialized_type_name="ltx_core.TransformerArgs",
        to_dumpable_context=lambda c: list(c), from_dumpable_context=lambda d: tuple(d))

    def _aoti_sid(v, a):
        def kk(t):
            return "N" if t is None else "".join("1" if getattr(t, f) is None else "0" for f in _AT)
        return f"{kk(v)}_{kk(a)}"

    _AOTI_PKGS = {}
    from huggingface_hub import HfApi as _HfApi
    for _f in _HfApi(token=HF_TOKEN).list_repo_files(AOTI_REPO, repo_type="dataset"):
        if _f.startswith("BasicAVTransformerBlock/pkg_") and _f.endswith(".pt2"):
            _sid = _f.split("pkg_")[1].rsplit(".pt2", 1)[0]
            _AOTI_PKGS[_sid] = hf_hub_download(AOTI_REPO, _f, repo_type="dataset", token=HF_TOKEN)
    if not _AOTI_PKGS:
        raise RuntimeError("no AoTI packages found")
    _STG = set(_params.video_guider_params.stg_blocks)

    def _aoti_wrap(block, orig):
        h = {"mat": False, "cm": {}, "sd": None}
        def fwd(video=None, audio=None):
            k = _aoti_sid(video, audio)
            if not h["mat"]:
                out = orig(video=video, audio=audio)  # eager -> weights materialise
                h["mat"] = True; h["sd"] = block.state_dict()
                return out
            if k not in _AOTI_PKGS:
                return orig(video=video, audio=audio)
            if k not in h["cm"]:
                try:
                    h["cm"][k] = _ZGCM(_AOTI_PKGS[k], _ZGW(h["sd"]))
                except Exception:
                    h["cm"][k] = False
            if h["cm"][k] is False:
                return orig(video=video, audio=audio)
            try:
                return h["cm"][k](video, audio)
            except Exception:
                h["cm"][k] = False
                return orig(video=video, audio=audio)
        return fwd

    def _aoti_patch(model):
        for i, b in enumerate(model.transformer_blocks):
            if i not in _STG:
                b.forward = _aoti_wrap(b, b.forward)
        return model

    _base_ops = pipeline.stage._transformer_builder.module_ops
    pipeline.stage._transformer_builder = pipeline.stage._transformer_builder.with_module_ops(
        (*_base_ops, _ModuleOps(name="aoti", matcher=lambda m: isinstance(m, _LTXModel), mutator=_aoti_patch)))
    print(f"[AoTI] enabled — {len(_AOTI_PKGS)} block variants", flush=True)
except Exception as _aoti_err:
    print(f"[AoTI] disabled ({_aoti_err!r}); running eager", flush=True)


def _validate_video(path):
    """Validate that the video meets LTX-2.3 constraints (divisible by 32, 8k+1 frames).
    If not, use ffmpeg to pad/trim. Returns the path to a compliant video."""
    if path is None:
        raise gr.Error("Please upload a silent video to add Foley sound to.")

    meta = get_videostream_metadata(path)
    print(f"[VIDEO] {meta.width}x{meta.height}, {meta.fps:.2f}fps, {meta.frames} frames")

    needs_fix = False
    target_w = meta.width
    target_h = meta.height
    target_frames = meta.frames

    # Round height/width up to nearest multiple of 32
    if meta.width % 32 != 0:
        target_w = ((meta.width + 31) // 32) * 32
        needs_fix = True
    if meta.height % 32 != 0:
        target_h = ((meta.height + 31) // 32) * 32
        needs_fix = True

    # Snap frames down to 8k+1
    base = 8
    if (meta.frames - 1) % base != 0:
        target_frames = ((meta.frames - 1) // base) * base + 1
        if target_frames < 1:
            target_frames = base + 1  # minimum 9 frames
        needs_fix = True

    if not needs_fix:
        return path

    # Use ffmpeg to pad/trim
    ff = imageio_ffmpeg.get_ffmpeg_exe()
    out = tempfile.mktemp(suffix=".mp4")
    vf_filters = []
    if target_w != meta.width or target_h != meta.height:
        pad_x = (target_w - meta.width) // 2
        pad_y = (target_h - meta.height) // 2
        vf_filters.append(f"pad={target_w}:{target_h}:{pad_x}:{pad_y}:black")

    cmd = [ff, "-y", "-i", path, "-an"]  # strip audio
    if vf_filters:
        cmd.extend(["-vf", ",".join(vf_filters)])
    cmd.extend(["-frames:v", str(target_frames), "-c:v", "libx264", "-crf", "19",
                "-preset", "veryfast", out])
    print(f"[FIX] padding/trimming video: {meta.width}x{meta.height} {meta.frames}f -> "
          f"{target_w}x{target_h} {target_frames}f")
    subprocess.run(cmd, check=True, capture_output=True)
    return out


@spaces.GPU(duration=_duration)
@torch.inference_mode()
def generate(video, prompt, negative, lora_scale, seed, randomize,
             steps, progress=gr.Progress(track_tqdm=True)):
    """Generate synchronized Foley audio for a silent video using LTX-2.3 + Foley LoRA."""
    if video is None:
        raise gr.Error("Please upload a silent video to add Foley sound to.")
    if not prompt.strip():
        raise gr.Error("Please enter a prompt describing the sounds in the video.")

    # Append the recommended suffix if not already present
    full_prompt = prompt.strip()
    if "no speech" not in full_prompt.lower():
        full_prompt += ". " + DEFAULT_PROMPT_SUFFIX

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

    # Validate/fix the video to meet LTX constraints
    video_path = _validate_video(video)

    meta = get_videostream_metadata(video_path)
    end_time = meta.frames / meta.fps
    print(f"[GEN] {meta.width}x{meta.height}, {meta.fps:.2f}fps, {meta.frames} frames, "
          f"duration={end_time:.2f}s, seed={seed}, steps={steps}")

    # Update LoRA scale if different from default
    if abs(lora_scale - LORA_SCALE) > 0.01:
        pipeline.stage._transformer_builder = pipeline.stage._transformer_builder.with_loras(
            (LoraPathStrengthAndSDOps(lora_path, float(lora_scale), LTXV_LORA_COMFY_RENAMING_MAP),)
        )

    tiling = TilingConfig.default()

    # Use detected params for guider, override steps
    v_guider = _params.video_guider_params
    a_guider = _params.audio_guider_params

    try:
        video_iter, audio_out = pipeline(
            video_path=video_path,
            prompt=full_prompt,
            start_time=0.0,
            end_time=end_time,
            seed=seed,
            negative_prompt=(negative or "").strip(),
            num_inference_steps=int(steps),
            video_guider_params=v_guider,
            audio_guider_params=a_guider,
            regenerate_video=False,  # keep original video
            regenerate_audio=True,  # generate new audio
            tiling_config=tiling,
        )

        out_path = tempfile.mktemp(suffix=".mp4")
        encode_video(
            video=video_iter,
            fps=int(meta.fps),
            audio=audio_out,
            output_path=out_path,
            video_chunks_number=get_video_chunks_number(meta.frames, tiling),
        )
        return out_path, seed
    except Exception as e:
        print(traceback.format_exc(), flush=True)
        raise gr.Error(
            "Foley generation failed. Try a shorter video, or adjust the prompt. "
            "(Full traceback in the Space logs)"
        )


CSS = """
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(css=CSS, title=TITLE, theme=gr.themes.Citrus()) as demo:
    gr.Markdown(
        "# \U0001F3A7 LTX-2.3 Foley LoRA\n"
        "Upload a **silent video** and describe the sounds you want — this demo generates "
        "**synchronized Foley and sound effects** that match the visual action, using "
        "[LTX-2.3](https://huggingface.co/Lightricks/LTX-2.3) with the "
        "[Foley LoRA](https://huggingface.co/FuzzPuppy/LTX-2.3-Foley-LoRA).\n\n"
        "**Tips:** Describe the visible action briefly (e.g. *'A door slams shut'*), then let "
        "the model handle the rest. The LoRA suppresses music and emphasizes realistic sound effects."
    )

    with gr.Row():
        with gr.Column():
            video_in = gr.Video(label="Silent video (upload a video without audio)")
            prompt = gr.Textbox(
                label="Prompt — describe the sounds in the video",
                lines=3,
                placeholder="e.g. A barista uses an espresso machine to steam milk",
                value="A door slams shut with a heavy thud",
            )
            with gr.Accordion("Advanced settings", open=False):
                negative = gr.Textbox(
                    label="Negative prompt",
                    value=DEFAULT_NEGATIVE,
                    lines=3,
                )
                lora_scale = gr.Slider(
                    0.5, 3.0, value=LORA_SCALE, step=0.1,
                    label="LoRA scale (higher = stronger Foley, suppresses music)",
                )
                steps = gr.Slider(
                    10, 50, value=NUM_STEPS, step=1,
                    label="Inference steps (more = higher quality, slower)",
                )
                randomize = gr.Checkbox(True, label="Randomize seed")
                seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
            run = gr.Button("Generate Foley audio", variant="primary")
        with gr.Column():
            video_out = gr.Video(label="Video with generated Foley audio")

    run.click(
        generate,
        inputs=[video_in, prompt, negative, lora_scale, seed, randomize, steps],
        outputs=[video_out, seed],
        api_name="generate",
    )

    gr.Examples(
        examples=[
            ["examples/door-close.mp4", "A door slams shut with a heavy thud", DEFAULT_NEGATIVE, 1.5, 42, False, 30],
            ["examples/racecar.mp4", "A race car speeds by with loud engine revving", DEFAULT_NEGATIVE, 1.5, 42, False, 30],
            ["examples/pineapple.mp4", "A knife slices through a pineapple with crisp cutting sounds", DEFAULT_NEGATIVE, 1.5, 42, False, 30],
            ["examples/squash.mp4", "A squash ball hits the wall with sharp impacts", DEFAULT_NEGATIVE, 1.5, 42, False, 30],
        ],
        inputs=[video_in, prompt, negative, lora_scale, seed, randomize, steps],
        outputs=[video_out, seed],
        fn=generate,
        # Cache on the first click of each example (deterministic seed), then serve
        # instantly to every later visitor — key for a ~10-min-per-run V2A model.
        # lazy: never runs at startup (ZeroGPU has no GPU attached then).
        cache_examples=True,
        cache_mode="lazy",
        label="Examples — silent videos to add Foley to (click to run)",
    )

if __name__ == "__main__":
    # mcp_server=True exposes the /generate endpoint as an MCP tool (docstring +
    # type hints on `generate` become its schema). Requires Gradio >= 5.28.
    demo.launch(show_error=True, mcp_server=True)