#!/usr/bin/env python3
"""把生成的音乐裁成指定长度的无缝循环 BGM，并做客观验收。

MiniMax 接口既没有 duration 参数、也没有循环淡化，游戏/短片要定长循环 BGM
只能后期做。本脚本负责：
  1. 测素材实际速度，把循环长度对齐到整数小节（只对齐电平不对齐节奏，
     循环起来会丢拍）
  2. 扫描起点，选首尾 2 秒在 RMS/频谱质心/低频占比上最接近的窗口，
     并避开渐入、渐出和能量凹陷
  3. 用 qsin 等功率曲线做尾→头交叉淡化
  4. 验收：峰值、静音、接缝跳变、立体声宽度

依赖：ffmpeg/ffprobe + numpy
用法：
    python3 loopify.py raw.mp3 -o bgm_loop.mp3 -L 45 --preview 3
"""
import argparse, os, shutil, subprocess, sys

try:
    import numpy as np
except ImportError:
    sys.exit("需要 numpy：pip3 install numpy")

if not shutil.which("ffmpeg"):
    sys.exit("需要 ffmpeg：brew install ffmpeg")

ANALYZE_SR = 22050
HOP = 512


def decode(path, sr, ch=1):
    r = subprocess.run(["ffmpeg", "-v", "error", "-i", path, "-ac", str(ch),
                        "-ar", str(sr), "-f", "f32le", "-"],
                       capture_output=True)
    if r.returncode != 0:
        sys.exit("解码失败：%s" % r.stderr.decode("utf-8", "replace")[:400])
    a = np.frombuffer(r.stdout, dtype=np.float32)
    return a.reshape(-1, ch) if ch > 1 else a


def beat_period(x):
    """谱通量 + 自相关，估计节拍周期（秒）。"""
    win = 1024
    n = (len(x) - win) // HOP
    if n < 64:
        return None
    idx = np.arange(n)[:, None] * HOP + np.arange(win)
    S = np.abs(np.fft.rfft(x[idx] * np.hanning(win), axis=1))
    flux = np.maximum(0, np.diff(S, axis=0)).sum(axis=1)
    flux = flux - flux.mean()
    fps = ANALYZE_SR / HOP
    ac = np.correlate(flux, flux, "full")[len(flux) - 1:]
    lo, hi = int(fps * 60 / 160), int(fps * 60 / 60)
    if hi >= len(ac):
        return None
    return (lo + int(np.argmax(ac[lo:hi]))) / fps


def pick_window(x, dur, target, tol, fade):
    """返回 (t0, L, score, 说明)。"""
    beat = beat_period(x)
    cands = []
    if beat:
        for bpb in (3, 4, 6, 8):
            bar = beat * bpb
            k = 1
            while bar * k <= target + tol:
                L = bar * k
                if target - tol <= L <= target + tol:
                    cands.append((L, "%d 小节 × %d 拍 @ %.1f BPM"
                                  % (k, bpb, 60 / beat)))
                k += 1
    if not cands:
        cands = [(float(target), "未测出稳定节拍，用目标长度")]

    def feat(t):
        a = x[int(t * ANALYZE_SR):int((t + fade) * ANALYZE_SR)]
        if len(a) < ANALYZE_SR // 2:
            return None
        sp = np.abs(np.fft.rfft(a * np.hanning(len(a))))
        fr = np.fft.rfftfreq(len(a), 1 / ANALYZE_SR)
        e = sp.sum() + 1e-9
        return np.array([20 * np.log10(np.sqrt((a ** 2).mean()) + 1e-9),
                         (sp * fr).sum() / e / 1000.0,
                         sp[fr < 300].sum() / e * 20])

    best = None
    for L, why in cands:
        if L + fade + 1.5 > dur:
            continue
        t0 = 1.0
        while t0 + L + fade <= dur - 0.5:
            h, t = feat(t0), feat(t0 + L)
            if h is not None and t is not None:
                seg = x[int(t0 * ANALYZE_SR):int((t0 + L) * ANALYZE_SR)]
                k = int(ANALYZE_SR * 0.5)
                quietest = min(np.sqrt((seg[i:i + k] ** 2).mean())
                               for i in range(0, max(1, len(seg) - k), k))
                penalty = max(0.0, -20 * np.log10(quietest + 1e-9) - 40) * 0.5
                d = float(np.abs(h - t).sum()) + penalty
                if best is None or d < best[2]:
                    best = (t0, L, d, why)
            t0 += 0.05
    if best is None:
        sys.exit("素材太短，做不出 %.1fs 的循环（需要至少 %.1fs）"
                 % (target, target + fade + 2.5))
    return best


