"""music2tap: 将音频文件的 RMS 音量包络转换为 RichTap 标准 .he 触觉文件。

支持批量模式（扫描 music/ 目录）与单文件模式（--input/--output）。
"""

from __future__ import annotations

import argparse
import json
import sys
import warnings
from datetime import datetime, timezone
from pathlib import Path

# 抑制 librosa 解码 mp3 时回退 audioread 引发的两条已知告警：
# - PySoundFile 对部分 mp3 头部有问题会失败回退（结果正确，不影响包络）
# - audioread 后端在 librosa 0.10 标记为 future-deprecated（1.0 才移除）
warnings.filterwarnings("ignore", message="PySoundFile failed.*")
warnings.filterwarnings("ignore", category=FutureWarning, module="librosa.*")

import librosa  # noqa: E402  warnings 过滤需在 librosa 导入前生效
import numpy as np  # noqa: E402

# 支持的音频扩展名集合，用于批量模式过滤
AUDIO_EXTS = {".mp3", ".wav", ".flac", ".ogg", ".m4a"}

# 默认参数：22050Hz 单声道、100Hz 帧率（10ms 一帧）——触觉反馈常用粒度
DEFAULT_SR = 22050
DEFAULT_FRAME_RATE = 100

# RichTap continuous 事件的固定中频，dense 模式整曲使用统一频率
DEFAULT_FREQUENCY = 50

# 谱心映射上限：人耳对 0-8kHz 最敏感、流行乐谱心绝大多数落在 200-6000Hz，
# 取 8000Hz 作上界后线性映射到 0-100，能避免少数尖锐高频把整体压扁
SPECTRAL_CENTROID_MAX_HZ = 8000.0

# 默认关键帧阈值：相对上一关键帧的振幅/频率最小变化（0-100 单位）
DEFAULT_INTENSITY_THRESHOLD = 2
DEFAULT_FREQUENCY_THRESHOLD = 2


def extract_envelope(
    audio_path: Path,
    sr: int = DEFAULT_SR,
    frame_rate: int = DEFAULT_FRAME_RATE,
) -> tuple[np.ndarray, int, int]:
    """提取音频文件的归一化 RMS 包络。

    使用 librosa 加载音频并计算逐帧 RMS，再按全曲最大值线性归一化到 0-100
    整数区间，匹配 RichTap .he 格式对 ``Intensity`` 字段的取值要求。

    :param audio_path: 输入音频路径，支持 mp3/wav/flac/ogg/m4a
    :param sr: 目标采样率（Hz），统一重采样以保证 hop_length 推导一致
    :param frame_rate: 期望的输出帧率（Hz），决定时间分辨率
    :returns: ``(intensities, frame_interval_ms, duration_ms)``
        intensities 为 ``np.uint8`` 数组，每元素是 0-100 的整数振幅；
        frame_interval_ms 为相邻两帧的时间间隔（毫秒）；
        duration_ms 为音频总时长（毫秒）。
    :raises FileNotFoundError: 音频文件不存在
    """
    if not audio_path.exists():
        raise FileNotFoundError(f"音频文件不存在: {audio_path}")

    # mono=True 以避免立体声两通道导致的 RMS 数值偏差；统一 sr 以稳定 hop_length
    y, sr = librosa.load(str(audio_path), sr=sr, mono=True)

    # 由采样率与目标帧率推导 hop_length：sr / frame_rate ≈ 每帧样本数
    # 例如 22050Hz、100Hz → hop_length=220，对应 10ms 一帧
    hop_length = max(1, int(round(sr / frame_rate)))
    frame_length = 2048  # librosa 默认值，提供稳定的能量估计窗口

    rms = librosa.feature.rms(y=y, frame_length=frame_length, hop_length=hop_length)[0]

    # 按全曲最大 RMS 归一化：保证最响段映射到 100，弱段保持比例关系
    # 防 0 除：极端情况下整段静音则全部输出 0
    peak = float(rms.max()) if rms.size > 0 else 0.0
    if peak > 0:
        normalized = np.clip(rms / peak * 100.0, 0, 100)
    else:
        normalized = np.zeros_like(rms)

    intensities = np.round(normalized).astype(np.uint8)

    # frame_interval_ms 用 hop_length 反推真实间隔，避免与目标帧率有舍入误差时漂移
    frame_interval_ms = int(round(hop_length * 1000 / sr))
    duration_ms = int(round(len(y) * 1000 / sr))

    return intensities, frame_interval_ms, duration_ms


