"""LTX-2.3 dev two-stage T2V backend (API Space for the Pre-Production Playground frontend).

TI2VidTwoStagesPipeline: dev checkpoint stage 1 with CFG/STG guidance, distilled-LoRA
stage 2 + x2 spatial upsampler. Text encoding (positive AND negative, for CFG) is
offloaded to the Gemma encoder Space (TEXT_ENCODER_SPACE) — encode_prompts is
monkeypatched with the precomputed [ctx_p, ctx_n]. Same ZeroGPU setup as the distilled
backend.
"""

import os
import subprocess
import sys

os.environ["TORCH_COMPILE_DISABLE"] = "1"
os.environ["TORCHDYNAMO_DISABLE"] = "1"

subprocess.run([sys.executable, "-m", "pip", "install", "xformers==0.0.32.post2",
                "--no-build-isolation"], check=False)

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_SHA = os.environ.get("LTX_COMMIT_SHA", "ae855f8538843825f9015a419cf4ba5edaf5eec2")

if not os.path.exists(LTX_REPO_DIR):
    os.makedirs(LTX_REPO_DIR)
    subprocess.run(["git", "init", LTX_REPO_DIR], check=True)
    subprocess.run(["git", "remote", "add", "origin", LTX_REPO_URL], cwd=LTX_REPO_DIR, check=True)
    subprocess.run(["git", "fetch", "--depth", "1", "origin", LTX_COMMIT_SHA], cwd=LTX_REPO_DIR, check=True)
    subprocess.run(["git", "checkout", LTX_COMMIT_SHA], cwd=LTX_REPO_DIR, 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 json
import logging
import random
import struct
import tempfile

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

import gradio as gr
import numpy as np
import spaces
from gradio_client import Client
from huggingface_hub import hf_hub_download

import ltx_pipelines.ti2vid_two_stages as ti2vid_module
from ltx_core.components.guiders import MultiModalGuiderParams
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessorOutput
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
from ltx_pipelines.utils.media_io import encode_video

# xformers attention: force-patch + disable Hopper-only FA3 dispatch (Blackwell ZeroGPU).
from ltx_core.model.transformer import attention as _attn_mod
try:
    from xformers.ops import memory_efficient_attention as _mea
    _attn_mod.memory_efficient_attention = _mea
    from xformers.ops.fmha import _set_use_fa3
    _set_use_fa3(False)
    print("[ATTN] xformers patched, FA3 disabled")
except Exception as e:
    print(f"[ATTN] xformers patch failed: {type(e).__name__}: {e}")

# Chunked-read safetensors loader: safe_open mmap deadlocks on FUSE-backed storage.
from ltx_core.loader.primitives import StateDict
from ltx_core.loader.sft_loader import SafetensorsStateDictLoader

_SAFETENSORS_DTYPE_MAP = {
    "F64": torch.float64, "F32": torch.float32, "F16": torch.float16,
    "BF16": torch.bfloat16, "F8_E5M2": torch.float8_e5m2, "F8_E4M3": torch.float8_e4m3fn,
    "I64": torch.int64, "I32": torch.int32, "I16": torch.int16, "I8": torch.int8,
    "U8": torch.uint8, "BOOL": torch.bool,
}


def _patched_load(self, path, sd_ops, device=None):
    sd, size, dtype = {}, 0, set()
    device = device or torch.device("cpu")
    for shard_path in (path if isinstance(path, list) else [path]):
        with open(shard_path, "rb") as f:
            header_len = struct.unpack("<Q", f.read(8))[0]
            header = json.loads(f.read(header_len).decode("utf-8"))
            data_base = 8 + header_len
            for name, meta in header.items():
                if name == "__metadata__":
                    continue
                expected_name = name if sd_ops is None else sd_ops.apply_to_key(name)
                if expected_name is None:
                    continue
                start, end = meta["data_offsets"]
                f.seek(data_base + start)
                buf = f.read(end - start)
                t = torch.frombuffer(bytearray(buf), dtype=_SAFETENSORS_DTYPE_MAP[meta["dtype"]]
                                     ).reshape(meta["shape"])
                t = t.to(device=device, non_blocking=True, copy=False)
                kvs = (((expected_name, t),) if sd_ops is None
                       else sd_ops.apply_to_key_value(expected_name, t))
                for key, v in kvs:
                    size += v.nbytes
                    dtype.add(v.dtype)
                    sd[key] = v
    return StateDict(sd=sd, device=device, size=size, dtype=dtype)


SafetensorsStateDictLoader.load = _patched_load

logging.getLogger().setLevel(logging.INFO)
MAX_SEED = np.iinfo(np.int32).max
FRAME_RATE = 24.0

DEFAULT_NEGATIVE = (
    "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, "
    "motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, "
    "transition, static"
)

# LTX-2.3 guider defaults
VIDEO_STG, VIDEO_RESCALE, A2V_SCALE, STG_BLOCKS = 1.0, 0.7, 3.0, [28]
AUDIO_CFG, AUDIO_STG, AUDIO_RESCALE, V2A_SCALE = 7.0, 1.0, 0.7, 3.0

LTX_REPO = os.environ.get("LTX_REPO", "Lightricks/LTX-2.3")
DEV_FILE = "ltx-2.3-22b-dev.safetensors"
DISTILLED_LORA_FILE = "ltx-2.3-22b-distilled-lora-384-1.1.safetensors"
UPSCALER_FILE = "ltx-2.3-spatial-upscaler-x2-1.1.safetensors"
TOKEN = os.environ.get("HF_TOKEN")
TEXT_ENCODER_SPACE = os.environ.get("TEXT_ENCODER_SPACE", "linoyts/ltx23-gemma-encoder-api")

print("Downloading checkpoints…")
checkpoint_path = hf_hub_download(LTX_REPO, DEV_FILE, token=TOKEN)
distilled_lora_path = hf_hub_download(LTX_REPO, DISTILLED_LORA_FILE, token=TOKEN)
upsampler_path = hf_hub_download(LTX_REPO, UPSCALER_FILE, token=TOKEN)

pipeline = TI2VidTwoStagesPipeline(
    checkpoint_path=checkpoint_path,
    distilled_lora=[LoraPathStrengthAndSDOps(path=distilled_lora_path, strength=1.0, sd_ops=None)],
    spatial_upsampler_path=upsampler_path,
    gemma_root=None,  # text encoding happens on TEXT_ENCODER_SPACE
    loras=[],
    quantization=QuantizationPolicy.fp8_cast(),
)

print("Preloading models (ZeroGPU tensor packing)…")
s1, s2 = pipeline.stage_1_model_ledger, pipeline.stage_2_model_ledger
_cached_s1 = {n: getattr(s1, n)() for n in ("transformer", "video_encoder")}
_cached_s2 = {n: getattr(s2, n)() for n in (
    "transformer", "video_decoder", "audio_decoder", "vocoder", "spatial_upsampler")}
for name, model in _cached_s1.items():
    setattr(s1, name, (lambda m: (lambda: m))(model))
for name, model in _cached_s2.items():
    setattr(s2, name, (lambda m: (lambda: m))(model))
print("Pipeline ready.")


# ---- remote text encoding (positive + negative, needed for CFG) ----

_emb_cache = {}


def fetch_embeddings(prompt, negative_prompt, enhance_prompt, seed):
    """Call the Gemma encoder Space; returns {'positive', 'negative', 'final_prompt'}."""
    key = (prompt, negative_prompt, bool(enhance_prompt), int(seed) if enhance_prompt else 0)
    if key in _emb_cache:
        return _emb_cache[key]
    client = Client(TEXT_ENCODER_SPACE, token=TOKEN)
    emb_file, final_prompt, status = client.predict(
        prompt=prompt, negative_prompt=negative_prompt or "", encode_negative=True,
        enhance_prompt=enhance_prompt, seed=int(seed), api_name="/encode",
    )
    print(f"[encoder] {status} | final prompt: {final_prompt[:80]}…")
    data = torch.load(emb_file, map_location="cpu", weights_only=True)
    if len(_emb_cache) > 32:
        _emb_cache.clear()
    _emb_cache[key] = data
    return data


def _to_output(pack, device):
    return EmbeddingsProcessorOutput(
        video_encoding=pack["video"].to(device),
        audio_encoding=pack["audio"].to(device) if pack["audio"] is not None else None,
        attention_mask=pack["mask"].to(device),
    )


def gpu_duration(embeddings, prompt, negative_prompt="", duration=3.0, *args, **kwargs):
    return min(420, 120 + int(duration) * 30)


@spaces.GPU(duration=gpu_duration)
@torch.inference_mode()
def _generate_gpu(embeddings, prompt, negative_prompt, duration, used_seed,
                  num_inference_steps, video_cfg_scale, width, height,
                  progress=gr.Progress(track_tqdm=True)):
    num_frames = ((int(duration * FRAME_RATE) + 1 - 1 + 7) // 8) * 8 + 1
    tiling_config = TilingConfig.default()
    video_chunks_number = get_video_chunks_number(num_frames, tiling_config)

    precomputed = [_to_output(embeddings["positive"], "cuda"),
                   _to_output(embeddings["negative"], "cuda")]
    original = ti2vid_module.encode_prompts
    ti2vid_module.encode_prompts = lambda *a, **kw: precomputed
    try:
        video, audio = pipeline(
            prompt=prompt,
            negative_prompt=negative_prompt,
            seed=used_seed,
            height=int(height),
            width=int(width),
            num_frames=num_frames,
            frame_rate=FRAME_RATE,
            num_inference_steps=int(num_inference_steps),
            video_guider_params=MultiModalGuiderParams(
                cfg_scale=video_cfg_scale, stg_scale=VIDEO_STG, rescale_scale=VIDEO_RESCALE,
                modality_scale=A2V_SCALE, skip_step=0, stg_blocks=STG_BLOCKS),
            audio_guider_params=MultiModalGuiderParams(
                cfg_scale=AUDIO_CFG, stg_scale=AUDIO_STG, rescale_scale=AUDIO_RESCALE,
                modality_scale=V2A_SCALE, skip_step=0, stg_blocks=STG_BLOCKS),
            images=[],
            tiling_config=tiling_config,
            enhance_prompt=False,  # already enhanced on the encoder Space
        )
    finally:
        ti2vid_module.encode_prompts = original

    output_path = tempfile.mktemp(suffix=".mp4")
    encode_video(video=video, fps=FRAME_RATE, audio=audio, output_path=output_path,
                 video_chunks_number=video_chunks_number)
    return output_path


def generate(
    prompt: str,
    negative_prompt: str = DEFAULT_NEGATIVE,
    duration: float = 5.0,
    seed: int = 42,
    randomize_seed: bool = False,
    num_inference_steps: int = 30,
    video_cfg_scale: float = 3.0,
    width: int = 1536,
    height: int = 864,
    enhance_prompt: bool = True,
    progress=gr.Progress(track_tqdm=True),
):
    """Dev two-stage T2V: returns (video_path, used_seed). Encoder call runs outside the GPU slot."""
    used_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
    progress(0.05, desc="encoding prompts (Gemma Space)…")
    embeddings = fetch_embeddings(prompt, negative_prompt, enhance_prompt, used_seed)
    output_path = _generate_gpu(embeddings, prompt, negative_prompt, duration, used_seed,
                                num_inference_steps, video_cfg_scale, width, height)
    return output_path, used_seed


with gr.Blocks(title="LTX-2.3 dev API") as demo:
    gr.Markdown("# 🎬 LTX-2.3 dev two-stage — hero backend\n"
                "API backend for the LTX-2.3 Pre-Production Playground.")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", lines=3)
            negative_prompt = gr.Textbox(label="Negative prompt", value=DEFAULT_NEGATIVE, lines=2)
            duration = gr.Slider(1, 10, value=5, step=0.5, label="Duration (s)")
            with gr.Row():
                seed = gr.Number(label="Seed", value=42, precision=0)
                randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
            with gr.Row():
                num_inference_steps = gr.Slider(10, 50, value=30, step=1, label="Steps")
                video_cfg_scale = gr.Slider(1, 8, value=3.0, step=0.1, label="CFG")
            with gr.Row():
                width = gr.Number(label="Width", value=1536, precision=0)
                height = gr.Number(label="Height", value=864, precision=0)
            enhance_prompt = gr.Checkbox(label="Enhance prompt", value=True)
            btn = gr.Button("Generate", variant="primary")
        with gr.Column():
            video = gr.Video(label="Video", autoplay=True)
            used_seed = gr.Number(label="Used seed", precision=0)

    btn.click(generate,
              [prompt, negative_prompt, duration, seed, randomize_seed, num_inference_steps,
               video_cfg_scale, width, height, enhance_prompt],
              [video, used_seed], api_name="generate")

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