"""engine.py — Stable Audio 3 Small Music continuation core.

CODA's job is one thing done well: take a short, unfinished clip and continue it
into a finished-sounding track in the same key, tempo and feel. SA3 does that in
a SINGLE call. Its `generate_diffusion_cond_inpaint` is a native audio-inpainting
diffusion sampler: place the user's clip at the front of the buffer, mask the
region after it, and the model fills the masked region conditioned on the kept
audio — true long-form continuation, 44.1 kHz stereo, no multi-pass chaining,
no energy guards, no re-roll logic.

This module is the whole generation core. It returns ONLY the newly generated
tail (the model's [source_end, total] region) plus the source length in seconds;
`stitch.py` joins that tail onto the user's *pristine* original so the real
recording (and any vocals) plays untouched up to the seam.

Mask convention (verified against the installed library source):
  inpaint_mask = ones(buffer); inpaint_mask[start:end] = 0
  -> 1 = keep the input audio, 0 = generate. So masking [L_src, L_total] keeps
  the source in [0, L_src] and generates everything after it.
"""
import numpy as np
import torch

MODEL_ID = "stabilityai/stable-audio-3-small-music"

SR = 44100               # SA3 native sample rate (model_config: sample_rate)
STEPS = 8                # SA3 Small is an 8-step adversarially-distilled model
SAMPLER = "pingpong"     # the sampler the distilled model was tuned for
DEFAULT_CFG = 1.0        # distilled-model guidance; the prompt still conditions
                         # at 1.0 (CFG amplification off, conditional path on)
MAX_TOTAL_SECONDS = 120  # SA3 Small duration cap (sample_size / sample_rate)
MIN_NEW_SECONDS = 5      # below this a "continuation" isn't worth a GPU call

_model = None
_model_config = None
_sample_size = None
_on_device = None        # which device the weights currently live on


def _device():
    return "cuda" if torch.cuda.is_available() else "cpu"


def preload():
    """Load model + autoencoder + T5Gemma conditioner into CPU RAM at process
    start. On ZeroGPU the per-call GPU window is the scarce resource, so weights
    must come off disk at boot, not inside the window. The CUDA placement + fp16
    cast is deferred to the first `continue_audio` call (i.e. the @spaces.GPU
    window), matching how Stability's own Space defers it."""
    global _model, _model_config, _sample_size
    if _model is None:
        from stable_audio_tools import get_pretrained_model
        _model, _model_config = get_pretrained_model(MODEL_ID)
        _sample_size = int(_model_config["sample_size"])
        _model.eval()
        print(f"[coda] preload: SA3 resident "
              f"(sr={_model_config['sample_rate']}, "
              f"sample_size={_sample_size} = "
              f"{_sample_size / int(_model_config['sample_rate']):.0f}s)",
              flush=True)
    return _model, _model_config


def _ensure_on_device():
    """Ensure weights are on the GPU in fp16. Called inside the @spaces.GPU
    window on every generation. `.to()` is a cheap no-op when the model is
    already placed, so we re-ensure each call rather than caching device state —
    that stays correct even if ZeroGPU detaches the GPU between calls. fp16 is
    only valid on CUDA; on CPU the model stays fp32."""
    global _model, _on_device
    dev = _device()
    _model = _model.to(dev)
    if dev == "cuda":
        _model = _model.to(torch.float16)
    _on_device = dev
    return _model


def _load_source(clip_path):
    """Load the clip as stereo float32 @44.1k as a (2, N) tensor. SA3's
    autoencoder is stereo; `prepare_audio` inside the sampler will pad/crop to
    the buffer length and place this at the FRONT (PadCrop, randomize=False)."""
    import librosa
    y, _ = librosa.load(clip_path, sr=SR, mono=False)
    y = np.asarray(y, dtype=np.float32)
    if y.ndim == 1:
        y = np.stack([y, y])           # mono -> stereo
    elif y.shape[0] > 2:
        y = y[:2]
    return torch.from_numpy(np.ascontiguousarray(y))


#: longest clip that still leaves room for MIN_NEW of continuation under the cap
MAX_SOURCE_SECONDS = MAX_TOTAL_SECONDS - MIN_NEW_SECONDS