def extract_envelope_with_freq(
    audio_path: Path,
    sr: int = DEFAULT_SR,
    frame_rate: int = DEFAULT_FRAME_RATE,
) -> tuple[np.ndarray, np.ndarray, int, int]:
    """同时提取归一化 RMS 振幅与谱心频率序列。

    与 :func:`extract_envelope` 共用 hop_length 推导逻辑，因此返回的
    intensities / frequencies 等长且时间戳一致，可直接用于 keyframe 抽稀。

    :param audio_path: 输入音频路径
    :param sr: 目标采样率（Hz）
    :param frame_rate: 期望帧率（Hz）
    :returns: ``(intensities, frequencies, frame_interval_ms, duration_ms)``
        intensities/frequencies 都是 ``np.uint8`` 数组，取值 0-100。
    :raises FileNotFoundError: 音频文件不存在
    """
    if not audio_path.exists():
        raise FileNotFoundError(f"音频文件不存在: {audio_path}")

    y, sr = librosa.load(str(audio_path), sr=sr, mono=True)

    hop_length = max(1, int(round(sr / frame_rate)))
    n_fft = 2048

    rms = librosa.feature.rms(y=y, frame_length=n_fft, hop_length=hop_length)[0]
    centroid = librosa.feature.spectral_centroid(
        y=y, sr=sr, n_fft=n_fft, hop_length=hop_length
    )[0]

    # RMS 按全曲峰值归一化；静音段全 0，避免除零
    peak = float(rms.max()) if rms.size > 0 else 0.0
    if peak > 0:
        intensity_norm = np.clip(rms / peak * 100.0, 0, 100)
    else:
        intensity_norm = np.zeros_like(rms)
    intensities = np.round(intensity_norm).astype(np.uint8)

    # 谱心按 0-8kHz 上界做硬截断映射，越界值压到 100
    freq_norm = np.clip(centroid / SPECTRAL_CENTROID_MAX_HZ * 100.0, 0, 100)
    frequencies = np.round(freq_norm).astype(np.uint8)

    # rms 与 spectral_centroid 在 librosa 里默认 center=True，长度一致；
    # 不一致时按较短者截断，避免下游索引越界
    n = min(len(intensities), len(frequencies))
    intensities = intensities[:n]
    frequencies = frequencies[:n]

    frame_interval_ms = int(round(hop_length * 1000 / sr))
    duration_ms = int(round(len(y) * 1000 / sr))

    return intensities, frequencies, frame_interval_ms, duration_ms


def select_keyframes(
    intensities: np.ndarray,
    frequencies: np.ndarray,
    intensity_threshold: int = DEFAULT_INTENSITY_THRESHOLD,
    frequency_threshold: int = DEFAULT_FREQUENCY_THRESHOLD,
) -> list[int]:
    """贪心抽稀，返回保留的关键帧索引列表。

    规则：相对上一个已保留关键帧，若 ``|ΔI| ≥ intensity_threshold`` 或
    ``|ΔF| ≥ frequency_threshold`` 则记录当前帧为新关键帧。
    首末必留——首点决定起始状态、末点决定 Duration 边界，缺失会让下位机
    的线性插值跨越整曲终点。

    选用贪心而非 RDP：下位机已对 continuous 事件做线性插值，贪心算法的
    “对上一关键帧最大偏差”语义恰好等于下位机插值误差的上界，与播放行为一致。

    :param intensities: 0-100 整数振幅序列
    :param frequencies: 0-100 整数频率序列，长度需与 intensities 一致
    :param intensity_threshold: 振幅变化阈值（含等号）
    :param frequency_threshold: 频率变化阈值（含等号）
    :returns: 保留索引升序列表（含首末）
    """
    n = len(intensities)
    if n == 0:
        return []
    if n == 1:
        return [0]

    kept: list[int] = [0]
    last_i = int(intensities[0])
    last_f = int(frequencies[0])

    # 仅遍历中间点，末点单独追加，保证 Duration 对齐
    for idx in range(1, n - 1):
        cur_i = int(intensities[idx])
        cur_f = int(frequencies[idx])
        if (
            abs(cur_i - last_i) >= intensity_threshold
            or abs(cur_f - last_f) >= frequency_threshold
        ):
            kept.append(idx)
            last_i, last_f = cur_i, cur_f

    if kept[-1] != n - 1:
        kept.append(n - 1)
    return kept


