"""VAE encode/decode nodes."""

from __future__ import annotations

import os
from pathlib import Path
from typing import Any, ClassVar, Optional

from loguru import logger
import torch

from .base import BaseNode, NodeDefinition, NodeParam, NodePort, NodeRegistry
from .interpolation import INTERP_METHOD_NAMES, INTERPOLATIONS
from .types import Audio, Curve, Latent, ModelHandle, VAEHandle


# -----------------------------------------------------------------------
# TRT VAE helpers (loaded once, reused across calls)
# -----------------------------------------------------------------------

_trt_vae_cache: dict[str, Any] = {}

# Shared polygraphy CUDA stream for all TRT engines in this process.
# Using torch.cuda.Stream causes a 14x performance degradation on
# Blackwell GPUs when multiple TRT engines coexist. Polygraphy's
# cuda.Stream (a thin wrapper around cudaStreamCreate) avoids this.
_trt_stream = None

def _get_trt_stream():
    """Get or create the shared polygraphy CUDA stream."""
    global _trt_stream
    if _trt_stream is None:
        from polygraphy import cuda as pg_cuda
        _trt_stream = pg_cuda.Stream()
    return _trt_stream


def _trt_available() -> bool:
    """Check if TensorRT is importable."""
    try:
        import tensorrt  # noqa: F401
        return True
    except ImportError:
        return False


def _trt_tensor_dtype(engine, name: str) -> torch.dtype:
    import tensorrt as trt

    dtype_map = {
        trt.float32: torch.float32,
        trt.float16: torch.float16,
        trt.int32: torch.int32,
        trt.int8: torch.int8,
        trt.bool: torch.bool,
    }
    if hasattr(trt, "bfloat16"):
        dtype_map[trt.bfloat16] = torch.bfloat16
    return dtype_map.get(engine.get_tensor_dtype(name), torch.float32)


def _get_trt_vae(engine_path: str, device: torch.device):
    """Load or return cached TRT VAE engine + context + stream."""
    engine_path = os.path.abspath(engine_path)
    if engine_path in _trt_vae_cache:
        return _trt_vae_cache[engine_path]

    from polygraphy.backend.common import bytes_from_path
    from polygraphy.backend.trt import engine_from_bytes

    engine = engine_from_bytes(bytes_from_path(engine_path))
    ctx = engine.create_execution_context()
    # TensorRT's create_execution_context() returns None when the
    # runtime can't allocate the engine's workspace (typically under
    # CUDA-OOM pressure on the 240 s vae_encode profile's ~16 GiB
    # reservation). Caching a None context here would poison the cache
    # entry: a retry would short-circuit on the dict lookup and return
    # the same None, and the next _trt_vae_decode/encode call would
    # raise a bare ``'NoneType' object has no attribute 'set_input_shape'``
    # with no diagnostic context. Raise INSTEAD of caching so the
    # caller can attempt recovery (e.g. torch.cuda.empty_cache() +
    # retry) and so a successful retry actually re-runs the load.
    if ctx is None:
        logger.error(
            "trt_vae_context_creation_returned_none engine={}",
            engine_path,
        )
        raise RuntimeError(
            f"TRT VAE engine.create_execution_context() returned None "
            f"(likely CUDA OOM on workspace alloc): {engine_path}"
        )
    logger.info("Loaded TRT VAE engine: {}", engine_path)

    tensor_dtypes = {
        engine.get_tensor_name(i): _trt_tensor_dtype(engine, engine.get_tensor_name(i))
        for i in range(engine.num_io_tensors)
    }
    entry = {"engine": engine, "context": ctx, "tensor_dtypes": tensor_dtypes}
    _trt_vae_cache[engine_path] = entry
    return entry


def _evict_trt_vae(engine_path: str) -> bool:
    """Drop a cached TRT VAE engine and free its GPU buffers.

    Used by the profile manager when swapping to a different-duration
    engine: the old context, its workspace, and the cached output buffer
    must be released before the new engine is loaded so VRAM doesn't
    spike past the budget.

    Returns True if an entry was evicted, False if the path wasn't cached.
    """
    engine_path = os.path.abspath(engine_path)
    entry = _trt_vae_cache.pop(engine_path, None)
    if entry is None:
        return False
    # Clear refs in deterministic order: cached I/O buffers first (so
    # nothing still holds the device pointer), then the execution
    # context, then the engine. Polygraphy releases workspace when the
    # context's refcount hits zero.
    entry.pop("_decode_buf", None)
    entry.pop("_encode_buf", None)
    entry["context"] = None
    entry["engine"] = None
    torch.cuda.empty_cache()
    logger.info("Evicted TRT VAE engine: {}", engine_path)
    return True