def plan_continuation(source_seconds, total_seconds):
    """Pure helper (unit-testable, no model): clamp the request to SA3's limits
    and return (total_seconds, new_seconds, mask_start, mask_end).

    - the mask runs from where the source ends to the total length: that masked
      region is what SA3 generates, so mask_end MUST exceed mask_start.
    - total is capped at MAX_TOTAL_SECONDS and floored so at least
      MIN_NEW_SECONDS of new audio is generated.
    - raises ValueError when the source is already so long there's no room to
      continue under the cap (otherwise the mask would invert and SA3 would
      silently generate nothing).
    """
    source_seconds = float(source_seconds)
    total_seconds = float(total_seconds)
    if source_seconds > MAX_SOURCE_SECONDS:
        raise ValueError(
            f"clip is {source_seconds:.0f}s — too long to continue under SA3's "
            f"{MAX_TOTAL_SECONDS:.0f}s cap (need room for at least "
            f"{MIN_NEW_SECONDS:.0f}s of new audio); trim it under "
            f"{MAX_SOURCE_SECONDS:.0f}s.")
    # source <= MAX_SOURCE_SECONDS, so source + MIN_NEW <= MAX_TOTAL: the floor
    # never pushes total past the cap, and mask_end (total) > mask_start (source).
    total_seconds = min(total_seconds, MAX_TOTAL_SECONDS)
    total_seconds = max(total_seconds, source_seconds + MIN_NEW_SECONDS)
    new_seconds = total_seconds - source_seconds
    return total_seconds, new_seconds, source_seconds, total_seconds


def continue_audio(clip_path, total_seconds, prompt="", cfg_scale=DEFAULT_CFG,
                   seed=-1, progress=None):
    """Continue `clip_path` up to `total_seconds` in one SA3 inpaint call.

    Returns (new_tail, source_seconds, SR) where:
      new_tail       : (2, M) float32 @44.1k — ONLY the generated region
                       [source_end, total]. Peak-normalized to <= 1.0.
      source_seconds : the clip's true length (the splice boundary, in seconds)
      SR             : 44100

    `progress(stage_name)` is called (best-effort) at each stage so the UI can
    paint a live status. SA3 is one diffusion call, so progress is stage-based.
    """
    from einops import rearrange
    from stable_audio_tools.inference.generation import (
        generate_diffusion_cond_inpaint)

    def _notify(stage):
        if progress is not None:
            try:
                progress(stage)
            except Exception as e:
                print(f"[coda] progress callback failed ({e})", flush=True)

    preload()
    model = _ensure_on_device()
    dev = _device()

    # The library does `np.random.randint(0, 2**32-1)` when seed == -1, which
    # overflows int32 on Windows/numpy<2. Draw a safe seed ourselves so the
    # default path works everywhere, not just on the Linux Space.
    if seed is None or seed < 0:
        seed = int(np.random.randint(0, 2 ** 31 - 1))

    _notify("reading")
    source = _load_source(clip_path)
    source_seconds = source.shape[-1] / SR
    # the autoencoder runs in the model's dtype (fp16 on CUDA); the conditioning
    # audio must match it or the encoder's conv1d rejects the input dtype.
    model_dtype = next(model.model.parameters()).dtype
    source = source.to(model_dtype)

    total_seconds, new_seconds, mask_start, mask_end = plan_continuation(
        source_seconds, total_seconds)
    prompt = (prompt or "").strip()

    print(f"[coda] continuation: source={source_seconds:.1f}s -> "
          f"total={total_seconds:.1f}s (+{new_seconds:.1f}s new), "
          f"mask=[{mask_start:.1f}s, {mask_end:.1f}s], steps={STEPS}, "
          f"cfg={cfg_scale}, prompt={prompt!r}", flush=True)

    _notify("composing")
    with torch.no_grad():
        output = generate_diffusion_cond_inpaint(
            model,
            steps=STEPS,
            cfg_scale=cfg_scale,
            conditioning=[{"prompt": prompt, "seconds_total": total_seconds}],
            sample_size=_sample_size,
            sampler_type=SAMPLER,
            inpaint_audio=(SR, source),
            inpaint_mask_start_seconds=mask_start,
            inpaint_mask_end_seconds=mask_end,
            seed=seed,
            device=dev,
        )

    _notify("finalizing")
    # (b, d, n) -> (d, b*n); peak-normalize like Stability's reference Space
    output = rearrange(output, "b d n -> d (b n)")
    audio = output.to(torch.float32).cpu().numpy()
    peak = float(np.abs(audio).max())
    if peak > 1e-9:
        audio = audio / peak

    if audio.shape[0] == 1:            # safety: ensure stereo
        audio = np.repeat(audio, 2, axis=0)

    boundary = int(round(source_seconds * SR))
    end = int(round(total_seconds * SR))
    end = min(end, audio.shape[-1])
    new_tail = audio[:, boundary:end]
    new_tail = np.ascontiguousarray(new_tail.astype(np.float32))

    print(f"[coda] generated tail: shape={new_tail.shape} "
          f"({new_tail.shape[-1] / SR:.1f}s), peak after norm "
          f"{float(np.abs(new_tail).max()):.3f}, "
          f"rms {float(np.sqrt(np.mean(new_tail ** 2))):.3f}", flush=True)
    return new_tail, source_seconds, SR