import math

import numpy as np

from .model import SAMPLE_RATE


class Compressor:
    """Feedforward dynamic range compressor + brick-wall limiter for speech.

    Per sample: track an attack/release-smoothed envelope of |x|, compute gain
    reduction above the threshold, apply makeup gain, then hard-limit to the ceiling.
    """

    def __init__(
        self,
        threshold_db: float = -20.0,
        ratio: float = 4.0,
        attack_ms: float = 5.0,
        release_ms: float = 80.0,
        makeup_db: float = 12.0,
        ceiling: float = 10 ** (-1.0 / 20),  # -1 dBFS
        sample_rate: int = SAMPLE_RATE,
    ):
        self.threshold_db = threshold_db
        self.ratio = ratio
        self.makeup_gain = 10 ** (makeup_db / 20)
        self.ceiling = ceiling
        self.knee = 1.0 - 1.0 / ratio
        self.a_att = math.exp(-1.0 / (attack_ms * 0.001 * sample_rate))
        self.a_rel = math.exp(-1.0 / (release_ms * 0.001 * sample_rate))
        self.envelope = 0.0

    def process(self, x: np.ndarray) -> np.ndarray:
        abs_x = np.abs(x)
        env_out = np.empty_like(x)
        e = self.envelope
        a_att = self.a_att
        a_rel = self.a_rel
        for i in range(len(x)):
            target = abs_x[i]
            coef = a_att if target > e else a_rel
            e = coef * e + (1.0 - coef) * target
            env_out[i] = e
        self.envelope = e

        env_db = 20.0 * np.log10(np.maximum(env_out, 1e-10))
        over = env_db - self.threshold_db
        gr_db = np.where(over > 0, over * self.knee, 0.0)
        gain = 10 ** (-gr_db / 20.0) * self.makeup_gain
        y = x * gain
        return np.clip(y, -self.ceiling, self.ceiling).astype(np.float32)