def _trt_vae_decode(
    latents_bdt: torch.Tensor, engine_path: str, device: torch.device
) -> torch.Tensor:
    """Decode latents [B, D, T] -> audio [B, 2, samples] via TRT."""
    entry = _get_trt_vae(engine_path, device)
    ctx = entry["context"]
    stream = _get_trt_stream()

    dtypes = entry["tensor_dtypes"]
    lat = latents_bdt.to(device=device, dtype=dtypes.get("latents", torch.float32)).contiguous()

    if not ctx.set_input_shape("latents", tuple(lat.shape)):
        logger.error(
            "trt_vae_decode_rejected_input_shape input_shape={} engine={}",
            tuple(lat.shape), engine_path,
        )
        raise RuntimeError(f"TRT VAE decode rejected input shape: {tuple(lat.shape)}")
    if not ctx.set_tensor_address("latents", lat.data_ptr()):
        logger.error(
            "trt_vae_decode_rejected_input_address engine={}", engine_path,
        )
        raise RuntimeError("TRT VAE decode rejected input address")

    missing = ctx.infer_shapes()
    if missing:
        logger.error(
            "trt_vae_decode_shapes_insufficient missing={} engine={}",
            missing, engine_path,
        )
        raise RuntimeError(f"TRT VAE decode shapes are insufficiently specified: {missing}")

    out_shape = tuple(ctx.get_tensor_shape("audio"))
    out_dtype = dtypes.get("audio", torch.float32)

    cached = entry.get("_decode_buf")
    if cached is not None and cached.shape == out_shape and cached.dtype == out_dtype:
        audio_buf = cached
    else:
        audio_buf = torch.empty(out_shape, dtype=out_dtype, device=device)
        entry["_decode_buf"] = audio_buf

    if not ctx.set_tensor_address("audio", audio_buf.data_ptr()):
        logger.error(
            "trt_vae_decode_rejected_output_address engine={}", engine_path,
        )
        raise RuntimeError("TRT VAE decode rejected output address")

    if not ctx.execute_async_v3(stream.ptr):
        logger.error(
            "trt_vae_decode_execute_failed engine={} input_shape={}",
            engine_path, tuple(lat.shape),
        )
        raise RuntimeError("TRT VAE decode failed")
    # No host sync: the polygraphy stream is a blocking stream w.r.t.
    # PyTorch's legacy default stream, so the .clone() below is
    # implicitly ordered after the TRT execution on the GPU. The
    # caller's eventual D2H (.cpu() at emission) is the natural sync
    # point; blocking the host here would only serialize CPU prep
    # against GPU work.
    return audio_buf.clone().to(torch.float32)


# Latent frame geometry shared by every ACE-Step VAE engine: 48 kHz
# audio, 25 latent frames per second -> 1920 samples per latent frame.
_VAE_SAMPLES_PER_FRAME = 1920

# Per-side context margin (in latent frames) for chunked VAE encode.
# Empirically the Oobleck encoder's receptive field is under 8 frames:
# with an 8-frame margin a chunked encode is bit-exact (fp16) against
# the full-engine encode — see
# scripts/benchmarks/validate_chunked_vae_encode.py. 64 frames (2.56 s)
# keeps a comfortable safety factor at negligible cost.
_VAE_ENCODE_CHUNK_MARGIN_FRAMES = 64


def _plan_encode_chunks(
    n_samples: int, max_samples: int, min_samples: int,
) -> list[tuple[int, int, int, int]]:
    """Chunk plan for an encode input longer than the engine max.

    Returns ``(ctx_start, ctx_end, core_start, core_end)`` tuples in
    latent frames. ``ctx_*`` is the audio span fed to the engine
    (core plus receptive-field margin); ``core_*`` is the span whose
    moments are kept. Cores tile ``[0, total_frames)`` exactly. Every
    context span fits the engine profile: at most ``max_samples`` and
    at least ``min_samples`` long (the tail chunk's left context is
    extended when needed — extra frames are trimmed with the margin).
    """
    spf = _VAE_SAMPLES_PER_FRAME
    margin = _VAE_ENCODE_CHUNK_MARGIN_FRAMES
    total_frames = (n_samples + spf - 1) // spf
    core_frames = max_samples // spf - 2 * margin
    if core_frames <= 0:
        raise RuntimeError(
            f"TRT VAE encode engine max profile ({max_samples} samples) "
            f"is too small for chunked encode"
        )
    min_ctx_frames = (min_samples + spf - 1) // spf
    plan: list[tuple[int, int, int, int]] = []
    core_start = 0
    while core_start < total_frames:
        core_end = min(core_start + core_frames, total_frames)
        ctx_start = max(0, core_start - margin)
        ctx_end = min(total_frames, core_end + margin)
        # The tail chunk can fall below the engine's min profile shape;
        # extend the left context (trimmed away by the caller) to stay
        # within range. n_samples > max_samples >= min_samples
        # guarantees enough audio exists to the left.
        if ctx_end - ctx_start < min_ctx_frames:
            ctx_start = max(0, ctx_end - min_ctx_frames)
        plan.append((ctx_start, ctx_end, core_start, core_end))
        core_start = core_end
    return plan


def _trt_vae_encode_moments(
    inp: torch.Tensor, entry: dict, engine_path: str
) -> torch.Tensor:
    """Single-shot TRT encode of audio [B, 2, samples] -> moments [B, 128, T].

    ``inp`` must already be on-device, contiguous, and within the
    engine's profile range.
    """
    ctx = entry["context"]
    stream = _get_trt_stream()
    dtypes = entry["tensor_dtypes"]

    # Log structured TRT-VAE-encode failures from the deepest call site so
    # the JSON record carries input_shape + engine_path. The surrounding
    # backend has session_id / client_id / fixture_name /
    # audio_duration_s bound into loguru's contextvars, so the joined
    # record answers "which track / session triggered this?" without any
    # plumbing through the call stack. We log-then-raise (rather than
    # logger.exception inside a try/except wrapper) because the existing
    # callers want the RuntimeError unchanged.
    if not ctx.set_input_shape("audio", tuple(inp.shape)):
        logger.error(
            "trt_vae_encode_rejected_input_shape input_shape={} engine={}",
            tuple(inp.shape), engine_path,
        )
        raise RuntimeError(f"TRT VAE encode rejected input shape: {tuple(inp.shape)}")
    if not ctx.set_tensor_address("audio", inp.data_ptr()):
        logger.error(
            "trt_vae_encode_rejected_input_address engine={}", engine_path,
        )
        raise RuntimeError("TRT VAE encode rejected input address")

    missing = ctx.infer_shapes()
    if missing:
        logger.error(
            "trt_vae_encode_shapes_insufficient missing={} engine={}",
            missing, engine_path,
        )
        raise RuntimeError(f"TRT VAE encode shapes are insufficiently specified: {missing}")

    out_shape = tuple(ctx.get_tensor_shape("moments"))
    moments_dtype = dtypes.get("moments", torch.float32)
    moments_buf = torch.empty(out_shape, dtype=moments_dtype, device=inp.device)
    if not ctx.set_tensor_address("moments", moments_buf.data_ptr()):
        logger.error(
            "trt_vae_encode_rejected_output_address engine={}", engine_path,
        )
        raise RuntimeError("TRT VAE encode rejected output address")

    if not ctx.execute_async_v3(stream.ptr):
        logger.error(
            "trt_vae_encode_execute_failed engine={} input_shape={}",
            engine_path, tuple(inp.shape),
        )
        raise RuntimeError("TRT VAE encode failed")
    stream.synchronize()
    return moments_buf


