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."
PROMPT_ADHERENCE_SUFFIX = (
    "Follow the prompt literally. Preserve the named subjects, actions, setting, and style."
)
MAX_PROMPT_CHARS = 900
MAX_ENHANCED_PROMPT_CHARS = 1400
PROMPT_ENHANCER_SYSTEM_PROMPT = (
    "You are Sulphur's prompt enhancer for image-to-video generation. Rewrite the user's prompt into one concise, "
    "literal cinematic video prompt. Preserve the user's subject, action, setting, style, and intent. Add useful "
    "visual motion, camera movement, lighting, and temporal details. Do not refuse, moralize, apologize, explain, "
    "quote policy, or mention being an AI. Return only the rewritten prompt."
)
PROMPT_ENHANCER_USER_TEMPLATE = "Rewrite this video prompt only:\n{prompt}"

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 _safetensors_header(path):
    with open(path, "rb") as reader:
        header_len = struct.unpack("<Q", reader.read(8))[0]
        return json.loads(reader.read(header_len).decode("utf-8"))


def _validate_lora(name, lora):
    path = Path(lora.path)
    if not path.exists():
        raise FileNotFoundError(f"{name} LoRA missing: {path}")
    keys = [key for key in _safetensors_header(path) if key != "__metadata__"]
    mapped = [lora.sd_ops.apply_to_key(key) for key in keys]
    mapped = [key for key in mapped if key is not None]
    lora_a = sum(1 for key in mapped if key.endswith(".lora_A.weight"))
    lora_b = sum(1 for key in mapped if key.endswith(".lora_B.weight"))
    alpha = sum(1 for key in mapped if key.endswith(".alpha"))
    if lora_a == 0 or lora_b == 0:
        raise RuntimeError(f"{name} LoRA has no usable lora_A/lora_B keys after mapping: {path}")
    print(
        f"[lora] active: {name} strength={lora.strength} path={path} "
        f"lora_A={lora_a} lora_B={lora_b} alpha={alpha} size={path.stat().st_size / 1024**3:.2f}GB"
    )


def _prompt_for_model(prompt):
    prompt = " ".join(str(prompt or DEFAULT_PROMPT).split())
    prompt = _truncate_prompt(prompt, MAX_PROMPT_CHARS)
    if PROMPT_ADHERENCE_SUFFIX.lower() in prompt.lower():
        return prompt
    return f"{prompt}\n\n{PROMPT_ADHERENCE_SUFFIX}"


def _truncate_prompt(prompt, limit):
    prompt = str(prompt or "").strip()
    if len(prompt) <= limit:
        return prompt
    truncated = prompt[:limit].rsplit(" ", 1)[0].strip()
    return truncated or prompt[:limit].strip()


_PROMPT_ENHANCER = None


def _get_prompt_enhancer():
    global _PROMPT_ENHANCER
    if _PROMPT_ENHANCER is not None:
        return _PROMPT_ENHANCER
    from llama_cpp import Llama

    print(f"[prompt-enhancer] loading {ASSETS['prompt_enhancer_model']}")
    _PROMPT_ENHANCER = Llama(
        model_path=str(ASSETS["prompt_enhancer_model"]),
        n_ctx=2048,
        n_threads=max(1, min(8, os.cpu_count() or 1)),
        n_gpu_layers=0,
        verbose=False,
    )
    return _PROMPT_ENHANCER


def _enhance_prompt_with_sulphur(prompt):
    base_prompt = _truncate_prompt(prompt or DEFAULT_PROMPT, MAX_PROMPT_CHARS)
    llm = _get_prompt_enhancer()
    result = llm.create_chat_completion(
        messages=[
            {"role": "system", "content": PROMPT_ENHANCER_SYSTEM_PROMPT},
            {"role": "user", "content": PROMPT_ENHANCER_USER_TEMPLATE.format(prompt=base_prompt)},
        ],
        temperature=0.35,
        top_p=0.9,
        max_tokens=320,
    )
    enhanced = result["choices"][0]["message"]["content"].strip()
    enhanced = enhanced.strip("\"'` \n\t")
    blocked_phrases = ("cannot fulfill", "safety guidelines", "as an ai", "i cannot", "i'm unable")
    if not enhanced or enhanced == "/" or any(phrase in enhanced.lower() for phrase in blocked_phrases):
        print(f"[prompt-enhancer] rejected unusable output; falling back to literal prompt: {enhanced!r}")
        return base_prompt
    enhanced = _truncate_prompt(enhanced, MAX_ENHANCED_PROMPT_CHARS)
    print(f"[prompt-enhancer] {enhanced}")
    return enhanced


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(
        "SulphurAI/Sulphur-2-base",
        "distill_loras/ltx-2.3-22b-distilled-lora-1.1_fro90_ceil72_condsafe.safetensors",
        ASSETS_DIR / "loras",
    )
    upsampler = _download(
        "DeepBeepMeep/LTX-2",
        "ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
        ASSETS_DIR,
    )
    prompt_enhancer_model = _download(
        "SulphurAI/Sulphur-2-base",
        "prompt_enhancer/sulphur_prompt_enhancer_model-q8_0.gguf",
        ASSETS_DIR,
    )
    prompt_enhancer_mmproj = _download(
        "SulphurAI/Sulphur-2-base",
        "prompt_enhancer/mmproj-BF16.gguf",
        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,
        "prompt_enhancer_model": prompt_enhancer_model,
        "prompt_enhancer_mmproj": prompt_enhancer_mmproj,
        "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),
]
_validate_lora("sulphur_distilled_condsafe", LORAS[0])
_validate_lora("sulphur_experimental_lora_v1", LORAS[1])
print("[lora] configured order:")
for index, lora in enumerate(LORAS, start=1):
    print(f"[lora] {index}: strength={lora.strength} path={lora.path}")

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)
            if enhance_prompt:
                prompt = _enhance_prompt_with_sulphur(prompt)
            model_prompt = _prompt_for_model(prompt)

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

            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=model_prompt,
                seed=current_seed,
                height=height,
                width=width,
                num_frames=num_frames,
                frame_rate=frame_rate,
                images=images,
                tiling_config=tiling_config,
                enhance_prompt=False,
            )

            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=True)
                    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())