"""Light signal building blocks — onsets + waveform + spectrum. CPU, NO torch.

The cheapest analysis tier: thin librosa wrappers that need no model at all.
Especially useful for hexa (a visuals platform): /waveform gives a render-ready
peak/energy envelope, /onsets gives rhythmic hit times for visual sync + slicing.

FULL-TRACK (#7): these stream the WHOLE upload in bounded windows and aggregate,
instead of only the first 60 s — so a long track is analyzed end-to-end without
buffering it all in RAM. waveform/spectrum aggregate EXACTLY (fixed output grid
filled by streaming reduction); onsets builds one full-track onset envelope (tiny
— hop-downsampled) from per-window pieces. Native sample rate throughout (no
per-block resample, which would seam at window edges). Self-contained like
ears_light/feats (the container has only armada/api/).
"""
import numpy as np

from . import streaming

HOP = 512                 # librosa default; fixes frame→time mapping across windows
_MIN_FRAMES_SR = 20       # need ≥ 1/20 s of audio to say anything
_FALLBACK_CAP_S = 600.0   # bounded whole-file decode when a format isn't frame-countable


def _load_capped(audio_path, sr_target=22050):
    """Bounded whole-file fallback for formats soundfile can't frame-count.
    Decodes at most _FALLBACK_CAP_S seconds — never the unbounded file."""
    import librosa
    y, sr = librosa.load(str(audio_path), sr=sr_target, mono=True, duration=_FALLBACK_CAP_S)
    if y.size < sr // _MIN_FRAMES_SR:
        raise ValueError("clip too short / undecodable (need ≥ 50 ms of audio)")
    return y, sr


def onsets(audio_path, win_s=streaming.DEFAULT_WIN_S) -> dict:
    """Onset hit times (s) + tempo + onset rate over the WHOLE track. Builds one
    full-track onset-strength envelope by concatenating per-window pieces (the
    envelope is hop-downsampled, so a 10-min track is only ~tens of KB), then runs
    detection + drift-tolerant tempo on it once."""
    import librosa
    meta = streaming.probe(audio_path)
    sr = meta["sr"]
    if meta["frames"] <= 0:                       # uncountable format → bounded fallback
        y, sr = _load_capped(audio_path)
        env = librosa.onset.onset_strength(y=y, sr=sr, hop_length=HOP)
        full = False
    else:
        pieces, full = [], True
        for block, sr, _ in streaming.windows(audio_path, win_s, mono=True):
            if block.size:
                pieces.append(librosa.onset.onset_strength(y=block, sr=sr, hop_length=HOP))
        if not pieces:
            raise ValueError("clip too short / undecodable (need ≥ 50 ms of audio)")
        env = np.concatenate(pieces)

    times = librosa.onset.onset_detect(onset_envelope=env, sr=sr, hop_length=HOP, units="time")
    tempo = float(np.atleast_1d(librosa.feature.tempo(onset_envelope=env, sr=sr, hop_length=HOP))[0])
    dur = meta["duration_s"] or (len(env) * HOP / sr)
    return {
        "tempo_bpm": round(tempo, 2),
        "onsets": [round(float(t), 4) for t in times],
        "n_onsets": int(len(times)),
        "onset_rate": round(len(times) / dur, 3) if dur else 0.0,
        "duration_s": round(dur, 3),
        "sr": int(sr),
        "engine": "light",
        "full_track": full,
    }