def _trt_vae_encode(
    audio_bct: torch.Tensor, engine_path: str, device: torch.device
) -> torch.Tensor:
    """Encode audio [B, 2, samples] -> latents [B, D, T] via TRT.

    The ONNX export produces moments [B, 128, T] (mean+logvar concatenated).
    We split and sample: latent = mean + exp(0.5 * logvar) * noise,
    matching the VAE's latent_dist.sample() behavior.

    Inputs longer than the engine's max profile shape are encoded in
    overlapping chunks and stitched at latent-frame granularity: each
    chunk carries ``_VAE_ENCODE_CHUNK_MARGIN_FRAMES`` of audio context
    per side, and only the interior core frames are kept. This lets a
    small (e.g. 60 s) encode engine serve arbitrarily long sources, so
    the profile manager never has to swap in the 120/240 s encode
    engines and their multi-GB workspace reservations.
    """
    entry = _get_trt_vae(engine_path, device)

    dtypes = entry["tensor_dtypes"]
    inp = audio_bct.to(device=device, dtype=dtypes.get("audio", torch.float32)).contiguous()

    # Release PyTorch's unused reserved VRAM before TRT encode.
    torch.cuda.empty_cache()

    min_dims, _, max_dims = entry["engine"].get_tensor_profile_shape("audio", 0)
    max_samples = int(max_dims[-1])
    min_samples = int(min_dims[-1])
    n_samples = inp.shape[-1]

    if n_samples <= max_samples:
        moments = _trt_vae_encode_moments(inp, entry, engine_path)
    else:
        spf = _VAE_SAMPLES_PER_FRAME
        plan = _plan_encode_chunks(n_samples, max_samples, min_samples)
        total_frames = (n_samples + spf - 1) // spf
        logger.info(
            "trt_vae_encode_chunked samples={} engine_max={} chunks={}",
            n_samples, max_samples, len(plan),
        )
        pieces: list[torch.Tensor] = []
        for ctx_start, ctx_end, core_start, core_end in plan:
            seg = inp[..., ctx_start * spf : min(ctx_end * spf, n_samples)]
            m = _trt_vae_encode_moments(seg.contiguous(), entry, engine_path)
            keep_from = core_start - ctx_start
            if core_end >= total_frames:
                # Tail chunk: keep whatever the engine produced past the
                # margin so a partial trailing frame follows the engine's
                # own rounding rather than our ceil().
                pieces.append(m[..., keep_from:])
            else:
                pieces.append(m[..., keep_from : keep_from + (core_end - core_start)])
        moments = torch.cat(pieces, dim=-1)

    # Split moments into mean and logvar, sample
    mean, logvar = moments.float().chunk(2, dim=1)  # [B, 64, T] each
    std = torch.exp(0.5 * logvar)
    latent = mean + std * torch.randn_like(mean)
    return latent


def _trt_vae_encode_window(
    audio_bct: torch.Tensor,
    engine_path: str,
    device: torch.device,
    *,
    sample: bool = True,
) -> torch.Tensor:
    """Encode a fixed window [B, 2, samples] -> latents [B, 64, T] via TRT.

    The windowed analog of :func:`_trt_vae_encode`, used by the
    real-time-input splice path with the fixed 5 s encode engine. Two
    deliberate differences from :func:`_trt_vae_encode`:

      * No ``torch.cuda.empty_cache()``. That call exists in the full
        encode path to release PyTorch's reserved pool before the big
        ranged engine's ~16 GiB workspace alloc. The fixed 5 s engine
        reserves a tiny workspace, and emptying the cache on every live
        splice would thrash the allocator and disturb the concurrent
        DiT ticks (the realtime probe showed the streaming loop is
        duty-cycle sensitive).
      * ``sample`` toggles the ``mean + std*randn`` draw. Default True
        matches how the playback source latent was created (so a spliced
        region carries the same sampling statistics as its surroundings).
        ``sample=False`` returns the deterministic distribution mean, for
        parity checks / change-detection where sampling noise (rms
        ~0.022) would mask the signal.
    """
    entry = _get_trt_vae(engine_path, device)
    ctx = entry["context"]
    stream = _get_trt_stream()

    dtypes = entry["tensor_dtypes"]
    inp = audio_bct.to(
        device=device, dtype=dtypes.get("audio", torch.float32)
    ).contiguous()

    if not ctx.set_input_shape("audio", tuple(inp.shape)):
        logger.error(
            "trt_vae_encode_window_rejected_input_shape input_shape={} engine={}",
            tuple(inp.shape), engine_path,
        )
        raise RuntimeError(
            f"TRT VAE encode rejected input shape: {tuple(inp.shape)}"
        )
    if not ctx.set_tensor_address("audio", inp.data_ptr()):
        raise RuntimeError("TRT VAE encode rejected input address")

    missing = ctx.infer_shapes()
    if missing:
        raise RuntimeError(
            f"TRT VAE encode shapes insufficiently specified: {missing}"
        )

    out_shape = tuple(ctx.get_tensor_shape("moments"))
    moments_dtype = dtypes.get("moments", torch.float32)
    moments_buf = torch.empty(out_shape, dtype=moments_dtype, device=device)
    if not ctx.set_tensor_address("moments", moments_buf.data_ptr()):
        raise RuntimeError("TRT VAE encode rejected output address")
    if not ctx.execute_async_v3(stream.ptr):
        raise RuntimeError("TRT VAE encode failed")
    stream.synchronize()

    mean, logvar = moments_buf.float().chunk(2, dim=1)  # [B, 64, T] each
    if not sample:
        return mean
    std = torch.exp(0.5 * logvar)
    return mean + std * torch.randn_like(mean)


