"""파일도 서브프로세스도 없이 메모리에서 처리하는 오디오 유틸."""
import io
from fractions import Fraction

import av
import numpy as np

SAMPLE_RATE = 44100                    # 렌더 믹스 기준
MP3_BIT_RATE = 192000
INT16_FULL_SCALE = 32768


def decode_mp3(data: bytes, sample_rate: int = SAMPLE_RATE) -> np.ndarray:
    """mp3 → 모노 s16 PCM. 소스 샘플레이트가 달라도 리샘플해 맞춘다."""
    container = av.open(io.BytesIO(data))
    resampler = av.audio.resampler.AudioResampler(
        format="s16", layout="mono", rate=sample_rate)
    chunks = []
    for frame in container.decode(audio=0):
        for resampled in resampler.resample(frame):
            chunks.append(resampled.to_ndarray().reshape(-1))
    for resampled in resampler.resample(None):        # 남은 것 밀어내기
        chunks.append(resampled.to_ndarray().reshape(-1))
    container.close()
    return np.concatenate(chunks) if chunks else np.zeros(0, np.int16)


def encode_mp3(pcm: np.ndarray, sample_rate: int = SAMPLE_RATE) -> bytes:
    buffer = io.BytesIO()
    container = av.open(buffer, "w", format="mp3")
    stream = container.add_stream("mp3", rate=sample_rate)
    stream.layout = "mono"
    stream.bit_rate = MP3_BIT_RATE
    frame = av.AudioFrame.from_ndarray(pcm.reshape(1, -1), format="s16", layout="mono")
    frame.rate = sample_rate
    for packet in stream.encode(frame):
        container.mux(packet)
    for packet in stream.encode():
        container.mux(packet)
    container.close()
    return buffer.getvalue()


def decode_audio(media: bytes, sample_rate: int = SAMPLE_RATE) -> np.ndarray | None:
    """비디오·오디오 컨테이너에서 오디오 트랙을 꺼낸다. 트랙이 없으면 None."""
    container = av.open(io.BytesIO(media))
    try:
        if not container.streams.audio:
            return None
        resampler = av.audio.resampler.AudioResampler(
            format="s16", layout="mono", rate=sample_rate)
        chunks = []
        for frame in container.decode(audio=0):
            for resampled in resampler.resample(frame):
                chunks.append(resampled.to_ndarray().reshape(-1))
        for resampled in resampler.resample(None):
            chunks.append(resampled.to_ndarray().reshape(-1))
    finally:
        container.close()
    return np.concatenate(chunks) if chunks else None


def apply_limiter(pcm: np.ndarray, sample_rate: int = SAMPLE_RATE,
                  limit: float = 0.94, attack: int = 5, release: int = 60) -> np.ndarray:
    """ffmpeg alimiter를 그대로 태운다. 리미터가 컴프레서처럼 작동해 목소리가 또렷해진다."""
    graph = av.filter.Graph()
    source = graph.add_abuffer(format="s16", layout="mono", sample_rate=sample_rate,
                               time_base=Fraction(1, sample_rate))
    limiter = graph.add("alimiter", f"limit={limit}:attack={attack}:release={release}")
    # alimiter는 부동소수로 처리하므로 싱크 앞에서 s16으로 되돌린다
    to_s16 = graph.add("aformat", "sample_fmts=s16:channel_layouts=mono")
    sink = graph.add("abuffersink")
    source.link_to(limiter)
    limiter.link_to(to_s16)
    to_s16.link_to(sink)
    graph.configure()

    frame = av.AudioFrame.from_ndarray(pcm.reshape(1, -1), format="s16", layout="mono")
    frame.rate = sample_rate
    frame.time_base = Fraction(1, sample_rate)
    graph.push(frame)
    graph.push(None)

    chunks = []
    while True:
        try:
            chunks.append(sink.pull().to_ndarray().reshape(-1))
        except (av.error.EOFError, av.error.BlockingIOError):
            break
    return np.concatenate(chunks).astype(np.int16) if chunks else pcm


def trim_silence(pcm: np.ndarray, threshold_db: float = -45,
                 keep_lead: float = 0.05, keep_trail: float = 0.12,
                 sample_rate: int = SAMPLE_RATE) -> np.ndarray:
    """앞뒤 무음만 깎는다. 말소리 사이의 쉼은 보존한다."""
    threshold = INT16_FULL_SCALE * 10 ** (threshold_db / 20)
    loud = np.flatnonzero(np.abs(pcm) > threshold)
    if len(loud) == 0:
        return pcm[:0]
    start = max(0, loud[0] - round(keep_lead * sample_rate))
    end = min(len(pcm), loud[-1] + 1 + round(keep_trail * sample_rate))
    return pcm[start:end]