import io
import os
import pickle

import soundfile as sf
import numpy as np
from pydub import AudioSegment
from pyloudnorm import Meter

os.chdir(os.path.dirname(os.path.abspath(__file__)))

def normalize_audio_loudness(data: bytes, target_loudness: float = -23.0) -> bytes:
    audio = AudioSegment.from_file(io.BytesIO(data), format='mp3')    
    meter = Meter(audio.frame_rate)  # 创建一个响度计量器
    sr = audio.frame_rate
    samples = audio.get_array_of_samples()
    audio = np.array(samples, dtype=np.float64)

    # 测量积分响度
    loudness = meter.integrated_loudness(audio)

    # 计算增益
    gain_db = target_loudness - loudness
    gain_linear = 10 ** (gain_db / 20.0)

    # 应用增益
    balanced_audio = audio * gain_linear

    # 应用软限幅以防止削波
    balanced_audio = np.tanh(balanced_audio)

    # 将numpy数组转换回bytes
    balanced_audio = (balanced_audio * 32767).astype(np.int16)
    byte_io = io.BytesIO()
    sf.write(byte_io, balanced_audio, sr, format='mp3')
    normalized_audio_bytes = byte_io.getvalue()

    return normalized_audio_bytes

if __name__ == "__main__":
    normalize_audio_loudness()