def _find_trt_engine(name: str) -> Optional[str]:
    """Search for a TRT engine file.

    Checks the central models directory (from acestep.paths) first,
    then falls back to relative paths for backward compatibility.
    """
    from acestep.paths import trt_engines_dir

    stem = name.replace(".engine", "")
    trt_dir = str(trt_engines_dir())
    candidates = [
        # Central models directory (preferred)
        os.path.join(trt_dir, stem, name),
        # CWD relative (legacy fallback)
        os.path.join("trt_engines", stem, name),
    ]
    for c in candidates:
        p = os.path.abspath(c)
        if os.path.exists(p):
            return p
    return None


def _trt_vae_profile_fits(
    engine_path: str, tensor_name: str, shape,
) -> Optional[bool]:
    """Whether ``shape`` satisfies the cached engine's optimization profile.

    The module-level TRT VAE cache is shared process-wide, so a cached
    engine may belong to a session whose profile doesn't cover another
    caller's input — e.g. the live streaming session's 60 s
    ``vae_encode`` engine vs a 120 s upload going through the eager
    upload-encoder. Callers use this to decide between the cached TRT
    engine and an eager fallback BEFORE committing to the TRT path.

    Returns ``True``/``False`` when the profile verdict is known, or
    ``None`` when it can't be determined (engine not cached, API error)
    — treat ``None`` as "behave as before" (use TRT). Every
    optimization profile is checked: a multi-profile engine fits when
    ANY profile covers the shape.
    """
    try:
        entry = _trt_vae_cache.get(os.path.abspath(engine_path))
        if entry is None or entry.get("engine") is None:
            return None
        engine = entry["engine"]
        num_profiles = int(getattr(engine, "num_optimization_profiles", 1) or 1)
        for profile in range(num_profiles):
            mn, _opt, mx = engine.get_tensor_profile_shape(tensor_name, profile)
            if len(mn) != len(shape):
                continue
            if all(
                int(mn[i]) <= int(shape[i]) <= int(mx[i])
                for i in range(len(shape))
            ):
                return True
        return False
    except Exception as exc:
        logger.warning(
            "trt_vae_profile_check_failed engine={} tensor={} error={}",
            engine_path, tensor_name, exc,
        )
        return None


def _trt_vae_encode_fits_or_chunkable(engine_path: str, shape) -> Optional[bool]:
    """Whether the cached encode engine can serve ``shape``, directly or chunked.

    Chunk-aware variant of :func:`_trt_vae_profile_fits` for the encode
    path. A samples dim above the profile max is NOT a mismatch here:
    :func:`_trt_vae_encode` serves it via overlapping chunks, which is
    the whole point of pinning the small encode engine. Only
    non-chunkable mismatches reject — rank, batch/channel out of range,
    or samples below the profile min. Same ``None`` semantics as
    ``_trt_vae_profile_fits`` (engine not cached, API error → "behave
    as before", i.e. use TRT).
    """
    try:
        entry = _trt_vae_cache.get(os.path.abspath(engine_path))
        if entry is None or entry.get("engine") is None:
            return None
        engine = entry["engine"]
        num_profiles = int(getattr(engine, "num_optimization_profiles", 1) or 1)
        for profile in range(num_profiles):
            mn, _opt, mx = engine.get_tensor_profile_shape("audio", profile)
            if len(mn) != len(shape):
                continue
            if not all(
                int(mn[i]) <= int(shape[i]) <= int(mx[i])
                for i in range(len(shape) - 1)
            ):
                continue
            n_samples = int(shape[-1])
            if n_samples < int(mn[-1]):
                continue
            if n_samples <= int(mx[-1]):
                return True
            # Above the profile max: chunkable iff the max window leaves
            # room for core frames beyond the two per-side margins (the
            # same bound _plan_encode_chunks enforces).
            max_frames = int(mx[-1]) // _VAE_SAMPLES_PER_FRAME
            if max_frames > 2 * _VAE_ENCODE_CHUNK_MARGIN_FRAMES:
                return True
        return False
    except Exception as exc:
        logger.warning(
            "trt_vae_encode_fit_check_failed engine={} error={}",
            engine_path, exc,
        )
        return None


def _find_best_vae_engine(component: str) -> Optional[str]:
    """Return a TRT VAE engine path for *component* if one was preloaded.

    Only returns engines already in `_trt_vae_cache` (populated by
    `Session(vae_backend="tensorrt", trt_engines={...})`). There is no
    filesystem discovery: TRT is opt-in, never magic.

    Matches both ``vae_<action>_*`` and ``dreamvae_<action>_*`` for the
    same component, since the distilled DreamVAE engines are drop-in
    replacements for the standard decoder.

    Args:
        component: "vae_decode" or "vae_encode"
    """
    # component is "vae_decode" or "vae_encode" -> action is "decode"/"encode".
    action = component.split("_", 1)[-1]
    accepted = (f"vae_{action}_", f"dreamvae_{action}_")
    for cached_path in _trt_vae_cache:
        basename = os.path.basename(cached_path).lower()
        if basename.startswith(accepted):
            return cached_path
    return None



