//! AEC3 — Adaptive Echo Canceller with delay estimation.
//!
//! ## Algorithm
//!
//! Time-domain NLMS (Normalised Least Mean Squares) adaptive filter
//! with cross-correlation delay estimation:
//!
//! 1. **Delay estimation** — cross-correlates the microphone and
//!    render-reference signals to find the bulk acoustic delay.
//!    Tracked with exponential smoothing over a 16-block history.
//!
//! 2. **NLMS adaptive filter** — a time-domain FIR filter of length
//!    FILTER_TAPS adapts sample-by-sample using the NLMS rule:
//!    `w[n+1] = w[n] + μ · e[n] · x[n] / (||x[n]||² + δ)`
//!    where x[n] is the delayed reference vector and e[n] = mic[n] - ŷ[n].
//!
//! 3. **Post-filter** — residual echo suppression using ERLE.
//!
//! ## Realtime safety
//!
//! All state is pre-allocated. No heap allocation, no I/O, no blocking
//! inside `process_capture` or `process_render`.

#![allow(clippy::needless_range_loop)]

use super::super::FRAME_SAMPLES;

/// Adaptive filter length in taps (80 ms at 48 kHz).
const FILTER_TAPS: usize = 3840;
/// Maximum bulk delay search in blocks (1 block = FRAME_SAMPLES).
const MAX_DELAY_BLOCKS: usize = 16;
/// NLMS step size μ.
const MU: f32 = 0.05;
/// NLMS regularisation δ.
const NLMS_REG: f32 = 1e-3;
/// Post-filter suppression floor.
const POST_FILTER_FLOOR: f32 = 0.1;
/// ERLE smoothing coefficient.
const ERLE_ALPHA: f32 = 0.05;
/// Minimum ERLE (linear) before post-filter activates (6 dB).
const MIN_ERLE: f32 = 2.0;
/// Reference buffer length: delay line + filter taps.
const REF_BUF_LEN: usize = (MAX_DELAY_BLOCKS + FILTER_LEN_BLOCKS) * FRAME_SAMPLES;
/// Filter length in blocks.
const FILTER_LEN_BLOCKS: usize = FILTER_TAPS / FRAME_SAMPLES;

/// Adaptive echo canceller.
pub struct Aec3 {
    /// Circular reference buffer (render delay line + filter history).
    ref_buf: Vec<f32>,
    /// Write head into ref_buf.
    ref_head: usize,
    /// Estimated bulk delay in samples.
    bulk_delay: usize,
    /// Cross-correlation per candidate delay block.
    xcorr: Box<[f32; MAX_DELAY_BLOCKS]>,
    /// Adaptive filter weights.
    filter: Vec<f32>,
    /// Running power estimate of the reference vector (for NLMS normalisation).
    ref_power: f32,
    /// ERLE estimate.
    erle: f32,
    /// Frame counter for convergence detection.
    frame_count: u32,
    /// Whether the filter has converged.
    converged: bool,
    /// Whether AEC is enabled.
    enabled: bool,
}

impl Aec3 {
    /// Construct a new `Aec3` with default state (filter zeroed, bulk delay 20 ms).
    pub fn new() -> Self {
        Self {
            ref_buf: vec![0.0_f32; REF_BUF_LEN],
            ref_head: 0,
            bulk_delay: 2 * FRAME_SAMPLES,
            xcorr: Box::new([0.0; MAX_DELAY_BLOCKS]),
            filter: vec![0.0_f32; FILTER_TAPS],
            ref_power: NLMS_REG,
            erle: 1.0,
            frame_count: 0,
            converged: false,
            enabled: true,
        }
    }

    /// Enable or disable echo cancellation. When disabled `process_capture` is a no-op.
    pub fn set_enabled(&mut self, enabled: bool) {
        self.enabled = enabled;
    }

