import json
import logging
import os
import random
import struct
import subprocess
import sys
import tempfile
import builtins
from pathlib import Path

import gradio as gr
import numpy as np
import spaces
import torch
from huggingface_hub import hf_hub_download


LTX_REPO_URL = "https://github.com/Lightricks/LTX-2.git"
LTX_COMMIT_SHA = "ae855f8538843825f9015a419cf4ba5edaf5eec2"

HF_TOKEN = os.environ.get("HF_TOKEN")
PERSISTENT = Path("/data") if Path("/data").exists() else Path(tempfile.gettempdir())
ASSETS_DIR = PERSISTENT / "sulphur_ltx"
LTX_REPO_DIR = PERSISTENT / "LTX-2"
OUTPUT_DIR = PERSISTENT / "sulphur_outputs"

MAX_SEED = np.iinfo(np.int32).max
DEFAULT_FRAME_RATE = 24.0
DEFAULT_PROMPT = "Make this image come alive with cinematic motion, smooth animation."

RESOLUTIONS = {
    "high": {"16:9": (1536, 1024), "9:16": (1024, 1536), "1:1": (1024, 1024)},
    "low": {"16:9": (768, 512), "9:16": (512, 768), "1:1": (768, 768)},
}


def _run(cmd, cwd=None, check=True):
    print("[setup]", " ".join(str(part) for part in cmd))
    return subprocess.run(cmd, cwd=cwd, check=check)


def install_ltx_runtime():
    if not (LTX_REPO_DIR / ".git").exists():
        LTX_REPO_DIR.parent.mkdir(parents=True, exist_ok=True)
        if LTX_REPO_DIR.exists():
            raise RuntimeError(f"{LTX_REPO_DIR} exists but is not a git checkout")
        _run(["git", "init", str(LTX_REPO_DIR)])
        _run(["git", "remote", "add", "origin", LTX_REPO_URL], cwd=LTX_REPO_DIR)

    _run(["git", "fetch", "--depth", "1", "origin", LTX_COMMIT_SHA], cwd=LTX_REPO_DIR)
    _run(["git", "checkout", LTX_COMMIT_SHA], cwd=LTX_REPO_DIR)

    _run(
        [
            sys.executable,
            "-m",
            "pip",
            "install",
            "--force-reinstall",
            "--no-deps",
            "-e",
            str(LTX_REPO_DIR / "packages" / "ltx-core"),
            "-e",
            str(LTX_REPO_DIR / "packages" / "ltx-pipelines"),
        ]
    )

    sys.path.insert(0, str(LTX_REPO_DIR / "packages" / "ltx-pipelines" / "src"))
    sys.path.insert(0, str(LTX_REPO_DIR / "packages" / "ltx-core" / "src"))


install_ltx_runtime()


_real_import = builtins.__import__


def _import_without_xformers(name, globals=None, locals=None, fromlist=(), level=0):
    if name == "xformers" or name.startswith("xformers."):
        raise ImportError("xformers is disabled for this Space runtime")
    return _real_import(name, globals, locals, fromlist, level)


builtins.__import__ = _import_without_xformers
try:
    from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
    from ltx_core.loader.primitives import StateDict
    from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps
    from ltx_core.loader.sft_loader import SafetensorsStateDictLoader
    from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
    from ltx_core.quantization import QuantizationPolicy
    import ltx_pipelines.utils.model_ledger as _model_ledger_mod
    from ltx_pipelines.distilled import DistilledPipeline
    from ltx_pipelines.utils.args import ImageConditioningInput
    from ltx_pipelines.utils.media_io import encode_video
finally:
    builtins.__import__ = _real_import


_model_ledger_mod.GEMMA_LLM_KEY_OPS = (
    SDOps("GEMMA_LLM_KEY_OPS_SULPHUR")
    .with_matching(prefix="model.language_model.")
    .with_matching(prefix="model.vision_tower.")
    .with_matching(prefix="model.multi_modal_projector.")
    .with_replacement("model.language_model.", "model.model.language_model.")
    .with_replacement("model.vision_tower.", "model.model.vision_tower.")
    .with_replacement("model.multi_modal_projector.", "model.model.multi_modal_projector.")
    .with_kv_operation(
        operation=lambda key, value: [
            KeyValueOperationResult(key, value),
            KeyValueOperationResult("model.lm_head.weight", value),
        ],
        key_prefix="model.model.language_model.embed_tokens.weight",
    )
)