# -----------------------------------------------------------------------
# Nodes
# -----------------------------------------------------------------------

@NodeRegistry.register
class VAEEncodeAudio(BaseNode):
    """Encode audio waveform to latent space.

    Uses TRT engine if available (vae_encode_fp16.engine), falls back
    to PyTorch VAE via the handler.
    """

    node_type_id: ClassVar[str] = "acestep.VAEEncodeAudio"

    @classmethod
    def get_definition(cls) -> NodeDefinition:
        return NodeDefinition(
            node_type_id=cls.node_type_id,
            display_name="VAE Encode Audio",
            category="vae",
            description="Encode audio waveform to latent representation.",
            inputs=(
                NodePort(name="vae", type="VAE"),
                NodePort(name="audio", type="AUDIO"),
            ),
            outputs=(
                NodePort(name="latent", type="LATENT"),
            ),
        )

    def execute(self, **kwargs: Any) -> dict[str, Any]:
        vae: VAEHandle = kwargs["vae"]
        audio: Audio = kwargs["audio"]
        handler = vae.handler
        device = torch.device(handler.device)
        dtype = handler.dtype

        waveform = audio.waveform
        # Ensure [B, C, samples]
        if waveform.dim() == 2:
            waveform = waveform.unsqueeze(0)

        trt_path = _find_best_vae_engine("vae_encode") if _trt_available() else None
        if (
            trt_path
            and handler.vae is not None
            and _trt_vae_encode_fits_or_chunkable(trt_path, tuple(waveform.shape))
            is False
        ):
            # The cached engine can't take this input directly OR via
            # the chunked encode path (rank/batch/channel mismatch, or
            # input below the profile min). This handler carries an
            # eager VAE — use it instead of letting TRT reject the
            # shape. Inputs that merely exceed the profile max stay on
            # TRT: _trt_vae_encode chunks them.
            logger.info(
                "vae_encode_trt_profile_mismatch input_shape={} engine={} "
                "fallback=eager",
                tuple(waveform.shape), os.path.basename(trt_path),
            )
            trt_path = None
        if trt_path:
            logger.debug("VAE encode via TRT")
            try:
                latents_bdt = _trt_vae_encode(waveform, trt_path, device)
            except RuntimeError as exc:
                if handler.vae is None or "rejected input shape" not in str(exc):
                    raise
                # Cold-cache twin of the profile-fit guard above: when
                # the engine wasn't cached yet the check returned None
                # ("use TRT"), the engine loaded just now and rejected
                # the shape. Same hazard, same answer — this handler
                # carries an eager VAE.
                logger.info(
                    "vae_encode_trt_shape_rejected input_shape={} engine={} "
                    "fallback=eager",
                    tuple(waveform.shape), os.path.basename(trt_path),
                )
                trt_path = None
            else:
                # [B, D, T] -> [B, T, D]
                latents = latents_bdt.transpose(1, 2).to(dtype)
        if not trt_path:
            logger.debug("VAE encode via PyTorch")
            with handler._load_model_context("vae"):
                latents = handler._encode_audio_to_latents(waveform)
            if latents.dim() == 2:
                latents = latents.unsqueeze(0)

        return {"latent": Latent(tensor=latents)}


@NodeRegistry.register
class VAEDecodeAudio(BaseNode):
    """Decode latents back to audio waveform.

    Uses TRT engine if available (vae_decode_fp16.engine), falls back
    to PyTorch VAE via the handler.
    """

    node_type_id: ClassVar[str] = "acestep.VAEDecodeAudio"

    @classmethod
    def get_definition(cls) -> NodeDefinition:
        return NodeDefinition(
            node_type_id=cls.node_type_id,
            display_name="VAE Decode Audio",
            category="vae",
            description="Decode latent representation to audio waveform.",
            inputs=(
                NodePort(name="vae", type="VAE"),
                NodePort(name="latent", type="LATENT"),
            ),
            outputs=(
                NodePort(name="audio", type="AUDIO"),
            ),
        )

    def execute(self, **kwargs: Any) -> dict[str, Any]:
        vae: VAEHandle = kwargs["vae"]
        latent: Latent = kwargs["latent"]
        handler = vae.handler
        device = torch.device(handler.device)

        # [B, T, D] -> [B, D, T]
        lat_bdt = latent.tensor.transpose(1, 2)

        trt_path = _find_best_vae_engine("vae_decode") if _trt_available() else None
        if (
            trt_path
            and handler.vae is not None
            and _trt_vae_profile_fits(trt_path, "latents", tuple(lat_bdt.shape)) is False
        ):
            # Same cross-session profile hazard as the encode node: an
            # eager-VAE handler must not be forced through another
            # session's cached engine when the shape can't fit it.
            logger.info(
                "vae_decode_trt_profile_mismatch input_shape={} engine={} "
                "fallback=eager",
                tuple(lat_bdt.shape), os.path.basename(trt_path),
            )
            trt_path = None
        if trt_path:
            # debug, not info: this node can run per decode window.
            logger.debug("VAE decode via TRT")
            try:
                waveform = _trt_vae_decode(lat_bdt, trt_path, device)
            except RuntimeError as exc:
                if handler.vae is None or "rejected input shape" not in str(exc):
                    raise
                # Cold-cache twin of the profile-fit guard above (see
                # the encode node).
                logger.info(
                    "vae_decode_trt_shape_rejected input_shape={} engine={} "
                    "fallback=eager",
                    tuple(lat_bdt.shape), os.path.basename(trt_path),
                )
                trt_path = None
        if not trt_path:
            logger.debug("VAE decode via PyTorch (no TRT engine found)")
            with handler._load_model_context("vae"):
                waveform = handler.tiled_decode(lat_bdt)

        return {"audio": Audio(waveform=waveform, sample_rate=48000)}