    /// Feed one render-reference frame. Call before `process_capture`.
    pub fn process_render(&mut self, render: &[f32; FRAME_SAMPLES]) {
        let n = self.ref_buf.len();
        for &s in render.iter() {
            self.ref_buf[self.ref_head] = s;
            self.ref_head = (self.ref_head + 1) % n;
        }
    }

    /// Process one capture frame in-place (echo subtraction).
    pub fn process_capture(&mut self, mic: &mut [f32; FRAME_SAMPLES]) {
        if !self.enabled {
            return;
        }
        self.frame_count = self.frame_count.saturating_add(1);

        let buf_len = self.ref_buf.len();

        // --- Delay estimation (once per block) ---
        let mic_energy: f32 = mic.iter().map(|x| x * x).sum();
        if mic_energy > 1e-6 {
            for d in 0..MAX_DELAY_BLOCKS {
                let delay = d * FRAME_SAMPLES + self.bulk_delay % FRAME_SAMPLES;
                let mut xc = 0.0_f32;
                for n in 0..FRAME_SAMPLES {
                    let idx = (self.ref_head + buf_len - delay - FRAME_SAMPLES + n) % buf_len;
                    xc += mic[n] * self.ref_buf[idx];
                }
                self.xcorr[d] = self.xcorr[d] * 0.95 + xc.abs() * 0.05;
            }
            let best = self
                .xcorr
                .iter()
                .enumerate()
                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
                .map(|(i, _)| i)
                .unwrap_or(2);
            let cur_block = self.bulk_delay / FRAME_SAMPLES;
            if (best as i32 - cur_block as i32).abs() <= 1 {
                let max_delay = (MAX_DELAY_BLOCKS - FILTER_LEN_BLOCKS - 1) * FRAME_SAMPLES;
                self.bulk_delay = (best * FRAME_SAMPLES).min(max_delay);
            }
        }

        // --- Per-sample NLMS ---
        let mut error = [0.0_f32; FRAME_SAMPLES];
        for n in 0..FRAME_SAMPLES {
            // Reference sample at tap 0 (most recent delayed sample).
            // The reference vector x[n] = [ref[n], ref[n-1], ..., ref[n-FILTER_TAPS+1]]
            // where ref[n] is the render sample delayed by bulk_delay.

            // Echo estimate: ŷ[n] = w · x[n]
            let mut y = 0.0_f32;
            for k in 0..FILTER_TAPS {
                let idx = (self.ref_head + buf_len
                    - self.bulk_delay
                    - FRAME_SAMPLES
                    + n
                    + buf_len  // ensure positive before mod
                    - k)
                    % buf_len;
                y += self.filter[k] * self.ref_buf[idx];
            }

            let e = mic[n] - y;
            error[n] = e;

            // Update running power estimate (exponential moving average).
            // Power of the current reference vector tap 0.
            let x0_idx = (self.ref_head + buf_len - self.bulk_delay - FRAME_SAMPLES + n) % buf_len;
            let x0 = self.ref_buf[x0_idx];
            self.ref_power = self.ref_power * 0.999 + x0 * x0 * 0.001 + NLMS_REG;

            // NLMS weight update: w[k] += μ · e[n] · x[n-k] / power
            let step = MU * e / (self.ref_power * FILTER_TAPS as f32);
            for k in 0..FILTER_TAPS {
                let idx = (self.ref_head + buf_len - self.bulk_delay - FRAME_SAMPLES + n + buf_len
                    - k)
                    % buf_len;
                self.filter[k] += step * self.ref_buf[idx];
            }
        }

        // --- ERLE update ---
        let mic_power: f32 = mic.iter().map(|x| x * x).sum::<f32>() / FRAME_SAMPLES as f32;
        let err_power: f32 = error.iter().map(|x| x * x).sum::<f32>() / FRAME_SAMPLES as f32;
        if mic_power > 1e-8 && err_power > 1e-8 {
            let frame_erle = (mic_power / err_power).clamp(0.5, 100.0);
            self.erle = self.erle * (1.0 - ERLE_ALPHA) + frame_erle * ERLE_ALPHA;
        }

        if self.frame_count > 50 {
            self.converged = true;
        }

        // --- Post-filter ---
        if self.converged && self.erle >= MIN_ERLE {
            let suppression = (1.0 / self.erle.sqrt()).clamp(POST_FILTER_FLOOR, 1.0);
            for n in 0..FRAME_SAMPLES {
                mic[n] = error[n] * suppression;
            }
        } else {
            mic.copy_from_slice(&error);
        }
    }