def build_he(
    intensities: np.ndarray,
    frame_interval_ms: int,
    duration_ms: int,
    source_name: str,
    frequencies: np.ndarray | None = None,
    indices: list[int] | None = None,
) -> dict:
    """根据归一化包络构造 RichTap 标准 .he 字典。

    使用单个 continuous 事件承载整曲包络，``Curve`` 数组的每点包含
    Time（绝对 ms）、Frequency（0-100）、Intensity（0-100）。

    :param intensities: 0-100 的整数振幅数组
    :param frame_interval_ms: 单帧时间步长（毫秒），用于把序号换算成 Time
    :param duration_ms: 整曲时长（毫秒），写入 Event.Duration
    :param source_name: 源文件名，用于 Metadata.Description
    :param frequencies: 可选，与 intensities 等长的 0-100 频率数组；
        为 None 时整曲使用 :data:`DEFAULT_FREQUENCY`（dense 模式行为）
    :param indices: 可选，要写入 Curve 的索引列表（关键帧抽稀结果）；
        为 None 时写入全部点
    :returns: 可直接 json.dump 的 dict
    """
    sel = indices if indices is not None else range(len(intensities))
    # Time 以 frame_interval_ms 为步长换算：dense 是等间距，keyframe 仅在选中索引上落点
    curve = [
        {
            "Time": int(i * frame_interval_ms),
            "Frequency": int(frequencies[i])
            if frequencies is not None
            else DEFAULT_FREQUENCY,
            "Intensity": int(intensities[i]),
        }
        for i in sel
    ]

    # librosa 的 center=True 会让最后一帧中心略超过音频长度，
    # Duration 必须 ≥ 最后关键帧的 Time，否则下位机会把末点视为越界
    if curve:
        duration_ms = max(duration_ms, curve[-1]["Time"])

    return {
        "Metadata": {
            "Version": 1,
            "Created": datetime.now(timezone.utc).isoformat(timespec="seconds"),
            "Description": f"Generated from {source_name} by music2tap",
            "ChannelNumber": 1,
        },
        "PatternList": [
            {
                "AbsoluteTime": 0,
                "Pattern": [
                    {
                        "Event": {
                            "Type": "continuous",
                            "RelativeTime": 0,
                            "Duration": duration_ms,
                            "Parameters": {
                                "Intensity": 100,
                                "Frequency": DEFAULT_FREQUENCY,
                                "Curve": curve,
                            },
                        }
                    }
                ],
            }
        ],
    }


def write_he(he_dict: dict, output_path: Path) -> None:
    """将 .he 字典以 UTF-8 JSON 形式写入磁盘。

    :param he_dict: 由 :func:`build_he` 生成的字典
    :param output_path: 输出路径，父目录会自动创建
    """
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with output_path.open("w", encoding="utf-8") as f:
        # ensure_ascii=False 保留 Description 中可能的中文；indent=2 便于人工排查
        json.dump(he_dict, f, ensure_ascii=False, indent=2)


def process_file(
    input_path: Path,
    output_path: Path,
    sr: int = DEFAULT_SR,
    frame_rate: int = DEFAULT_FRAME_RATE,
    method: str = "dense",
    intensity_threshold: int = DEFAULT_INTENSITY_THRESHOLD,
    frequency_threshold: int = DEFAULT_FREQUENCY_THRESHOLD,
) -> None:
    """单文件转换流水线：提包络 → 构造 .he → 写盘。

    :param input_path: 输入音频文件
    :param output_path: 输出 .he 文件路径
    :param sr: 目标采样率
    :param frame_rate: 期望帧率
    :param method: ``"dense"`` 等间距逐帧；``"keyframe"`` 仅写振幅/频率拐点
    :param intensity_threshold: keyframe 模式下的 |ΔI| 阈值
    :param frequency_threshold: keyframe 模式下的 |ΔF| 阈值
    """
    if method == "dense":
        intensities, frame_interval_ms, duration_ms = extract_envelope(
            input_path, sr=sr, frame_rate=frame_rate
        )
        he_dict = build_he(intensities, frame_interval_ms, duration_ms, input_path.name)
        write_he(he_dict, output_path)
        print(
            f"[OK] {input_path.name} -> {output_path} "
            f"(dense, {len(intensities)} points, {duration_ms} ms, "
            f"step {frame_interval_ms} ms)"
        )
        return

    if method == "keyframe":
        intensities, frequencies, frame_interval_ms, duration_ms = (
            extract_envelope_with_freq(input_path, sr=sr, frame_rate=frame_rate)
        )
        kept = select_keyframes(
            intensities,
            frequencies,
            intensity_threshold=intensity_threshold,
            frequency_threshold=frequency_threshold,
        )
        he_dict = build_he(
            intensities,
            frame_interval_ms,
            duration_ms,
            input_path.name,
            frequencies=frequencies,
            indices=kept,
        )
        write_he(he_dict, output_path)
        ratio = (len(kept) / len(intensities) * 100.0) if len(intensities) else 0.0
        print(
            f"[OK] {input_path.name} -> {output_path} "
            f"(keyframe, kept {len(kept)}/{len(intensities)} = {ratio:.1f}%, "
            f"{duration_ms} ms, step {frame_interval_ms} ms, "
            f"I_thr={intensity_threshold}, F_thr={frequency_threshold})"
        )
        return

    raise ValueError(f"未知 method: {method!r}（应为 'dense' 或 'keyframe'）")