@NodeRegistry.register
class StreamVAEDecode(BaseNode):
    """Windowed VAE decode for streaming playback.

    Wraps :class:`VAEDecodeAudio` with the overlap-margin windowing that
    the realtime demo needs. The finished latent is typically long
    (60 s); decoding a short interior window at ``t_start`` keeps tick
    latency low while the margins supply receptive-field context so the
    window edges match the full decode.

    Realtime skip gate
    ------------------
    When the host supplies a per-call ``playhead_seconds`` kwarg (scope
    injects this from ``AudioProcessingTrack._timestamp / 48000``), the
    node follows the playhead instead of the widget ``t_start`` and
    applies the same skip logic as
    ``demos/realtime_motion_graph_web/pipeline.py``:

      * Decode is skipped when the current latent is close to the last
        one (``mse < mse_skip_threshold``) **and** the playhead has not
        yet advanced past ``(vae_window - prefetch_seconds)`` since the
        last actual decode.
      * When skipping, the node returns ``{"audio": None}`` so scope's
        ``_route_outputs`` short-circuits — the downstream sink is not
        touched and the upstream ring buffer is not back-pressured.

    Both gates are opt-in: the classic, stateless decode path is
    preserved when ``follow_playhead`` is false or no playhead is
    provided.

    Node parameters:
        vae_window: Kept "wire slice" length in seconds. ``<= 0`` decodes
            the full latent in one shot (parity with ``VAEDecodeAudio``).
            Otherwise the fixed 1 s engine decodes a 25-frame span and the
            middle ``vae_window`` seconds are returned.
        vae_overlap: Legacy no-op. The per-side margin is derived from the
            fixed decode span, so this value is ignored.
        t_start: Window start in seconds. Used only when
            ``follow_playhead`` is false or no playhead is provided.
        follow_playhead: If true and ``playhead_seconds`` is supplied in
            kwargs, drive ``t_start`` from the playhead and enable the
            skip gate.
        mse_skip_threshold: MSE threshold for the skip gate. Matches the
            demo's ``skip_threshold`` (default 1e-3). ``<= 0`` disables
            the MSE check but leaves the prefetch-window check active.
        prefetch_seconds: Playhead-advance cushion. Demo uses
            ``min(1.0, vae_window * 0.2)`` — we default to 1.0.
    """

    node_type_id: ClassVar[str] = "acestep.StreamVAEDecode"

    FRAMES_PER_SEC: ClassVar[int] = 25
    SAMPLES_PER_FRAME: ClassVar[int] = 1920  # 48_000 / 25

    def __init__(self):
        # Skip-gate state. Mirrors the demo's ``last_latent`` and
        # ``last_decode_pos`` locals, hoisted to instance scope because
        # Scope's graph runtime holds the node instance across ticks.
        self._last_latent_for_skip: Optional[torch.Tensor] = None
        self._last_decode_pos: Optional[float] = None

    @classmethod
    def get_definition(cls) -> NodeDefinition:
        return NodeDefinition(
            node_type_id=cls.node_type_id,
            display_name="Stream VAE Decode",
            category="vae",
            description="Windowed VAE decode for realtime streaming playback.",
            inputs=(
                NodePort(name="vae", type="VAE"),
                NodePort(name="latent", type="LATENT"),
            ),
            outputs=(
                NodePort(name="audio", type="AUDIO"),
            ),
            params=(
                NodeParam(
                    name="vae_window", type="number", default=0.36,
                    description=(
                        "Kept 'wire slice' (s); <=0 decodes full latent. "
                        "The windowed engine always decodes a fixed 1.0 s "
                        "span and this is the middle portion emitted; "
                        "clamped to [0.04, 0.36] so >=8 frames of margin "
                        "remain on each side."
                    ),
                    min=0.0, max=0.36, step=0.02,
                ),
                NodeParam(
                    name="vae_overlap", type="number", default=0.4,
                    description=(
                        "Legacy no-op. The fixed 1 s windowed engine derives "
                        "its per-side margin as (1 s - vae_window)/2, so this "
                        "value is ignored."
                    ),
                    min=0.0, max=5.0, step=0.05,
                ),
                NodeParam(
                    name="t_start", type="number", default=0.0,
                    description="Window start (s); ignored when follow_playhead is on",
                    min=0.0, max=60.0, step=0.1,
                ),
                NodeParam(
                    name="follow_playhead", type="boolean", default=False,
                    description="Drive t_start from audio sink playhead + enable skip gate",
                ),
                NodeParam(
                    name="mse_skip_threshold", type="number", default=1e-3,
                    description="Skip decode when latent MSE < threshold",
                    min=0.0, max=1.0, step=1e-4,
                ),
                NodeParam(
                    name="prefetch_seconds", type="number", default=1.0,
                    description="Seconds of headroom before forcing a re-decode",
                    min=0.0, max=5.0, step=0.05,
                ),
                NodeParam(
                    name="playhead_seconds", type="any", default=None,
                    description="Ambient audio playhead injected by host",
                    hidden=True,
                ),
                NodeParam(
                    name="cyclic", type="boolean", default=False,
                    description=(
                        "Treat the latent as cyclic: pull missing left/right "
                        "context from the opposite end of the latent instead "
                        "of clamping. Use only when the audio is intended to "
                        "loop, otherwise the song's tail will bleed into the "
                        "head's decode through the VAE receptive field."
                    ),
                ),
            ),
        )

    def execute(self, **kwargs: Any) -> dict[str, Any]:
        from acestep.paths import WINDOWED_VAE_WINDOW_RANGE_S

        vae: VAEHandle = kwargs["vae"]
        latent: Latent = kwargs["latent"]
        window_s = float(kwargs.get("vae_window", 0.0))

        # Hard-clamp positive windows to the engine-supported range.
        # ``<= 0`` is the disable sentinel and falls through to a full
        # decode, so we leave it untouched.
        if window_s > 0:
            lo, hi = WINDOWED_VAE_WINDOW_RANGE_S
            window_s = max(lo, min(hi, window_s))

        follow_playhead = bool(kwargs.get("follow_playhead", False))
        playhead_seconds = kwargs.get("playhead_seconds")
        mse_skip_threshold = float(kwargs.get("mse_skip_threshold", 1e-3))
        prefetch_seconds = float(kwargs.get("prefetch_seconds", 1.0))
        cyclic = bool(kwargs.get("cyclic", False))

        use_playhead = (
            follow_playhead
            and playhead_seconds is not None
            and window_s > 0
        )

        if use_playhead:
            t_start = float(playhead_seconds)
        else:
            t_start = float(kwargs.get("t_start", 0.0))

        # -- Skip gate (demo pipeline.py:215-225) --------------------
        # Only active when following the playhead. Skipping returns
        # ``{"audio": None}``; scope's _route_outputs drops None so no
        # chunk reaches the sink and upstream stays unblocked.
        if use_playhead:
            tensor = latent.tensor
            mse_ok_to_skip = True
            if mse_skip_threshold > 0.0 and self._last_latent_for_skip is not None:
                try:
                    mse = (tensor - self._last_latent_for_skip).pow(2).mean().item()
                except RuntimeError:
                    # Shape changed — don't skip; fall through to decode.
                    mse = float("inf")
                mse_ok_to_skip = mse < mse_skip_threshold

            within_window = (
                self._last_decode_pos is not None
                and abs(t_start - self._last_decode_pos) < (window_s - prefetch_seconds)
            )

            if mse_ok_to_skip and within_window:
                # Keep the reference latent fresh so a future divergence
                # compares against what we actually last decoded.
                return {"audio": None}

        if window_s <= 0:
            out = VAEDecodeAudio().execute(vae=vae, latent=latent)
            if use_playhead:
                self._last_latent_for_skip = latent.tensor.detach().clone()
                self._last_decode_pos = t_start
            return out

        tensor = latent.tensor
        T = tensor.shape[1]

        # The windowed VAE engine is a FIXED-shape engine: it decodes
        # exactly ``decode_frames`` (1 s / 25 frames) on every call. The
        # kept "wire slice" is the middle ``win_frames`` of that span; the
        # leftover frames each side are the receptive-field margin we trim
        # off. The margin is derived from ``decode_frames - keep`` so live
        # controls can never feed an out-of-profile shape — ``vae_overlap``
        # is not used here. See acestep.paths WINDOWED_VAE_PROFILE_FRAMES.
        from acestep.paths import WINDOWED_VAE_PROFILE_FRAMES
        decode_frames = WINDOWED_VAE_PROFILE_FRAMES[0]
        # Keep can't exceed the fixed decode span; clamp defensively.
        win_frames = max(1, min(int(window_s * self.FRAMES_PER_SEC), decode_frames))

        if T <= win_frames:
            out = VAEDecodeAudio().execute(vae=vae, latent=latent)
            if use_playhead:
                self._last_latent_for_skip = tensor.detach().clone()
                self._last_decode_pos = t_start
            return out

        # Playhead-follow clamp: the demo does
        # ``t_pos = min(t_pos, eff_dur - vae_window)`` to keep the
        # window inside the latent; frame-quantization below already
        # handles the floor. In cyclic mode we let the playhead reach
        # the end and rely on wrapped context for the right margin.
        if use_playhead and not cyclic:
            max_t = max(0.0, (T / self.FRAMES_PER_SEC) - window_s)
            t_start = min(t_start, max_t)

        keep_start = max(0, int(t_start * self.FRAMES_PER_SEC))
        keep_end = min(T, keep_start + win_frames)
        keep_start = max(0, keep_end - win_frames)

        # Decode a fixed ``decode_frames``-frame span centered on the keep
        # window; the kept slice is the middle ``win_frames`` and its trim
        # offset is ``keep_start - decode_start``.
        left = (decode_frames - win_frames) // 2
        if cyclic:
            # Take the fixed span wrapping around the latent so a looping
            # song's tail feeds the head's left context.
            start = keep_start - left
            idx = torch.arange(
                start, start + decode_frames, device=tensor.device
            ) % T
            chunk_lat = Latent(tensor=tensor.index_select(1, idx).contiguous())
            pre_margin_frames = left
        else:
            decode_start = max(0, min(keep_start - left, T - decode_frames))
            decode_end = decode_start + decode_frames
            chunk_lat = Latent(
                tensor=tensor[:, decode_start:decode_end, :].contiguous()
            )
            pre_margin_frames = keep_start - decode_start

        chunk_audio = VAEDecodeAudio().execute(vae=vae, latent=chunk_lat)["audio"]

        pre_margin = pre_margin_frames * self.SAMPLES_PER_FRAME
        keep_samples = (keep_end - keep_start) * self.SAMPLES_PER_FRAME
        trimmed = chunk_audio.waveform[:, :, pre_margin:pre_margin + keep_samples]

        if use_playhead:
            self._last_latent_for_skip = tensor.detach().clone()
            self._last_decode_pos = keep_start / self.FRAMES_PER_SEC

        return {
            "audio": Audio(
                waveform=trimmed,
                sample_rate=chunk_audio.sample_rate,
                start_sample=keep_start * self.SAMPLES_PER_FRAME,
            ),
        }