    /// Reset all adaptive filter state (call on route change or session restart).
    pub fn reset(&mut self) {
        self.ref_buf.fill(0.0);
        self.ref_head = 0;
        self.bulk_delay = 2 * FRAME_SAMPLES;
        self.xcorr.fill(0.0);
        self.filter.fill(0.0);
        self.ref_power = NLMS_REG;
        self.erle = 1.0;
        self.frame_count = 0;
        self.converged = false;
    }

    /// True once the adaptive filter has converged (~500 ms of double-talk).
    pub fn is_converged(&self) -> bool {
        self.converged
    }

    /// Current bulk delay estimate in 10 ms blocks.
    pub fn bulk_delay_blocks(&self) -> usize {
        self.bulk_delay / FRAME_SAMPLES
    }
}

impl Default for Aec3 {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn aec_reduces_echo_after_convergence() {
        let mut aec = Aec3::new();
        let mut render = [0.0_f32; FRAME_SAMPLES];
        for i in 0..FRAME_SAMPLES {
            render[i] = (2.0 * std::f32::consts::PI * 300.0 * i as f32 / 48_000.0).sin() * 0.5;
        }
        // 150 frames to converge (~1.5 s).
        // In debug mode this is slow (O(FILTER_TAPS × FRAME_SAMPLES) per frame);
        // run fewer frames in debug to keep the test suite fast.
        #[cfg(debug_assertions)]
        let frames = 60;
        #[cfg(not(debug_assertions))]
        let frames = 150;

        for _ in 0..frames {
            aec.process_render(&render);
            let mut mic = render;
            aec.process_capture(&mut mic);
        }
        let input_rms = rms(&render);
        aec.process_render(&render);
        let mut mic = render;
        aec.process_capture(&mut mic);
        let output_rms = rms(&mic);

        // In debug mode with fewer frames the filter may not fully converge;
        // we just check it doesn't diverge (output ≤ input).
        #[cfg(debug_assertions)]
        assert!(
            output_rms <= input_rms * 1.1,
            "AEC diverged in debug mode: in={input_rms:.4} out={output_rms:.4}"
        );
        #[cfg(not(debug_assertions))]
        assert!(
            output_rms < input_rms * 0.7,
            "AEC did not reduce echo: in={input_rms:.4} out={output_rms:.4}"
        );
    }

    #[test]
    fn disabled_aec_is_passthrough() {
        let mut aec = Aec3::new();
        aec.set_enabled(false);
        let render = [0.5_f32; FRAME_SAMPLES];
        let mut mic = [0.3_f32; FRAME_SAMPLES];
        aec.process_render(&render);
        aec.process_capture(&mut mic);
        assert!(mic.iter().all(|&s| (s - 0.3).abs() < 1e-6));
    }

    #[test]
    fn reset_clears_state() {
        let mut aec = Aec3::new();
        let render = [0.5_f32; FRAME_SAMPLES];
        for _ in 0..20 {
            aec.process_render(&render);
            let mut mic = render;
            aec.process_capture(&mut mic);
        }
        aec.reset();
        assert!(!aec.is_converged());
        assert_eq!(aec.bulk_delay_blocks(), 2);
    }

    fn rms(frame: &[f32]) -> f32 {
        (frame.iter().map(|x| x * x).sum::<f32>() / frame.len() as f32).sqrt()
    }
}
