"""Pinned native Kimodo inference; only paths and CPU results cross ZeroGPU's queue."""

from __future__ import annotations

import json
import math
from pathlib import Path

import numpy as np
import spaces

MODEL = "nvidia/Kimodo-SOMA-RP-v1.1"
_MODEL = None
_REVISIONS = {}


def load_motion_model(runtime: Path, lock: dict):
    """Load once at Space startup, before ZeroGPU forks its inference worker."""
    from kimodo.model import LLM2VecEncoder
    from kimodo.model.llm2vec.llm2vec import LLM2Vec
    from kimodo.model.llm2vec.models.bidirectional_llama import LlamaBiModel
    from kimodo.model.loading import instantiate_from_dict
    from omegaconf import OmegaConf
    from peft import PeftModel
    from transformers import AutoTokenizer
    import torch

    global _MODEL, _REVISIONS
    checkpoint = runtime / "kimodo-model"
    class PinnedTextEncoder(LLM2VecEncoder):
        def __init__(self):
            mntp = runtime / "kimodo-text/mntp"
            base = json.loads((mntp / "adapter_config.json").read_text())["base_model_name_or_path"]
            # Transformers 5 no longer treats an adapter directory as base weights.
            # Merge MNTP into the base before attaching the supervised adapter.
            model = LlamaBiModel.from_pretrained(base, dtype=torch.bfloat16)
            model = PeftModel.from_pretrained(model, str(mntp), torch_device="cpu").merge_and_unload()
            model = PeftModel.from_pretrained(
                model, str(runtime / "kimodo-text/supervised"), torch_device="cpu",
            )
            tokenizer = AutoTokenizer.from_pretrained(base)
            tokenizer.pad_token = tokenizer.eos_token
            tokenizer.padding_side = "left"
            self.model = LLM2Vec(model=model, tokenizer=tokenizer)
            # LLM2Vec uses this canonical identity for Llama's instruction template.
            self.model.model.config._name_or_path = "meta-llama/Meta-Llama-3-8B-Instruct"
            self.llm_dim = 4096
            self._device = "cuda"
            self.model = self.model.to(self._device).eval()
            for parameter in self.model.parameters():
                parameter.requires_grad = False

    encoder = PinnedTextEncoder()
    config = OmegaConf.load(checkpoint / "config.yaml")
    config = OmegaConf.merge(config, {"checkpoint_dir": str(checkpoint)})
    values = OmegaConf.to_container(config, resolve=True)
    values.pop("checkpoint_dir", None)
    values["text_encoder"] = None
    _MODEL = instantiate_from_dict(values, overrides={"device": "cuda"}).eval()
    _MODEL.text_encoder = encoder
    _REVISIONS = {key: value for key, value in lock.items()
                  if key == "kimodo_code" or key == MODEL or "LLM2Vec" in key or "Meta-Llama-3" in key}
    if float(_MODEL.fps) != 30.0 or _MODEL.output_skeleton.dim != 77:
        raise RuntimeError("Pinned Kimodo must produce the SOMA77 skeleton at 30 fps.")


def validate_motion(arrays: dict, frames: int):
    expected = {"posed_joints": (frames, 77, 3), "global_rot_mats": (frames, 77, 3, 3),
                "local_rot_mats": (frames, 77, 3, 3), "root_positions": (frames, 3)}
    for key, shape in expected.items():
        if key not in arrays or arrays[key].shape != shape or not np.isfinite(arrays[key]).all():
            raise RuntimeError(f"Kimodo returned invalid {key}; expected finite {shape}.")
    contacts = arrays.get("foot_contacts")
    if contacts is None or contacts.shape[0] != frames or not np.isfinite(contacts).all():
        raise RuntimeError("Kimodo returned invalid foot contacts.")


@spaces.GPU(duration=120)
def generate_motion(prompt: str, duration: float, seed: int, output: Path) -> dict:
    """Generate one complete prompt as a clip, preserving its full requested duration."""
    import torch
    from kimodo.exports.bvh import save_motion_bvh
    from kimodo.exports.motion_io import save_kimodo_npz
    from kimodo.skeleton import global_rots_to_local_rots
    from kimodo.tools import seed_everything

    if _MODEL is None:
        raise RuntimeError("Kimodo was not loaded during Space startup.")
    if not prompt.strip() or len(prompt) > 2000:
        raise ValueError("Animation prompt must contain 1–2000 characters.")
    if not math.isfinite(duration) or not 2 <= duration <= 12:
        raise ValueError("Animation duration must be between 2 and 12 seconds.")
    output = Path(output)
    output.mkdir(parents=True, exist_ok=True)
    frames = round(duration * _MODEL.fps)
    seed_everything(int(seed))
    with torch.inference_mode():
        batch = _MODEL(
            prompts=[prompt.strip()], num_frames=[frames], num_samples=1,
            num_denoising_steps=100, multi_prompt=False, constraint_lst=[],
            post_processing=True, return_numpy=True,
        )
    arrays = {key: value[0] for key, value in batch.items()}
    validate_motion(arrays, frames)
    npz, bvh = output / "motion.npz", output / "motion.bvh"
    save_kimodo_npz(str(npz), arrays)
    skeleton = _MODEL.output_skeleton
    rotations = torch.from_numpy(arrays["global_rot_mats"]).to(_MODEL.device)
    roots = torch.from_numpy(arrays["posed_joints"][:, skeleton.root_idx]).to(_MODEL.device)
    save_motion_bvh(str(bvh), global_rots_to_local_rots(rotations, skeleton), roots,
                    skeleton=skeleton, fps=_MODEL.fps, standard_tpose=False)
    provenance = {
        "provider": "native NVIDIA Kimodo", "model": MODEL, "revisions": _REVISIONS,
        "prompt": prompt.strip(), "seed": int(seed), "frames": frames, "fps": float(_MODEL.fps),
        "duration_seconds": frames / _MODEL.fps, "diffusion_steps": 100,
        "post_processing": "official foot-skate correction", "skeleton": "somaskel77",
    }
    (output / "provenance.json").write_text(json.dumps(provenance, indent=2))
    return {"npz": npz, "bvh": bvh, "provenance": provenance}