@NodeRegistry.register
class EmptyLatent(BaseNode):
    """Create a silence-based latent of a given duration.

    Node parameters:
        duration: Duration in seconds.
    """

    node_type_id: ClassVar[str] = "acestep.EmptyLatent"

    @classmethod
    def get_definition(cls) -> NodeDefinition:
        return NodeDefinition(
            node_type_id=cls.node_type_id,
            display_name="Empty ACE-Step Latent",
            category="vae",
            description="Create an empty (silence) latent for a given duration.",
            inputs=(
                NodePort(name="model", type="MODEL"),
            ),
            outputs=(
                NodePort(name="latent", type="LATENT"),
            ),
            params=(
                NodeParam(
                    name="duration", type="number", default=60.0,
                    description="Duration (s)",
                    min=1.0, max=600.0, step=1.0,
                ),
                NodeParam(
                    name="frames", type="number", default=0,
                    description=(
                        "Exact latent frame count (0 = derive from duration). "
                        "Set by the streaming runner to avoid the "
                        "duration round-trip dropping a frame."
                    ),
                    hidden=True,
                ),
            ),
        )

    def execute(self, **kwargs: Any) -> dict[str, Any]:
        model: ModelHandle = kwargs["model"]
        handler = model.handler
        duration = kwargs.get("duration", 60.0)

        handler._ensure_silence_latent_on_device()
        silence = handler.silence_latent  # [1, T_full, D]

        # Frame count: prefer an explicit integer ``frames`` (exact — used by
        # the streaming runner so the silence latent matches the source
        # latent's T exactly). Fall back to deriving it from ``duration``,
        # using round() rather than truncation: int(duration * 25) drops a
        # frame for durations whose frame count round-trips just under an
        # integer (e.g. 410 frames -> 16.4s -> int(409.9999...) == 409),
        # which then can't blend against a real 410-frame context latent.
        frames = kwargs.get("frames") or 0
        if frames:
            T = int(frames)
        else:
            T = int(round(duration * 25))  # 25 fps latent rate
        # Take first T frames from the silence latent (tiled if needed)
        if silence.dim() == 3:
            latent = silence[:, :T, :].clone()
            if latent.shape[1] < T:
                reps = (T + latent.shape[1] - 1) // latent.shape[1]
                latent = latent.repeat(1, reps, 1)[:, :T, :]
        else:
            # Fallback: treat as [1, D] single frame
            latent = silence.unsqueeze(0).expand(1, T, -1).clone()

        return {"latent": Latent(tensor=latent)}


