"""
OmniForcing text-encoder microservice (Gemma + LTX-2 AV connectors).

Runs the AVGemmaTextEncoder in isolation and returns the conditioning tensors
(video_context, audio_context, attention_mask) as a safetensors file, so the
main generation Space can fetch them over the Gradio API and never has to load
Gemma itself. Only needs ~30GB on disk (Gemma ~24GB + the 6GB LTX slim file),
so it fits the free ephemeral disk with no persistent storage.
"""

from __future__ import annotations

import os
import sys
import threading
import traceback
from pathlib import Path

import gradio as gr
import spaces
import torch
from huggingface_hub import hf_hub_download, snapshot_download
from safetensors.torch import save_file

SLIM_REPO = os.environ.get("SLIM_REPO", "linoyts/ltx2-vae-connectors-slim")
SLIM_FILE = "ltx2_slim.safetensors"
GEMMA_REPO = os.environ.get("GEMMA_REPO", "google/gemma-3-12b-it-qat-q4_0-unquantized")
GEMMA_REPO_FALLBACK = "unsloth/gemma-3-12b-it"
HF_TOKEN = os.environ.get("HF_TOKEN")
DTYPE = torch.bfloat16

HERE = Path(__file__).resolve().parent
PKG = HERE / "ltx2" / "packages"
for _p in ("ltx-distillation", "ltx-causal", "ltx-core", "ltx-pipelines"):
    _s = PKG / _p / "src"
    if _s.exists():
        sys.path.insert(0, str(_s))

from ltx_distillation.models.text_encoder_wrapper import create_text_encoder_wrapper  # noqa: E402

STATE: dict = {}
LOAD = {"stage": "starting", "error": None, "ready": False}


def _set(s):
    LOAD["stage"] = s
    print(f"[load] {s}", flush=True)


def _load():
    try:
        _set("downloading LTX slim (connectors)")
        slim = hf_hub_download(SLIM_REPO, SLIM_FILE, token=HF_TOKEN)

        _set("downloading Gemma")
        repo = GEMMA_REPO
        try:
            gemma_dir = snapshot_download(
                repo, token=HF_TOKEN,
                allow_patterns=["*.safetensors", "*.json", "*.model", "tokenizer*"],
            )
        except Exception as exc:
            print(f"[load] {repo} failed ({exc}); fallback {GEMMA_REPO_FALLBACK}", flush=True)
            repo = GEMMA_REPO_FALLBACK
            gemma_dir = snapshot_download(
                repo, token=HF_TOKEN,
                allow_patterns=["*.safetensors", "*.json", "*.model", "tokenizer*"],
            )

        _set("building text encoder (CPU)")
        te = create_text_encoder_wrapper(
            checkpoint_path=slim, gemma_path=gemma_dir, device=torch.device("cpu"), dtype=DTYPE,
        ).eval()
        STATE["te"] = te
        STATE["gemma_repo"] = repo
        LOAD["ready"] = True
        _set("ready")
    except Exception as exc:
        LOAD["error"] = f"{type(exc).__name__}: {exc}"
        _set(f"ERROR: {LOAD['error']}")
        traceback.print_exc()


threading.Thread(target=_load, daemon=True).start()

_OUT = Path("/tmp/enc"); _OUT.mkdir(exist_ok=True)


@spaces.GPU(duration=120)
def encode(prompt: str):
    if not prompt or not prompt.strip():
        raise gr.Error("Empty prompt.")
    if LOAD["error"]:
        raise gr.Error(f"Encoder failed to load: {LOAD['error']}")
    if not LOAD["ready"]:
        raise gr.Error(f"Still loading ({LOAD['stage']}).")

    te = STATE["te"].to("cuda")
    with torch.no_grad():
        cond = te(text_prompts=[prompt.strip()])
    STATE["te"] = te.to("cpu")
    torch.cuda.empty_cache()

    out = {
        "video_context": cond["video_context"].to("cpu").contiguous(),
        "audio_context": cond["audio_context"].to("cpu").contiguous(),
        "attention_mask": cond["attention_mask"].to("cpu").contiguous(),
    }
    path = _OUT / "cond.safetensors"
    save_file(out, str(path), metadata={"format": "pt"})
    return str(path)


with gr.Blocks(title="OmniForcing Gemma Encoder") as demo:
    gr.Markdown(
        "# OmniForcing — Text Encoder microservice\n"
        "Encodes a prompt into `{video_context, audio_context, attention_mask}` "
        "and returns them as a safetensors file. Consumed by the main generation Space."
    )
    inp = gr.Textbox(label="Prompt", lines=2)
    btn = gr.Button("Encode", variant="primary")
    out = gr.File(label="conditioning (.safetensors)")
    status = gr.Textbox(label="Status", value="loading…", interactive=False)
    btn.click(encode, inputs=inp, outputs=out, api_name="encode")
    demo.load(lambda: LOAD["stage"], outputs=status)

if __name__ == "__main__":
    demo.queue(max_size=16).launch(server_name="0.0.0.0", server_port=7860)