#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""原创五声音阶配乐、低量音效、水声及中文 TTS 混音（48kHz 双声道）。"""
import json
import math
import subprocess
import wave
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
STORY = json.loads((ROOT / "story.json").read_text(encoding="utf-8"))
TIMING = json.loads((ROOT / "timing.json").read_text(encoding="utf-8"))
SR = 48000
N = round(TIMING["duration"] * SR)
rng = np.random.default_rng(20261010)  # 固定种子，输出可重复。
voice = np.zeros((N, 2), dtype=np.float32)
score = np.zeros((N, 2), dtype=np.float32)
fx = np.zeros((N, 2), dtype=np.float32)


def add(dest: np.ndarray, samples: np.ndarray, start_s: float, gain=1.0, pan=0.0) -> None:
    start = int(round(start_s * SR))
    if start < 0 or start >= N:
        return
    length = min(len(samples), N - start)
    samples = np.asarray(samples[:length], dtype=np.float32) * gain
    dest[start:start + length, 0] += samples * math.sqrt((1.0 - pan) / 2)
    dest[start:start + length, 1] += samples * math.sqrt((1.0 + pan) / 2)


def pluck(freq: float, seconds=2.0) -> np.ndarray:
    t = np.arange(round(seconds * SR), dtype=np.float32) / SR
    ph = 2 * np.pi * freq * t
    return ((np.sin(ph) + .27 * np.sin(2 * ph) + .08 * np.sin(3 * ph)) *
            np.exp(-2.0 * t) * np.clip(t / .008, 0, 1)).astype(np.float32)


def flute(freq: float, seconds=2.3) -> np.ndarray:
    t = np.arange(round(seconds * SR), dtype=np.float32) / SR
    ph = 2 * np.pi * freq * t + .04 * np.sin(2 * np.pi * 5.0 * t)
    fade = np.clip(t / .14, 0, 1) * np.clip((seconds - t) / .65, 0, 1)
    return ((np.sin(ph) + .17 * np.sin(2 * ph)) * fade).astype(np.float32)


def add_noise(at: float, seconds: float, gain: float, freq_smooth=18) -> None:
    n = round(seconds * SR)
    z = rng.normal(0, 1, n).astype(np.float32)
    if freq_smooth > 1:
        z = np.convolve(z, np.ones(freq_smooth, dtype=np.float32) / freq_smooth, "same")
    t = np.arange(n, dtype=np.float32) / SR
    env = np.exp(-t * 4.0) * np.clip(t / .012, 0, 1)
    add(fx, z * env, at, gain=gain, pan=.13)


melody = [293.66, 329.63, 392.0, 440.0, 392.0, 329.63, 293.66, 261.63]
for i, start in enumerate(np.arange(0.3, TIMING["duration"] - 1.0, 1.38)):
    add(score, pluck(melody[i % len(melody)], 2.4), float(start), gain=.049,
        pan=((-1) ** i) * .18)
    if i % 4 == 0:
        add(score, pluck(melody[i % len(melody)] / 2, 3.3), float(start), gain=.034, pan=-.23)
    if i % 5 == 1:
        add(score, flute(melody[(i + 3) % len(melody)], 2.5), float(start) + .15,
            gain=.017, pan=.24)

for sc, seg in zip(STORY["scenes"], TIMING["scenes"]):
    event = sc.get("sfx", "")
    at = float(seg["start"]) + min(2.4, max(0.6, seg["duration"] * .39))
    if event == "splash":
        add_noise(at, 1.7, gain=.46, freq_smooth=15)
        # 水面扩散，低频音效含能量小，避免喧宾夺主。
        t = np.arange(round(.6 * SR), dtype=np.float32) / SR
        ripple = (np.sin(2 * np.pi * (48 - 20 * t) * t) * np.exp(-8 * t)).astype(np.float32)
        add(fx, ripple, at, gain=.045)
    elif event in ("whoosh", "ripple"):
        add_noise(at, .85, gain=.18, freq_smooth=55 if event == "ripple" else 27)
    elif event == "sparkle":
        add(fx, flute(780, 1.7), at, gain=.045, pan=.23)

# 将每个 TTS 语音转成统一的单声道浮点格式，再按真实开始时间叠加。
for scene in TIMING["scenes"]:
    # 每幕可有任意数量的角色对白：使用逐段、实测的 TTS 时间轴混音。
    # 兼容旧版单句场景；不能同时混两套，否则会导致旁白被重复播放。
    lines = scene.get("lines") or [scene]
    for seg in lines:
        source = ROOT / seg["voice_file"]
        raw = subprocess.check_output([
            "ffmpeg", "-hide_banner", "-loglevel", "error", "-i", str(source),
            "-ac", "1", "-ar", str(SR), "-f", "f32le", "-"
        ])
        audio = np.frombuffer(raw, dtype="<f4").copy()
        if not len(audio):
            raise ValueError(f"空语音: {source} / 角色：{seg.get('role', 'narrator')}")
        amp = float(np.max(np.abs(audio)))
        gain = min(1.7, .63 / max(amp, .0001))
        add(voice, audio, float(seg["voice_start"]), gain=gain, pan=0)

# 配乐整体偏低，语音活跃时额外自动压低背景；在片头片尾淡入淡出。
voice_rms = np.sqrt(np.mean(voice.astype(np.float64) ** 2, axis=1))
active = voice_rms > .009
# 近似语音包络，与语音交叉处更平顺。
gate = active.astype(np.float32)
width = 7200
# 使用前缀和 O(N) 平滑，避免百万样本 × 7200 的昂贵直接卷积。
cumulative = np.concatenate((np.zeros(1, dtype=np.float64), np.cumsum(gate, dtype=np.float64)))
idx = np.arange(N)
lo = np.maximum(idx - width // 2, 0)
hi = np.minimum(idx + width // 2, N)
smooth_gate = ((cumulative[hi] - cumulative[lo]) / np.maximum(hi - lo, 1)).astype(np.float32)
score *= (.84 - .26 * np.clip(smooth_gate, 0, 1))[:, None]
times = np.arange(N, dtype=np.float32) / SR
fade = np.minimum(1.0, times / .72) * np.minimum(1.0, (TIMING["duration"] - times) / 1.6)
mix = (voice + score + fx) * fade[:, None]
max_abs = float(np.max(np.abs(mix)))
if max_abs > .91:
    mix *= .91 / max_abs
out = ROOT / "assets/film-mix.wav"
out.parent.mkdir(exist_ok=True)
with wave.open(str(out), "wb") as f:
    f.setnchannels(2)
    f.setsampwidth(2)
    f.setframerate(SR)
    f.writeframes((mix * 32767).astype("<i2").tobytes())
print(json.dumps({"audio": str(out), "seconds": TIMING["duration"],
                  "mock_voice": TIMING.get("mock", False), "peak_before_limiter": round(max_abs, 4)},
                 ensure_ascii=False, indent=2))