def waveform(audio_path, bins=800, win_s=streaming.DEFAULT_WIN_S) -> dict:
    """Downsampled waveform for rendering: per-bin [min, max] in [-1,1] + a 0..1
    RMS energy envelope, spanning the WHOLE track. Exact: each `bins` bucket is
    filled by an unsegmented min/max/mean-square reduction over the streamed
    samples, so the result equals a single-pass computation."""
    bins = max(1, min(int(bins), 4000))
    meta = streaming.probe(audio_path)
    total, sr = meta["frames"], meta["sr"]

    if total <= 0:                                 # uncountable format → bounded fallback
        y, sr = _load_capped(audio_path)
        return _waveform_array(y, sr, bins, full=False)

    mn = np.full(bins, np.inf, dtype="float64")
    mx = np.full(bins, -np.inf, dtype="float64")
    sq = np.zeros(bins, dtype="float64")
    cnt = np.zeros(bins, dtype="int64")
    for block, sr, start in streaming.windows(audio_path, win_s, mono=True):
        if not block.size:
            continue
        idx = ((np.arange(len(block), dtype="int64") + start) * bins // total)
        np.clip(idx, 0, bins - 1, out=idx)
        np.minimum.at(mn, idx, block)
        np.maximum.at(mx, idx, block)
        np.add.at(sq, idx, block.astype("float64") ** 2)
        np.add.at(cnt, idx, 1)
    if cnt.sum() < sr // _MIN_FRAMES_SR:
        raise ValueError("clip too short / undecodable (need ≥ 50 ms of audio)")
    return _finish_waveform(mn, mx, sq, cnt, bins, total, sr, full=True)


def _waveform_array(y, sr, bins, full):
    edges = np.linspace(0, len(y), bins + 1).astype(int)
    mn = np.zeros(bins); mx = np.zeros(bins); sq = np.zeros(bins); cnt = np.zeros(bins, dtype="int64")
    for i in range(bins):
        seg = y[edges[i]:edges[i + 1]]
        if seg.size:
            mn[i], mx[i] = seg.min(), seg.max()
            sq[i] = np.sum(seg.astype("float64") ** 2); cnt[i] = seg.size
    return _finish_waveform(mn, mx, sq, cnt, bins, len(y), sr, full=full)


def _finish_waveform(mn, mx, sq, cnt, bins, total, sr, full):
    empty = cnt == 0
    mn = np.where(empty, 0.0, mn)
    mx = np.where(empty, 0.0, mx)
    rms = np.sqrt(np.where(cnt > 0, sq / np.maximum(cnt, 1), 0.0))
    norm = float(rms.max()) or 1.0
    return {
        "bins": bins,
        "peaks": [[round(float(lo), 4), round(float(hi), 4)] for lo, hi in zip(mn, mx)],
        "rms": [round(float(r / norm), 4) for r in rms],
        "duration_s": round(total / sr, 3) if sr else 0.0,
        "sr": int(sr),
        "engine": "light",
        "full_track": full,
    }


def spectrum(audio_path, bands=64, frames=200, mel=True, win_s=streaming.DEFAULT_WIN_S) -> dict:
    """FFT-as-a-service: a downsampled, render-ready spectrogram — `bands` ×
    `frames`, 0..1 normalized, spanning the WHOLE track. Streams windows, bins each
    window's spectrogram columns into the global `frames` grid by absolute time,
    accumulates mean POWER per (band, frame), then converts to dB (ref = the
    track's global peak power) once at the end. `frames=1` → a single time-averaged
    spectrum. Torch-free."""
    import librosa
    bands = max(1, min(int(bands), 256))
    frames = max(1, min(int(frames), 1000))
    meta = streaming.probe(audio_path)
    total_frames, sr = meta["frames"], meta["sr"]

    if total_frames <= 0:                          # uncountable format → bounded fallback
        y, sr = _load_capped(audio_path)
        total_frames = len(y)
        blocks = [(y, sr, 0)]
        full = False
    else:
        blocks = streaming.windows(audio_path, win_s, mono=True)
        full = True

    acc = np.zeros((bands, frames), dtype="float64")   # summed power per (band, global-frame)
    cnt = np.zeros(frames, dtype="int64")
    gmax = 1e-12                                        # global peak power for the dB ref
    band_hz = None
    total_dur = (total_frames / sr) if sr else 0.0

    for block, sr, start in blocks:
        if block.size < HOP:
            continue
        if mel:
            S = librosa.feature.melspectrogram(y=block, sr=sr, n_mels=bands, hop_length=HOP)
            if band_hz is None:
                band_hz = [round(float(f), 1) for f in librosa.mel_frequencies(n_mels=bands, fmax=sr / 2)]
        else:
            mag = np.abs(librosa.stft(block, n_fft=2048, hop_length=HOP)) ** 2
            freqs = librosa.fft_frequencies(sr=sr, n_fft=2048)
            edges = np.logspace(np.log10(20), np.log10(sr / 2), bands + 1)
            S = np.stack([mag[(freqs >= edges[i]) & (freqs < edges[i + 1])].sum(axis=0)
                          for i in range(bands)])
            if band_hz is None:
                band_hz = [round(float(np.sqrt(edges[i] * edges[i + 1])), 1) for i in range(bands)]
        gmax = max(gmax, float(S.max()))
        # map each column to its absolute time → global frame bucket
        col_t = (start + np.arange(S.shape[1]) * HOP) / sr
        gf = np.zeros(S.shape[1], dtype="int64") if total_dur <= 0 else \
            np.clip((col_t / total_dur * frames).astype("int64"), 0, frames - 1)
        for b in range(frames):
            sel = gf == b
            if sel.any():
                acc[:, b] += S[:, sel].sum(axis=1)
                cnt[b] += int(sel.sum())

    if cnt.sum() == 0:
        raise ValueError("clip too short / undecodable (need ≥ 50 ms of audio)")
    # empty frame buckets (sparse coverage) borrow the nearest filled one's mean
    mean_p = np.zeros((bands, frames), dtype="float64")
    filled = cnt > 0
    mean_p[:, filled] = acc[:, filled] / cnt[filled]
    if not filled.all():
        idx = np.where(filled)[0]
        for b in np.where(~filled)[0]:
            mean_p[:, b] = mean_p[:, idx[np.argmin(np.abs(idx - b))]]
    Sdb = librosa.power_to_db(mean_p + 1e-12, ref=gmax)
    spec = (np.clip(Sdb, -80, 0) + 80) / 80.0
    return {
        "bands": bands,
        "frames": frames,
        "mel": bool(mel),
        "band_hz": band_hz or [],
        "spec": [[round(float(v), 4) for v in row] for row in spec],
        "duration_s": round(total_dur, 3),
        "sr": int(sr),
        "engine": "light",
        "full_track": full,
    }