def build(src, out, t0, L, fade, bitrate, peak_dbfs):
    wav = os.path.splitext(out)[0] + ".wav"
    fc = ("[0:a]atrim=start=%.4f:duration=%.4f,asetpts=PTS-STARTPTS[tail];"
          "[1:a]atrim=start=%.4f:duration=%.4f,asetpts=PTS-STARTPTS[body];"
          "[tail][body]acrossfade=d=%.2f:c1=qsin:c2=qsin[out]"
          % (t0 + L, fade, t0, L, fade))
    subprocess.run(["ffmpeg", "-hide_banner", "-v", "error", "-y",
                    "-i", src, "-i", src, "-filter_complex", fc,
                    "-map", "[out]", "-c:a", "pcm_s24le", wav], check=True)

    # MiniMax 的输出电平不稳定（实测有 -2.3 dBFS 的，也有 0.0 dBFS 顶格的）。
    # 顶格素材经交叉淡化两路叠加必然溢出，所以这里统一压到目标峰值。
    target = 10 ** (peak_dbfs / 20.0)
    p = float(np.abs(decode(wav, 44100, 2)).max())
    if p > target:
        g = target / max(p, 1e-9)
        tmp = wav + ".tmp.wav"
        subprocess.run(["ffmpeg", "-hide_banner", "-v", "error", "-y", "-i", wav,
                        "-af", "volume=%.6f" % g, "-c:a", "pcm_s24le", tmp],
                       check=True)
        os.replace(tmp, wav)
        print("留余量: 峰值 %.1f → %.1f dBFS（衰减 %.1f dB）"
              % (20 * np.log10(p + 1e-9), peak_dbfs, 20 * np.log10(g)))

    subprocess.run(["ffmpeg", "-hide_banner", "-v", "error", "-y", "-i", wav,
                    "-c:a", "libmp3lame", "-b:a", bitrate, out], check=True)
    return wav


def verify(wav, sr=44100):
    x = decode(wav, sr, 2)
    ok = True
    print("\n=== 验收 ===")
    print("时长 %.3fs · %d 声道 · %d Hz" % (len(x) / sr, x.shape[1], sr))

    peak = float(np.abs(x).max())
    good = peak < 0.999
    ok &= good
    print("峰值 %.4f (%.1f dBFS) %s" % (peak, 20 * np.log10(peak + 1e-9),
                                        "OK" if good else "削波!"))

    def rms_db(a):
        return 20 * np.log10(np.sqrt((a ** 2).mean()) + 1e-9)

    k = sr // 2
    q = min(rms_db(x[i:i + k]) for i in range(0, len(x) - k, k // 2))
    good = q > -45
    ok &= good
    print("最静 0.5s %.1f dB %s" % (q, "OK" if good else "有静音段!"))

    mono = x.mean(axis=1)
    typ = float(np.percentile(np.abs(np.diff(mono)), 99.9))
    seam = float(abs(mono[0] - mono[-1]))
    good = seam <= typ
    ok &= good
    print("接缝跳变 %.6f vs 曲内 99.9 分位 %.6f（比值 %.2f）%s"
          % (seam, typ, seam / (typ + 1e-12), "OK" if good else "有咔哒声!"))

    d = abs(rms_db(mono[-k:]) - rms_db(mono[:k]))
    good = d < 3
    ok &= good
    print("接缝前后 RMS 差 %.1f dB %s" % (d, "OK" if good else "电平不匹配"))

    w = float(np.abs(x[:, 0] - x[:, 1]).mean() / (np.abs(x).mean() + 1e-9))
    print("立体声宽度 %.3f %s" % (w, "有空间感" if w > 0.1 else "接近单声道"))
    print("=== %s ===" % ("全项通过" if ok else "有项未通过，见上"))
    return ok


def main():
    ap = argparse.ArgumentParser(description="裁成无缝循环 BGM")
    ap.add_argument("input")
    ap.add_argument("-o", "--output", default="bgm_loop.mp3")
    ap.add_argument("-L", "--length", type=float, default=45.0, help="目标秒数")
    ap.add_argument("--tol", type=float, default=1.0, help="长度容差秒")
    ap.add_argument("--fade", type=float, default=2.0, help="交叉淡化秒")
    ap.add_argument("--bitrate", default="256k")
    ap.add_argument("--peak", type=float, default=-1.0,
                    help="目标峰值 dBFS，超了自动衰减留余量")
    ap.add_argument("--preview", type=int, default=0,
                    help="额外导出连播 N 遍的试听文件，用来听接缝")
    ap.add_argument("--keep-wav", action="store_true", help="保留无损 wav")
    a = ap.parse_args()

    src = os.path.expanduser(a.input)
    out = os.path.expanduser(a.output)
    x = decode(src, ANALYZE_SR)
    dur = len(x) / ANALYZE_SR
    print("素材 %s · %.2fs" % (os.path.basename(src), dur))

    t0, L, score, why = pick_window(x, dur, a.length, a.tol, a.fade)
    print("循环长度 %.3fs（%s）· 起点 %.2fs · 接缝差异分 %.3f" % (L, why, t0, score))

    wav = build(src, out, t0, L, a.fade, a.bitrate, a.peak)
    verify(wav)

    if a.preview > 1:
        pv = os.path.splitext(out)[0] + "_x%d.mp3" % a.preview
        subprocess.run(["ffmpeg", "-hide_banner", "-v", "error", "-y",
                        "-stream_loop", str(a.preview - 1), "-i", wav,
                        "-c:a", "libmp3lame", "-b:a", a.bitrate, pv], check=True)
        print("接缝试听（连播 %d 遍）: %s" % (a.preview, pv))
    if not a.keep_wav:
        os.remove(wav)
    else:
        print("无损: %s" % wav)
    print("成品: %s" % out)


if __name__ == "__main__":
    main()