"""Chromagram 特征提取。

流程：
1. 音频标准化：ffmpeg 转 22050Hz / Mono / WAV
2. librosa 加载音频
3. librosa.feature.chroma_cens() 提取 12×T Chromagram（CENS，对速度/音色鲁棒）
4. 主音对齐：将能量最大的音级滚至第 0 行，实现转调不变性
5. scipy.signal.resample(chroma, 128, axis=1) 时间归一化到 12×128
6. .flatten() 展开为 1536 维向量
"""

import logging
import os
import subprocess
import tempfile

import librosa
import numpy as np
from scipy.signal import resample

logger = logging.getLogger(__name__)

# 目标采样率和时间帧数
TARGET_SR = 22050
TARGET_FRAMES = 128
VECTOR_DIM = 12 * TARGET_FRAMES  # 1536


def _normalize_audio_ffmpeg(audio_path: str, output_path: str) -> None:
    """使用 ffmpeg 将音频标准化为 22050Hz / Mono / WAV。"""
    cmd = [
        "ffmpeg",
        "-y",
        "-i", audio_path,
        "-ar", str(TARGET_SR),
        "-ac", "1",
        "-f", "wav",
        output_path,
    ]
    result = subprocess.run(
        cmd,
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        raise RuntimeError(f"ffmpeg 转换失败: {result.stderr}")


def extract_chroma_feature(audio_path: str) -> np.ndarray:
    """从音频文件提取 1536 维 Chromagram 特征向量。

    Args:
        audio_path: 音频文件路径。

    Returns:
        shape 为 (1536,) 的 numpy 数组。

    Raises:
        FileNotFoundError: 音频文件不存在。
        RuntimeError: ffmpeg 转换失败。
    """
    if not os.path.isfile(audio_path):
        raise FileNotFoundError(f"音频文件不存在: {audio_path}")

    # 1. 音频标准化：ffmpeg 转 WAV
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
        tmp_wav = tmp.name

    try:
        _normalize_audio_ffmpeg(audio_path, tmp_wav)

        # 2. librosa 加载音频
        y, _sr = librosa.load(tmp_wav, sr=TARGET_SR, mono=True)

        # 3. 提取 CENS Chromagram (12×T)，对速度变化和音色具有更强鲁棒性
        chroma = librosa.feature.chroma_cens(y=y, sr=TARGET_SR)

        # 4. 主音对齐：将全局能量最大的音级循环滚至第 0 行，实现转调不变性
        tonic = int(np.argmax(chroma.sum(axis=1)))
        if tonic != 0:
            chroma = np.roll(chroma, -tonic, axis=0)

        # 5. 时间归一化到 12×128
        if chroma.shape[1] != TARGET_FRAMES:
            chroma = resample(chroma, TARGET_FRAMES, axis=1)

        # 6. 展开为 1536 维向量
        feature = chroma.flatten().astype(np.float32)

        assert feature.shape == (VECTOR_DIM,), (
            f"特征维度错误: 期望 {VECTOR_DIM}, 实际 {feature.shape}"
        )

        return feature
    finally:
        # 清理临时文件
        if os.path.exists(tmp_wav):
            os.remove(tmp_wav)


def extract_chroma_matrix(audio_path: str) -> np.ndarray:
    """从音频文件提取 12×128 Chromagram 矩阵（未展平，供 DTW 精排使用）。

    Returns:
        shape 为 (12, 128) 的 numpy 数组，已做主音对齐。
    """
    feature = extract_chroma_feature(audio_path)
    return feature.reshape(12, TARGET_FRAMES)