jam_worker.py

# jam_worker.py - Bar-locked spool rewrite
from __future__ import annotations

import os

import threading, time
from dataclasses import dataclass
from fractions import Fraction
from typing import Optional, Dict, Tuple, List

import numpy as np
from magenta_rt import audio as au

from utils import (
    StreamingResampler,
    match_loudness_to_reference,
    make_bar_aligned_context,
    take_bar_aligned_tail,
    wav_bytes_base64,
)

def _dbg_rms_dbfs(x: np.ndarray) -> float:
    
    if x.ndim == 2:
        x = x.mean(axis=1)
    r = float(np.sqrt(np.mean(x * x) + 1e-12))
    return 20.0 * np.log10(max(r, 1e-12))

def _dbg_rms_dbfs_model(x: np.ndarray) -> float:
    # x is model-rate, shape [S,C] or [S]
    
    if x.ndim == 2:
        x = x.mean(axis=1)
    r = float(np.sqrt(np.mean(x * x) + 1e-12))
    return 20.0 * np.log10(max(r, 1e-12))

def _dbg_shape(x):
    return tuple(x.shape) if hasattr(x, "shape") else ("-",)

# -----------------------------
# Data classes
# -----------------------------

@dataclass
class JamParams:
    bpm: float
    beats_per_bar: int
    bars_per_chunk: int
    target_sr: int
    loudness_mode: str = "auto"
    headroom_db: float = 1.0
    style_vec: Optional[np.ndarray] = None
    ref_loop: Optional[au.Waveform] = None
    combined_loop: Optional[au.Waveform] = None
    guidance_weight: float = 1.1
    temperature: float = 1.1
    topk: int = 40
    style_ramp_seconds: float = 8.0  # 0 => instant (current behavior), try 6.0–10.0 for gentle glides


@dataclass
class JamChunk:
    index: int
    audio_base64: str
    metadata: dict


# -----------------------------
# Helpers
# -----------------------------

class BarClock:
    """Sample-domain bar clock with drift-free absolute boundaries."""
    def __init__(self, target_sr: int, bpm: float, beats_per_bar: int, base_offset_samples: int = 0):
        self.sr = int(target_sr)
        self.bpm = Fraction(str(bpm))  # exact decimal to avoid FP drift
        self.beats_per_bar = int(beats_per_bar)
        self.bar_samps = Fraction(self.sr * 60 * self.beats_per_bar, 1) / self.bpm
        self.base = int(base_offset_samples)

    def bounds_for_chunk(self, chunk_index: int, bars_per_chunk: int) -> Tuple[int, int]:
        start_f = self.base + self.bar_samps * (chunk_index * bars_per_chunk)
        end_f   = self.base + self.bar_samps * ((chunk_index + 1) * bars_per_chunk)
        return int(round(start_f)), int(round(end_f))

    def seconds_per_bar(self) -> float:
        return float(self.beats_per_bar) * (60.0 / float(self.bpm))


# -----------------------------
# Worker
# -----------------------------

class JamWorker(threading.Thread):
    FRAMES_PER_SECOND: float | None = None  # filled in __init__ once codec is available
    """Generates continuous audio with MagentaRT, spools it at target SR,
    and emits *sample-accurate*, bar-aligned chunks (no FPS drift)."""

    def __init__(self, mrt, params: JamParams):
        super().__init__(daemon=True)
        self.mrt = mrt
        self.params = params

        # external callers (FastAPI endpoints) use this for atomic updates
        self._lock = threading.RLock()

        # generation state
        self.state = self.mrt.init_state()
        self.mrt.guidance_weight = float(self.params.guidance_weight)
        self.mrt.temperature     = float(self.params.temperature)
        self.mrt.topk            = int(self.params.topk)

        

        # codec/setup
        self._codec_fps = float(self.mrt.codec.frame_rate)
        JamWorker.FRAMES_PER_SECOND = self._codec_fps
        self._ctx_frames = int(self.mrt.config.context_length_frames)
        self._ctx_seconds = self._ctx_frames / self._codec_fps

        # model stream (model SR) for internal continuity/crossfades
        self._model_stream: Optional[np.ndarray] = None
        self._model_sr = int(self.mrt.sample_rate)

        # style vector (already normalized upstream)
        self._style_vec = (None if self.params.style_vec is None
                   else np.array(self.params.style_vec, dtype=np.float32, copy=True))
        self._chunk_secs = (
            self.mrt.config.chunk_length_frames * self.mrt.config.frame_length_samples
        ) / float(self._model_sr)  # ≈ 2.0 s by default

        # target-SR in-RAM spool (what we cut loops from)
        if int(self.params.target_sr) != int(self._model_sr):
            self._rs = StreamingResampler(self._model_sr, int(self.params.target_sr), channels=2)
        else:
            self._rs = None
        self._spool = np.zeros((0, 2), dtype=np.float32)   # (S,2) target SR
        self._spool_written = 0                            # absolute frames written into spool

        self._pending_tail_model = None      # type: Optional[np.ndarray]  # last tail at model SR
        self._pending_tail_target_len = 0    # number of target-SR samples last tail contributed

        # bar clock: start with offset 0; if you have a downbeat estimator, set base later
        self._bar_clock = BarClock(self.params.target_sr, self.params.bpm, self.params.beats_per_bar, base_offset_samples=0)

        # emission counters
        self.idx = 0  # next chunk index to *produce*
        self._next_to_deliver = 0  # next chunk index to hand out via get_next_chunk()
        self._last_consumed_index = -1  # updated via mark_chunk_consumed(); generation throttle uses this

        # outbox and synchronization
        self._outbox: Dict[int, JamChunk] = {}
        self._cv = threading.Condition()

        # control flags
        self._stop_event = threading.Event()
        self._max_buffer_ahead = 1

        # reseed queues (install at next bar boundary after emission)
        self._pending_reseed: Optional[dict] = None           # legacy full reset path (kept for fallback)
        self._pending_token_splice: Optional[dict] = None     # seamless token splice

        # Prepare initial context from combined loop (best musical alignment)
        if self.params.combined_loop is not None:
            self._install_context_from_loop(self.params.combined_loop)

    # ---------- lifecycle ----------

    def set_buffer_seconds(self, seconds: float):
        """Clamp how far ahead we allow, in *seconds* of audio."""
        chunk_secs = float(self.params.bars_per_chunk) * self._bar_clock.seconds_per_bar()
        max_chunks = max(0, int(round(seconds / max(chunk_secs, 1e-6))))
        with self._cv:
            self._max_buffer_ahead = max_chunks

    def set_buffer_chunks(self, k: int):
        with self._cv:
            self._max_buffer_ahead = max(0, int(k))

    def stop(self):
        self._stop_event.set()

    # FastAPI reads this to block until the next sequential chunk is ready
    def get_next_chunk(self, timeout: float = 30.0) -> Optional[JamChunk]:
        deadline = time.time() + timeout
        with self._cv:
            while True:
                c = self._outbox.get(self._next_to_deliver)
                if c is not None:
                    self._next_to_deliver += 1
                    return c
                remaining = deadline - time.time()
                if remaining <= 0:
                    return None
                self._cv.wait(timeout=min(0.25, remaining))

    def mark_chunk_consumed(self, chunk_index: int):
        # This lets the generator run ahead, but not too far
        with self._cv:
            self._last_consumed_index = max(self._last_consumed_index, int(chunk_index))
            # purge old chunks to cap memory
            for k in list(self._outbox.keys()):
                if k < self._last_consumed_index - 1:
                    self._outbox.pop(k, None)

    def update_knobs(self, *, guidance_weight=None, temperature=None, topk=None):
        with self._lock:
            if guidance_weight is not None:
                self.params.guidance_weight = float(guidance_weight)
            if temperature is not None:
                self.params.temperature = float(temperature)
            if topk is not None:
                self.params.topk = int(topk)
            # push into mrt
            self.mrt.guidance_weight = float(self.params.guidance_weight)
            self.mrt.temperature     = float(self.params.temperature)
            self.mrt.topk            = int(self.params.topk)

    # ---------- context / reseed ----------

    def _expected_token_shape(self) -> Tuple[int, int]:
        F = int(self._ctx_frames)
        D = int(self.mrt.config.decoder_codec_rvq_depth)
        return F, D

    def _coerce_tokens(self, toks: np.ndarray) -> np.ndarray:
        """Force tokens to (context_length_frames, rvq_depth), padding/trimming as needed.
        Pads missing frames by repeating the last frame (safer than zeros for RVQ stacks)."""
        F, D = self._expected_token_shape()
        if toks.ndim != 2:
            toks = np.atleast_2d(toks)
        # depth first
        if toks.shape[1] > D:
            toks = toks[:, :D]
        elif toks.shape[1] < D:
            pad_cols = np.tile(toks[:, -1:], (1, D - toks.shape[1]))
            toks = np.concatenate([toks, pad_cols], axis=1)
        # frames
        if toks.shape[0] < F:
            if toks.shape[0] == 0:
                toks = np.zeros((1, D), dtype=np.int32)
            pad = np.repeat(toks[-1:, :], F - toks.shape[0], axis=0)
            toks = np.concatenate([pad, toks], axis=0)
        elif toks.shape[0] > F:
            toks = toks[-F:, :]
        if toks.dtype != np.int32:
            toks = toks.astype(np.int32, copy=False)
        return toks

    def _encode_exact_context_tokens(self, loop: au.Waveform) -> np.ndarray:
        """Build *exactly* context_length_frames worth of tokens (e.g., 250 @ 25fps),
        while ensuring the *end* of the audio lands on a bar boundary.
        Strategy: take the largest integer number of bars <= ctx_seconds as the tail,
        then left-fill from just before that tail (wrapping if needed) to reach exactly
        ctx_seconds; finally, pad/trim to exact samples and, as a last resort, pad/trim
        tokens to the expected frame count.
        """
        wav = loop.as_stereo().resample(self._model_sr)
        data = wav.samples.astype(np.float32, copy=False)
        if data.ndim == 1:
            data = data[:, None]

        spb = self._bar_clock.seconds_per_bar()
        ctx_sec = float(self._ctx_seconds)
        sr = int(self._model_sr)

        # bars that fit fully inside ctx_sec (at least 1)
        bars_fit = max(1, int(ctx_sec // spb))
        tail_len_samps = int(round(bars_fit * spb * sr))

        # ensure we have enough source by tiling
        need = int(round(ctx_sec * sr)) + tail_len_samps
        if data.shape[0] == 0:
            data = np.zeros((1, 2), dtype=np.float32)
        reps = int(np.ceil(need / float(data.shape[0])))
        tiled = np.tile(data, (reps, 1))

        end = tiled.shape[0]
        tail = tiled[end - tail_len_samps:end]

        # left-fill to reach exact ctx samples (keeps end-of-bar alignment)
        ctx_samps = int(round(ctx_sec * sr))
        pad_len = ctx_samps - tail.shape[0]
        if pad_len > 0:
            pre = tiled[end - tail_len_samps - pad_len:end - tail_len_samps]
            ctx = np.concatenate([pre, tail], axis=0)
        else:
            ctx = tail[-ctx_samps:]

        # final snap to *exact* ctx samples
        if ctx.shape[0] < ctx_samps:
            pad = np.zeros((ctx_samps - ctx.shape[0], ctx.shape[1]), dtype=np.float32)
            ctx = np.concatenate([pad, ctx], axis=0)
        elif ctx.shape[0] > ctx_samps:
            ctx = ctx[-ctx_samps:]

        exact = au.Waveform(ctx, sr)
        tokens_full = self.mrt.codec.encode(exact).astype(np.int32)
        depth = int(self.mrt.config.decoder_codec_rvq_depth)
        tokens = tokens_full[:, :depth]

        # Force expected (F,D) at *return time*
        tokens = self._coerce_tokens(tokens)
        return tokens

    def _encode_exact_context_tokens(self, loop: au.Waveform) -> np.ndarray:
        """Build *exactly* context_length_frames worth of tokens (e.g., 250 @ 25fps),
        while ensuring the *end* of the audio lands on a bar boundary.
        Strategy: take the largest integer number of bars <= ctx_seconds as the tail,
        then left-fill from just before that tail (wrapping if needed) to reach exactly
        ctx_seconds; finally, pad/trim to exact samples and, as a last resort, pad/trim
        tokens to the expected frame count.
        """
        wav = loop.as_stereo().resample(self._model_sr)
        data = wav.samples.astype(np.float32, copy=False)
        if data.ndim == 1:
            data = data[:, None]

        spb = self._bar_clock.seconds_per_bar()
        ctx_sec = float(self._ctx_seconds)
        sr = int(self._model_sr)

        # bars that fit fully inside ctx_sec (at least 1)
        bars_fit = max(1, int(ctx_sec // spb))
        tail_len_samps = int(round(bars_fit * spb * sr))

        # ensure we have enough source by tiling
        need = int(round(ctx_sec * sr)) + tail_len_samps
        if data.shape[0] == 0:
            data = np.zeros((1, 2), dtype=np.float32)
        reps = int(np.ceil(need / float(data.shape[0])))
        tiled = np.tile(data, (reps, 1))

        end = tiled.shape[0]
        tail = tiled[end - tail_len_samps:end]

        # left-fill to reach exact ctx samples (keeps end-of-bar alignment)
        ctx_samps = int(round(ctx_sec * sr))
        pad_len = ctx_samps - tail.shape[0]
        if pad_len > 0:
            pre = tiled[end - tail_len_samps - pad_len:end - tail_len_samps]
            ctx = np.concatenate([pre, tail], axis=0)
        else:
            ctx = tail[-ctx_samps:]

        # final snap to *exact* ctx samples
        if ctx.shape[0] < ctx_samps:
            pad = np.zeros((ctx_samps - ctx.shape[0], ctx.shape[1]), dtype=np.float32)
            ctx = np.concatenate([pad, ctx], axis=0)
        elif ctx.shape[0] > ctx_samps:
            ctx = ctx[-ctx_samps:]

        exact = au.Waveform(ctx, sr)
        tokens_full = self.mrt.codec.encode(exact).astype(np.int32)
        depth = int(self.mrt.config.decoder_codec_rvq_depth)
        tokens = tokens_full[:, :depth]

        # Last defense: force expected frame count
        frames = tokens.shape[0]
        exp = int(self._ctx_frames)
        if frames < exp:
            # repeat last frame
            pad = np.repeat(tokens[-1:, :], exp - frames, axis=0)
            tokens = np.concatenate([pad, tokens], axis=0)
        elif frames > exp:
            tokens = tokens[-exp:, :]
        return tokens


    def _install_context_from_loop(self, loop: au.Waveform):
        # Build exact-length, bar-locked context tokens
        context_tokens = self._encode_exact_context_tokens(loop)
        s = self.mrt.init_state()
        s.context_tokens = context_tokens
        self.state = s
        self._original_context_tokens = np.copy(context_tokens)

    def reseed_from_waveform(self, wav: au.Waveform):
        """Immediate reseed: replace context from provided wave (bar-locked, exact length)."""
        context_tokens = self._encode_exact_context_tokens(wav)
        with self._lock:
            s = self.mrt.init_state()
            s.context_tokens = context_tokens
            self.state = s
            self._model_stream = None  # drop model-domain continuity so next chunk starts cleanly
            self._original_context_tokens = np.copy(context_tokens)

    def reseed_splice(self, recent_wav: au.Waveform, anchor_bars: float):
        """Queue a *seamless* reseed by token splicing instead of full restart.
        We compute a fresh, bar-locked context token tensor of exact length
        (e.g., 250 frames), then splice only the *tail* corresponding to
        `anchor_bars` so generation continues smoothly without resetting state.
        """
        new_ctx = self._encode_exact_context_tokens(recent_wav)  # coerce to (F,D)
        F, D = self._expected_token_shape()

        # how many frames correspond to the requested anchor bars
        spb = self._bar_clock.seconds_per_bar()
        frames_per_bar = max(1, int(round(self._codec_fps * spb)))
        splice_frames = max(1, min(int(round(max(1.0, float(anchor_bars)) * frames_per_bar)), F))

        with self._lock:
            # snapshot current context
            cur = getattr(self.state, "context_tokens", None)
            if cur is None:
                # fall back to full reseed (still coerced)
                self._pending_reseed = {"ctx": new_ctx}
                return
            cur = self._coerce_tokens(cur)

            # build the spliced tensor: keep left (F - splice) from cur, take right (splice) from new
            left = cur[:F - splice_frames, :]
            right = new_ctx[F - splice_frames:, :]
            spliced = np.concatenate([left, right], axis=0)
            spliced = self._coerce_tokens(spliced)

            # queue for install at the *next bar boundary* right after emission
            self._pending_token_splice = {
                "tokens": spliced,
                "debug": {"F": F, "D": D, "splice_frames": splice_frames, "frames_per_bar": frames_per_bar}
            }
            


    def reseed_from_waveform(self, wav: au.Waveform):
        """Immediate reseed: replace context from provided wave (bar-aligned tail)."""
        wav = wav.as_stereo().resample(self._model_sr)
        tail = take_bar_aligned_tail(wav, self.params.bpm, self.params.beats_per_bar, self._ctx_seconds)
        tokens_full = self.mrt.codec.encode(tail).astype(np.int32)
        depth = int(self.mrt.config.decoder_codec_rvq_depth)
        context_tokens = tokens_full[:, :depth]

        s = self.mrt.init_state()
        s.context_tokens = context_tokens
        self.state = s
        # reset model stream so next generate starts cleanly
        self._model_stream = None

        # optional loudness match will be applied per-chunk on emission

        # also remember this as new "original"
        self._original_context_tokens = np.copy(context_tokens)

    # ---------- core streaming helpers ----------

    def _append_model_chunk_and_spool(self, wav: au.Waveform) -> None:
        """
        Append one MagentaRT chunk into the target-SR spool with an energy-aware,
        deferred-overwrite crossfade to avoid writing near-silence at bar edges.

        Key behavior:
        - Append BODY and TAIL of *this* chunk right away (resampled to target SR).
        - Keep THIS chunk's model-rate TAIL (+ its target-SR length if appended) to repair the
        previous boundary on the *next* call by mixing (prev_tail*cos + new_head*sin).
        - When the correction length Lpop would be 0 (e.g., tail produced no target samples last time),
        we APPEND the mixed-overlap to bridge the gap instead of overwriting 0 samples.
        - Before overwriting/appending the mixed-overlap, we guard against writing ultra-quiet audio
        by normalizing it up (bounded) if it's >20 dB below the existing spool end.

        This keeps your bar clock and external timing the same, but removes "bad starts" and fizzles.
        """
        import math
        import numpy as np

        # ---- helpers ----
        def _rms_dbfs(x: np.ndarray) -> float:
            if x.size == 0:
                return -120.0
            if x.ndim == 2 and x.shape[1] > 1:
                x_m = x.mean(axis=1, dtype=np.float32)
            else:
                x_m = x.astype(np.float32, copy=False).reshape(-1)
            # guard for NaNs
            x_m = np.nan_to_num(x_m, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32, copy=False)
            r = float(np.sqrt(np.mean(x_m * x_m) + 1e-12))
            return 20.0 * math.log10(max(r, 1e-12))

        def _rms_dbfs_model(x: np.ndarray) -> float:
            # same metric; named for clarity in logs
            return _rms_dbfs(x)

        def to_target(y: np.ndarray) -> np.ndarray:
            return y if self._rs is None else self._rs.process(y, final=False)

        # ---- unpack model-rate samples ----
        s = wav.samples.astype(np.float32, copy=False)
        if s.ndim == 1:
            s = s[:, None]
        if s.shape[1] == 1:
            # ensure stereo shape for consistency with your spool (S,2)
            s = np.repeat(s, 2, axis=1)

        n_samps = int(s.shape[0])

        # crossfade length in model samples
        try:
            xfade_s = float(self.mrt.config.crossfade_length)
        except Exception:
            xfade_s = 0.0
        xfade_n = int(round(max(0.0, xfade_s) * float(self._model_sr)))

        # carve head/body/tail in model domain
        if xfade_n > 0 and n_samps >= (2 * xfade_n):
            head_m = s[:xfade_n, :]
            body_m = s[xfade_n:n_samps - xfade_n, :]
            tail_m = s[n_samps - xfade_n:, :]
        else:
            # too short or no xfade configured — treat everything as body
            head_m = np.zeros((0, 2), dtype=np.float32)
            body_m = s
            tail_m = np.zeros((0, 2), dtype=np.float32)

        # ------------------------------------------
        # (A) Repair the PREVIOUS boundary if we have a pending model-tail
        # ------------------------------------------
        did_boundary_mix = False
        if (self._pending_tail_model is not None) and (xfade_n > 0) and (n_samps >= xfade_n):
            # adaptive crossfade length when either side is very quiet
            tail_prev_m = self._pending_tail_model
            head_now_m  = head_m

            # safety: match shapes
            if tail_prev_m.shape[1] != 2:
                if tail_prev_m.ndim == 1:
                    tail_prev_m = tail_prev_m[:, None]
                tail_prev_m = np.repeat(tail_prev_m[:, :1], 2, axis=1)
            if head_now_m.shape[1] != 2:
                if head_now_m.ndim == 1:
                    head_now_m = head_now_m[:, None]
                head_now_m = np.repeat(head_now_m[:, :1], 2, axis=1)

            # compute energy to decide whether to shorten xfade
            tail_r = _rms_dbfs_model(tail_prev_m)
            head_r = _rms_dbfs_model(head_now_m)
            xfade_use = int(xfade_n)
            if min(tail_r, head_r) < -45.0:
                xfade_use = max(1, xfade_n // 4)

            # windowed overlap (model domain)
            Lm = min(xfade_use, tail_prev_m.shape[0], head_now_m.shape[0])
            if Lm > 0:
                t = np.linspace(0.0, math.pi / 2.0, Lm, endpoint=False, dtype=np.float32)[:, None]
                cosw = np.cos(t, dtype=np.float32)
                sinw = np.sin(t, dtype=np.float32)
                mixed_m = tail_prev_m[-Lm:, :] * cosw + head_now_m[:Lm, :] * sinw

                # resample to target and correct the end of the spool
                y_mixed = to_target(mixed_m)
                Lcorr = int(y_mixed.shape[0])

                if Lcorr > 0:
                    # how many samples from last time's tail did we append?
                    # (may be zero if resampler yielded nothing then)
                    Lpop = int(min(self._pending_tail_target_len, self._spool.shape[0], Lcorr))

                    if Lpop > 0:
                        # energy-aware overwrite of last Lpop samples
                        prev_end = self._spool[-Lpop:, :]
                        new_seg  = y_mixed[-Lpop:, :]

                        prev_r = _rms_dbfs(prev_end)
                        new_r  = _rms_dbfs(new_seg)

                        # If the new overlap is >20 dB quieter than what's there, lift it (bounded)
                        if new_r < (prev_r - 20.0):
                            lift_db = max(0.0, min(20.0, (prev_r - 6.0) - new_r))  # cap boost; leave ~6 dB headroom
                            scale   = 10.0 ** (lift_db / 20.0)
                            new_seg = np.clip(new_seg * scale, -1.0, 1.0).astype(np.float32, copy=False)

                        self._spool[-Lpop:, :] = new_seg
                        print(f"[append] mixedOverlap len={Lpop} rms={_rms_dbfs(new_seg):+.1f} dBFS")
                    else:
                        # Nothing to overwrite (e.g., last tail produced 0 target samples).
                        # Bridge by APPENDING the mixed-overlap.
                        self._spool = np.concatenate([self._spool, y_mixed], axis=0)
                        self._spool_written += int(y_mixed.shape[0])
                        print(f"[append] mixedOverlap len={y_mixed.shape[0]} rms={_rms_dbfs(y_mixed):+.1f} dBFS")

                    did_boundary_mix = True

            # clear pending once we attempted the repair
            self._pending_tail_model = None
            self._pending_tail_target_len = 0

        # ------------------------------------------
        # (B) Append this chunk's BODY then TAIL (target SR)
        # ------------------------------------------
        # BODY
        y_body = to_target(body_m) if body_m.size else np.zeros((0, 2), dtype=np.float32)
        if y_body.size:
            self._spool = np.concatenate([self._spool, y_body], axis=0)
            self._spool_written += int(y_body.shape[0])
        print(f"[append] body len={y_body.shape[0] if y_body.size else 0} rms={_rms_dbfs(y_body):+.1f} dBFS")

        # TAIL (we append now to keep continuity; on next call we'll correct the end)
        y_tail = to_target(tail_m) if tail_m.size else np.zeros((0, 2), dtype=np.float32)
        if y_tail.size:
            self._spool = np.concatenate([self._spool, y_tail], axis=0)
            self._spool_written += int(y_tail.shape[0])
            self._pending_tail_target_len = int(y_tail.shape[0])   # how much we just added at target SR
        else:
            # resampler returned nothing for the tail; mark 0 so next Lpop==0
            self._pending_tail_target_len = 0
        print(f"[append] tail len={y_tail.shape[0] if y_tail.size else 0} rms={_rms_dbfs(y_tail):+.1f} dBFS")

        # keep THIS chunk's model tail to mix with next chunk's head
        # (even if y_tail had 0 target samples; in that case we'll bridge by appending mixed overlap)
        self._pending_tail_model = tail_m if tail_m.size else None




    def _should_generate_next_chunk(self) -> bool:
        # Allow running ahead relative to whichever is larger: last *consumed*
        # (explicit ack from client) or last *delivered* (implicit ack).
        implicit_consumed = self._next_to_deliver - 1  # last chunk handed to client
        horizon_anchor = max(self._last_consumed_index, implicit_consumed)
        return self.idx <= (horizon_anchor + self._max_buffer_ahead)

    def _emit_ready(self):
        """Emit next chunk(s) if the spool has enough samples."""
        while True:
            start, end = self._bar_clock.bounds_for_chunk(self.idx, self.params.bars_per_chunk)
            if end > self._spool_written:
                break  # need more audio
            loop = self._spool[start:end]

            # Loudness match per chunk (bar-aligned reference)
            if self.params.loudness_mode != "none" and self.params.combined_loop is not None:
                sr = int(self.params.target_sr)

                # 1) Get the combined loop at target SR (stereo, float32)
                comb = self.params.combined_loop.as_stereo().resample(sr).samples.astype(np.float32, copy=False)
                if comb.ndim == 1:
                    comb = comb[:, None]
                if comb.shape[1] == 1:
                    comb = np.repeat(comb, 2, axis=1)

                # 2) Build a reference slice aligned to this outgoing chunk [start:end]
                #    We wrap/tile the combined loop so it always covers the needed range.
                need = end - start
                if comb.shape[0] > 0 and need > 0:
                    s = start % comb.shape[0]
                    if s + need <= comb.shape[0]:
                        ref_slice = comb[s:s+need]
                    else:
                        part1 = comb[s:]
                        part2 = comb[:max(0, need - part1.shape[0])]
                        ref_slice = np.vstack([part1, part2])

                    ref = au.Waveform(ref_slice, sr)
                    tgt = au.Waveform(loop.copy(), sr)

                    matched, _stats = match_loudness_to_reference(
                        ref, tgt,
                        method=self.params.loudness_mode,
                        headroom_db=self.params.headroom_db
                    )
                    loop = matched.samples

            audio_b64, total_samples, channels = wav_bytes_base64(loop, int(self.params.target_sr))
            meta = {
                "bpm": float(self.params.bpm),
                "bars": int(self.params.bars_per_chunk),
                "beats_per_bar": int(self.params.beats_per_bar),
                "sample_rate": int(self.params.target_sr),
                "channels": int(channels),
                "total_samples": int(total_samples),
                "seconds_per_bar": self._bar_clock.seconds_per_bar(),
                "loop_duration_seconds": self.params.bars_per_chunk * self._bar_clock.seconds_per_bar(),
                "guidance_weight": float(self.params.guidance_weight),
                "temperature": float(self.params.temperature),
                "topk": int(self.params.topk),
            }
            chunk = JamChunk(index=self.idx, audio_base64=audio_b64, metadata=meta)

            if os.getenv("MRT_DEBUG_RMS", "0") == "1":
                spb = self._bar_clock.bar_samps
                seg = int(max(1, spb // 4))  # quarter-bar window
                
                rms = [float(np.sqrt(np.mean(loop[i:i+seg]**2))) for i in range(0, loop.shape[0], seg)]
                print(f"[emit idx={self.idx}] quarter-bar RMS: {rms[:8]}")

            with self._cv:
                self._outbox[self.idx] = chunk
                self._cv.notify_all()
            self.idx += 1

            # If a reseed is queued, install it *right after* we finish a chunk
            with self._lock:
                # Prefer seamless token splice when available
                if self._pending_token_splice is not None:
                    spliced = self._coerce_tokens(self._pending_token_splice["tokens"])
                    try:
                        # inplace update (no reset)
                        self.state.context_tokens = spliced
                        self._pending_token_splice = None
                    except Exception:
                        # fallback: full reseed using spliced tokens
                        new_state = self.mrt.init_state()
                        new_state.context_tokens = spliced
                        self.state = new_state
                        self._model_stream = None
                        self._pending_token_splice = None
                elif self._pending_reseed is not None:
                    ctx = self._coerce_tokens(self._pending_reseed["ctx"])
                    new_state = self.mrt.init_state()
                    new_state.context_tokens = ctx
                    self.state = new_state
                    self._model_stream = None
                    self._pending_reseed = None

    # ---------- main loop ----------

    def run(self):
        # generate until stopped
        while not self._stop_event.is_set():
            # throttle generation if we are far ahead
            if not self._should_generate_next_chunk():
                # still try to emit if spool already has enough
                self._emit_ready()
                time.sleep(0.01)
                continue

            # generate next model chunk
            # snapshot current style vector under lock for this step
            with self._lock:
                target = self.params.style_vec
                if target is None:
                    style_to_use = None
                else:
                    if self._style_vec is None:  # first use: start exactly at initial style (no glide)
                        self._style_vec = np.array(target, dtype=np.float32, copy=True)
                    else:
                        ramp = float(self.params.style_ramp_seconds or 0.0)
                        step = 1.0 if ramp <= 0.0 else min(1.0, self._chunk_secs / ramp)
                        # linear ramp in embedding space
                        self._style_vec += step * (target.astype(np.float32, copy=False) - self._style_vec)
                    style_to_use = self._style_vec

            wav, self.state = self.mrt.generate_chunk(state=self.state, style=style_to_use)
            # append and spool
            self._append_model_chunk_and_spool(wav)
            # try emitting zero or more chunks if available
            self._emit_ready()

        # finalize resampler (flush) — not strictly necessary here
        tail = self._rs.process(np.zeros((0,2), np.float32), final=True)
        if tail.size:
            self._spool = np.concatenate([self._spool, tail], axis=0)
            self._spool_written += tail.shape[0]
        # one last emit attempt
        self._emit_ready()