def process_directory(
    input_dir: Path,
    output_dir: Path,
    sr: int = DEFAULT_SR,
    frame_rate: int = DEFAULT_FRAME_RATE,
    method: str = "dense",
    intensity_threshold: int = DEFAULT_INTENSITY_THRESHOLD,
    frequency_threshold: int = DEFAULT_FREQUENCY_THRESHOLD,
) -> int:
    """批量转换目录下所有音频文件。

    :param input_dir: 待扫描的音频目录
    :param output_dir: 输出 .he 文件目录
    :param sr: 目标采样率
    :param frame_rate: 期望帧率
    :param method: ``"dense"`` 或 ``"keyframe"``
    :param intensity_threshold: keyframe 模式下的 |ΔI| 阈值
    :param frequency_threshold: keyframe 模式下的 |ΔF| 阈值
    :returns: 成功处理的文件数
    """
    if not input_dir.is_dir():
        print(f"[ERR] 输入目录不存在: {input_dir}", file=sys.stderr)
        return 0

    count = 0
    for audio_file in sorted(input_dir.iterdir()):
        if audio_file.suffix.lower() not in AUDIO_EXTS:
            continue
        out_file = output_dir / (audio_file.stem + ".he")
        try:
            process_file(
                audio_file,
                out_file,
                sr=sr,
                frame_rate=frame_rate,
                method=method,
                intensity_threshold=intensity_threshold,
                frequency_threshold=frequency_threshold,
            )
            count += 1
        except Exception as exc:
            # 单个文件失败不阻断批量任务，但需明确报错以便排查
            print(f"[ERR] {audio_file.name}: {exc}", file=sys.stderr)
    return count


def main() -> None:
    """CLI 入口：解析参数并分发到批量或单文件流水线。"""
    parser = argparse.ArgumentParser(
        description="提取音频 RMS 包络并生成 RichTap 标准 .he 触觉文件"
    )
    parser.add_argument("--input", type=Path, help="单文件模式：输入音频路径")
    parser.add_argument("--output", type=Path, help="单文件模式：输出 .he 路径")
    parser.add_argument(
        "--input-dir",
        type=Path,
        default=Path("music"),
        help="批量模式：输入目录（默认 music/）",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("output"),
        help="批量模式：输出目录（默认 output/）",
    )
    parser.add_argument("--sr", type=int, default=DEFAULT_SR, help="目标采样率 (Hz)")
    parser.add_argument(
        "--frame-rate",
        type=int,
        default=DEFAULT_FRAME_RATE,
        help="输出帧率 (Hz)，10ms→100",
    )
    parser.add_argument(
        "--method",
        choices=("dense", "keyframe"),
        default="dense",
        help="dense=逐帧；keyframe=只写振幅/频率拐点（下位机线性插值）",
    )
    parser.add_argument(
        "--intensity-threshold",
        type=int,
        default=DEFAULT_INTENSITY_THRESHOLD,
        help="keyframe 模式下的 |ΔIntensity| 阈值（0-100）",
    )
    parser.add_argument(
        "--frequency-threshold",
        type=int,
        default=DEFAULT_FREQUENCY_THRESHOLD,
        help="keyframe 模式下的 |ΔFrequency| 阈值（0-100）",
    )
    args = parser.parse_args()

    common = dict(
        sr=args.sr,
        frame_rate=args.frame_rate,
        method=args.method,
        intensity_threshold=args.intensity_threshold,
        frequency_threshold=args.frequency_threshold,
    )

    # 单文件模式优先：只要给了 --input 就走单文件分支
    if args.input is not None:
        out = args.output or (args.output_dir / (args.input.stem + ".he"))
        process_file(args.input, out, **common)
        return

    # 批量模式：扫描 input-dir
    n = process_directory(args.input_dir, args.output_dir, **common)
    print(f"[DONE] processed {n} file(s) into {args.output_dir}")


if __name__ == "__main__":
    main()