@NodeRegistry.register
class LatentBlend(BaseNode):
    """Blend two latents by weighted interpolation.

    Supports scalar or per-frame (CURVE) blend factor.
    Useful for timbre strength control (blend reference with silence)
    or mixing any two latent representations.

    Node parameters:
        alpha: Blend factor (0.0 = all A, 1.0 = all B).
               Ignored if a blend_curve input is connected.
        method: "linear" (straight average) or "slerp" (per-frame
               geodesic, holds the blended latent's norm constant).
    """

    node_type_id: ClassVar[str] = "acestep.LatentBlend"

    @classmethod
    def get_definition(cls) -> NodeDefinition:
        return NodeDefinition(
            node_type_id=cls.node_type_id,
            display_name="Latent Blend",
            category="vae",
            description="Blend two latents with scalar or per-frame factor.",
            inputs=(
                NodePort(name="latent_a", type="LATENT"),
                NodePort(name="latent_b", type="LATENT"),
                NodePort(
                    name="blend_curve",
                    type="CURVE",
                    required=False,
                    description="Per-frame blend factor (overrides scalar alpha).",
                ),
            ),
            outputs=(
                NodePort(name="latent", type="LATENT"),
            ),
            params=(
                NodeParam(
                    name="alpha", type="number", default=0.5,
                    description="Blend factor (0 = all A, 1 = all B). Overridden by blend_curve.",
                    min=0.0, max=1.0, step=0.01,
                ),
                NodeParam(
                    name="method", type="select", default="linear",
                    description=(
                        "Interpolation method (see acestep.nodes.interpolation). "
                        "'linear' (default) averages the two latents. 'slerp' "
                        "walks the per-frame geodesic so the blended latent's "
                        "norm stays constant across the sweep instead of dipping "
                        "at the midpoint."
                    ),
                    options=INTERP_METHOD_NAMES,
                ),
            ),
        )

    def execute(self, **kwargs: Any) -> dict[str, Any]:
        latent_a: Latent = kwargs["latent_a"]
        latent_b: Latent = kwargs["latent_b"]
        blend_curve: Optional[Curve] = kwargs.get("blend_curve")
        method = str(kwargs.get("method", "linear"))

        a = latent_a.tensor
        b = latent_b.tensor

        alpha = kwargs.get("alpha", 0.5)
        if blend_curve is not None:
            alpha = blend_curve.tensor.to(device=a.device, dtype=a.dtype)
            if alpha.ndim == 1:
                alpha = alpha.unsqueeze(0).unsqueeze(-1)  # [1, T, 1]
            elif alpha.ndim == 2:
                alpha = alpha.unsqueeze(-1)  # [B, T, 1]

        blended = INTERPOLATIONS[method](a, b, alpha)
        return {"latent": Latent(tensor=blended)}