_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_safetensors_load(self, path, sd_ops, device=None):
    sd = {}
    size = 0
    dtype = set()
    device = device or torch.device("cpu")
    model_paths = path if isinstance(path, list | tuple) else [path]
    for shard_path in model_paths:
        with open(shard_path, "rb") as reader:
            header_len = struct.unpack("<Q", reader.read(8))[0]
            header = json.loads(reader.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"]
                reader.seek(data_base + start)
                raw = reader.read(end - start)
                tensor = torch.frombuffer(
                    bytearray(raw), dtype=_SAFETENSORS_DTYPE_MAP[meta["dtype"]]
                ).reshape(meta["shape"])
                tensor = tensor.to(device=device, non_blocking=True, copy=False)
                kvs = ((expected_name, tensor),) if sd_ops is None else sd_ops.apply_to_key_value(expected_name, tensor)
                for key, value in kvs:
                    size += value.nbytes
                    dtype.add(value.dtype)
                    sd[key] = value
    return StateDict(sd=sd, device=device, size=size, dtype=dtype)


SafetensorsStateDictLoader.load = _patched_safetensors_load
print("[FUSE-PATCH] Safetensors loader uses chunked reads")

logging.getLogger().setLevel(logging.INFO)


def _download(repo_id, filename, local_dir):
    local_dir = Path(local_dir)
    local_dir.mkdir(parents=True, exist_ok=True)
    path = Path(
        hf_hub_download(
            repo_id=repo_id,
            filename=filename,
            local_dir=str(local_dir),
            token=HF_TOKEN,
        )
    )
    print(f"[download] ready: {repo_id}/{filename}")
    return path


def _ensure_gemma_model_alias(gemma_weight):
    gemma_weight = Path(gemma_weight)
    alias = gemma_weight.with_name("model.safetensors")
    if gemma_weight == alias:
        return alias
    if alias.exists() or alias.is_symlink():
        alias.unlink()
    gemma_weight.replace(alias)
    print(f"[download] ready: {alias}")
    return alias


def download_assets():
    checkpoint = _download("SulphurAI/Sulphur-2-base", "sulphur_distil_bf16.safetensors", ASSETS_DIR)
    sulphur_lora = _download(
        "SulphurAI/Sulphur-2-base",
        "experimental/sulphur_experimental_lora_v1.safetensors",
        ASSETS_DIR / "loras",
    )
    distilled_lora = _download(
        "DeepBeepMeep/LTX-2",
        "ltx-2.3-22b-distilled-lora-384.safetensors",
        ASSETS_DIR / "loras",
    )
    upsampler = _download(
        "DeepBeepMeep/LTX-2",
        "ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
        ASSETS_DIR,
    )

    gemma_folder = "gemma-3-12b-it-qat-q4_0-unquantized"
    gemma_files = [
        f"{gemma_folder}/{gemma_folder}.safetensors",
        f"{gemma_folder}/added_tokens.json",
        f"{gemma_folder}/chat_template.json",
        f"{gemma_folder}/config.json",
        f"{gemma_folder}/config_light.json",
        f"{gemma_folder}/generation_config.json",
        f"{gemma_folder}/preprocessor_config.json",
        f"{gemma_folder}/processor_config.json",
        f"{gemma_folder}/special_tokens_map.json",
        f"{gemma_folder}/tokenizer.json",
        f"{gemma_folder}/tokenizer.model",
        f"{gemma_folder}/tokenizer_config.json",
    ]
    gemma_weight = None
    for filename in gemma_files:
        path = _download("DeepBeepMeep/LTX-2", filename, ASSETS_DIR)
        if filename.endswith(".safetensors"):
            gemma_weight = path

    old_quanto = ASSETS_DIR / gemma_folder / f"{gemma_folder}_quanto_bf16_int8.safetensors"
    if old_quanto.exists():
        old_quanto.unlink()
    _ensure_gemma_model_alias(gemma_weight)

    return {
        "checkpoint": checkpoint,
        "sulphur_lora": sulphur_lora,
        "distilled_lora": distilled_lora,
        "upsampler": upsampler,
        "gemma_root": ASSETS_DIR / gemma_folder,
    }


ASSETS = download_assets()

LORAS = [
    LoraPathStrengthAndSDOps(str(ASSETS["distilled_lora"]), 1.0, LTXV_LORA_COMFY_RENAMING_MAP),
    LoraPathStrengthAndSDOps(str(ASSETS["sulphur_lora"]), 1.0, LTXV_LORA_COMFY_RENAMING_MAP),
]

pipeline = DistilledPipeline(
    distilled_checkpoint_path=str(ASSETS["checkpoint"]),
    spatial_upsampler_path=str(ASSETS["upsampler"]),
    gemma_root=str(ASSETS["gemma_root"]),
    loras=LORAS,
    quantization=QuantizationPolicy.fp8_cast(),
)

print("[startup] LTX assets ready; model loading is deferred to the GPU worker")


def log_memory(tag):
    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / 1024**3
        peak = torch.cuda.max_memory_allocated() / 1024**3
        free, total = torch.cuda.mem_get_info()
        print(
            f"[VRAM {tag}] allocated={allocated:.2f}GB peak={peak:.2f}GB "
            f"free={free / 1024**3:.2f}GB total={total / 1024**3:.2f}GB"
        )


def detect_aspect_ratio(image):
    if image is None:
        return "16:9"
    width, height = image.size
    ratio = width / height
    candidates = {"16:9": 16 / 9, "9:16": 9 / 16, "1:1": 1.0}
    return min(candidates, key=lambda key: abs(ratio - candidates[key]))


def update_resolution(image, high_res):
    aspect = detect_aspect_ratio(image)
    tier = "high" if high_res else "low"
    width, height = RESOLUTIONS[tier][aspect]
    return gr.update(value=width), gr.update(value=height)


@spaces.GPU
def generate_video(
    input_image,
    prompt,
    duration,
    enhance_prompt,
    seed,
    randomize_seed,
    height,
    width,
    progress=gr.Progress(track_tqdm=True),
):
    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
    try:
        with torch.no_grad():
            if torch.cuda.is_available():
                torch.cuda.reset_peak_memory_stats()
            log_memory("start")

            frame_rate = DEFAULT_FRAME_RATE
            num_frames = int(float(duration) * frame_rate) + 1
            num_frames = ((num_frames - 1 + 7) // 8) * 8 + 1
            height = int(height)
            width = int(width)

            print(f"[generate] {width}x{height}, frames={num_frames}, seed={current_seed}")

            images = []
            if input_image is not None:
                OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
                input_path = OUTPUT_DIR / f"input_{current_seed}.png"
                input_image.save(input_path)
                images = [ImageConditioningInput(path=str(input_path), frame_idx=0, strength=1.0)]

            tiling_config = TilingConfig.default()
            chunks = get_video_chunks_number(num_frames, tiling_config)
            log_memory("before pipeline")

            video, audio = pipeline(
                prompt=prompt,
                seed=current_seed,
                height=height,
                width=width,
                num_frames=num_frames,
                frame_rate=frame_rate,
                images=images,
                tiling_config=tiling_config,
                enhance_prompt=bool(enhance_prompt),
            )

            output_path = tempfile.mktemp(suffix=".mp4")
            encode_video(video=video, fps=frame_rate, audio=audio, output_path=output_path, video_chunks_number=chunks)
            log_memory("after encode")
            return output_path, current_seed
    except Exception as exc:
        import traceback

        log_memory("error")
        print(f"[ERROR] {exc}\n{traceback.format_exc()}")
        raise gr.Error(str(exc))


with gr.Blocks(title="Sulphur LTX Image to Video") as demo:
    gr.Markdown("# Sulphur Image to Video")

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Input Image", type="pil")
            prompt = gr.Textbox(
                label="Prompt",
                value=DEFAULT_PROMPT,
                lines=3,
                placeholder="Describe the motion and animation...",
            )
            with gr.Row():
                duration = gr.Slider(1.0, 5.0, value=3.0, step=0.1, label="Duration")
                with gr.Column():
                    enhance_prompt = gr.Checkbox(label="Enhance Prompt", value=False)
                    high_res = gr.Checkbox(label="High Resolution", value=False)

            generate_btn = gr.Button("Generate", variant="primary")

            with gr.Accordion("Advanced", open=False):
                seed = gr.Slider(0, MAX_SEED, value=10, step=1, label="Seed")
                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
                with gr.Row():
                    width = gr.Number(label="Width", value=768, precision=0)
                    height = gr.Number(label="Height", value=512, precision=0)

        with gr.Column():
            output_video = gr.Video(label="Output Video", autoplay=True)

    input_image.change(fn=update_resolution, inputs=[input_image, high_res], outputs=[width, height])
    high_res.change(fn=update_resolution, inputs=[input_image, high_res], outputs=[width, height])

    generate_btn.click(
        fn=generate_video,
        inputs=[input_image, prompt, duration, enhance_prompt, seed, randomize_seed, height, width],
        outputs=[output_video, seed],
    )


if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus())