"""stitch.py — splice the SA3 continuation onto the user's pristine original.

CODA's identity: your real recording plays untouched up to the seam, then the
generated tail takes over. Everything here is 44.1 kHz stereo — SA3's native
format and the deliverable's — so the original is resampled up to 44.1k and made
stereo, and the engine's tail is already there.

Unlike the old MusicGen path, the SA3 tail does NOT contain a re-encoded copy of
the source; it is fresh audio that begins exactly where the source ends. The seam
is therefore a join between sequential content, so:
  - level: the tail is loudness-matched to the original's tail RMS so the seam
    doesn't pump (gain bounded so a quiet lo-fi clip can't crush the tail);
  - click: a short equal-power crossfade smooths the join;
  - close: a cos^2 fade to true silence ends the track like a finished song;
  - then one peak-normalize lifts the whole (now level-consistent) track to a
    confident listening level with -1 dBFS of headroom.
"""
import librosa
import numpy as np

SR = 44100


def to_stereo_44k(audio, sr):
    """(channels, N) or (N,) @sr  ->  (2, N') float32 @44.1k stereo."""
    a = np.asarray(audio, dtype=np.float32)
    if a.ndim == 1:
        a = a[None, :]
    if sr != SR:
        a = np.stack([
            librosa.resample(ch, orig_sr=sr, target_sr=SR, res_type="soxr_hq")
            for ch in a
        ])
    if a.shape[0] == 1:
        a = np.repeat(a, 2, axis=0)
    elif a.shape[0] > 2:
        a = a[:2]
    return np.ascontiguousarray(a.astype(np.float32))


def _rms(x):
    return float(np.sqrt(np.mean(np.asarray(x, dtype=np.float64) ** 2)) + 1e-12)


def stitch(original, original_sr, new_tail, source_seconds,
           crossfade_seconds=0.10, end_fade_seconds=4.0, match_seconds=2.0,
           peak_ceiling=0.891):
    """Join the user's original to the SA3-generated tail.

    original       : (channels, N) or (N,) float32 @original_sr — pristine clip
    new_tail       : (2, M) float32 @44.1k from engine.continue_audio
    source_seconds : the splice boundary (clip length) in seconds
    returns        : (2, T) float32 @44.1k, peak == peak_ceiling (-1 dBFS)
    """
    orig = to_stereo_44k(original, original_sr)
    tail = to_stereo_44k(new_tail, SR)

    boundary = min(int(round(source_seconds * SR)), orig.shape[-1])
    orig = orig[:, :boundary]          # only the real recording up to the seam

    # loudness-match the tail to the original's level at the seam (continuity).
    # bound the gain: a very quiet lo-fi clip shouldn't drag the full-bodied
    # tail down to a whisper, and we never amplify the tail wildly either.
    m = min(int(match_seconds * SR), orig.shape[-1], tail.shape[-1])
    if m > 0:
        gain = float(np.clip(_rms(orig[:, -m:]) / _rms(tail[:, :m]), 0.4, 2.5))
        tail = tail * gain

    # short equal-power crossfade across the join. the two sides are sequential
    # (not time-aligned copies), so equal-power — not equal-gain — keeps the
    # energy flat through the blend.
    fade = min(int(crossfade_seconds * SR), orig.shape[-1], tail.shape[-1])
    if fade > 0:
        t = np.linspace(0.0, 1.0, fade, dtype=np.float32)
        fout, fin = np.cos(t * np.pi / 2), np.sin(t * np.pi / 2)
        seam = orig[:, -fade:] * fout + tail[:, :fade] * fin
        out = np.concatenate([orig[:, :-fade], seam, tail[:, fade:]], axis=-1)
    else:
        out = np.concatenate([orig, tail], axis=-1)
    out = out.astype(np.float32)

    # closing fade to true silence so the song ends instead of cutting off
    end_fade = min(int(end_fade_seconds * SR), out.shape[-1])
    if end_fade > 0:
        curve = np.cos(np.linspace(0.0, np.pi / 2, end_fade,
                                   dtype=np.float32)) ** 2
        out[:, -end_fade:] *= curve

    # lift the whole, now level-consistent, track to a confident level
    peak = float(np.abs(out).max())
    if peak > 1e-9:
        out = out * (peak_ceiling / peak)
    return out.astype(np.float32)