"""LTX-2.3 Gemma text-encoder backend (API Space for the Pre-Production Playground).

Runs Gemma-3-12B + the LTX-2.3 embeddings connectors and returns final
EmbeddingsProcessorOutput tensors (video/audio encodings + attention mask) as a .pt file,
for positive and (optionally) negative prompts — so the generation backends don't need
to hold Gemma in VRAM. Mirrors helpers.encode_prompts at the pinned LTX-2 commit.
Adapted from linoyts/gemma-text-encoder (LTX-2) to LTX-2.3.
"""

import os
import subprocess
import sys

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

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 struct
import tempfile
import time

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

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

from ltx_pipelines.utils.helpers import generate_enhanced_prompt
from ltx_pipelines.utils.model_ledger import ModelLedger

# 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

LTX_REPO = os.environ.get("LTX_REPO", "Lightricks/LTX-2.3")
CHECKPOINT_FILE = "ltx-2.3-22b-distilled-1.1.safetensors"  # connector weights only
GEMMA_REPO = "google/gemma-3-12b-it-qat-q4_0-unquantized"
TOKEN = os.environ.get("HF_TOKEN")

print("Downloading checkpoint (connectors) + Gemma…")
checkpoint_path = hf_hub_download(LTX_REPO, CHECKPOINT_FILE, token=TOKEN)
gemma_root = snapshot_download(GEMMA_REPO, token=TOKEN)

ledger = ModelLedger(
    dtype=torch.bfloat16,
    device="cuda",
    checkpoint_path=checkpoint_path,
    gemma_root_path=gemma_root,
)

print("Preloading text encoder + embeddings processor…")
_text_encoder = ledger.text_encoder()
_embeddings_processor = ledger.gemma_embeddings_processor()
print("Encoder ready.")


def _pack(output):
    return {"video": output.video_encoding.cpu(),
            "audio": output.audio_encoding.cpu() if output.audio_encoding is not None else None,
            "mask": output.attention_mask.cpu()}


@spaces.GPU(duration=60)
@torch.inference_mode()
def encode(
    prompt: str,
    negative_prompt: str = "",
    encode_negative: bool = False,
    enhance_prompt: bool = True,
    seed: int = 42,
):
    """Encode prompt (and optionally negative prompt) → .pt with final LTX-2.3 embeddings.

    Returns (embeddings_file, final_prompt, status).
    """
    start = time.time()
    final_prompt = prompt
    if enhance_prompt:
        final_prompt = generate_enhanced_prompt(_text_encoder, prompt, None, seed=int(seed))

    prompts = [final_prompt] + ([negative_prompt or ""] if encode_negative else [])
    raw = [_text_encoder.encode(p) for p in prompts]
    outputs = [_embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw]

    data = {"positive": _pack(outputs[0]),
            "negative": _pack(outputs[1]) if encode_negative else None,
            "final_prompt": final_prompt}
    out_path = tempfile.mktemp(suffix=".pt")
    torch.save(data, out_path)

    status = f"encoded in {time.time() - start:.1f}s"
    return out_path, final_prompt, status


with gr.Blocks(title="LTX-2.3 Gemma encoder API") as demo:
    gr.Markdown("# 🧠 LTX-2.3 Gemma text encoder\n"
                "API backend for the LTX-2.3 Pre-Production Playground — returns final "
                "embeddings (video/audio encodings + mask) for positive and negative prompts.")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", lines=4)
            negative_prompt = gr.Textbox(label="Negative prompt", lines=2)
            encode_negative = gr.Checkbox(label="Encode negative (CFG / two-stage)", value=False)
            enhance_prompt = gr.Checkbox(label="Enhance prompt", value=True)
            seed = gr.Number(label="Enhancement seed", value=42, precision=0)
            btn = gr.Button("Encode", variant="primary")
        with gr.Column():
            emb_file = gr.File(label="Embeddings (.pt)")
            final_prompt = gr.Textbox(label="Final prompt used", lines=4)
            status = gr.Textbox(label="Status")

    btn.click(encode, [prompt, negative_prompt, encode_negative, enhance_prompt, seed],
              [emb_file, final_prompt, status], api_name="encode")

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