"""
wivexx_aec.py — echo-aware barge-in gating + optional true AEC  (v1.17)
═══════════════════════════════════════════════════════════════════════
Stops WIVEXX from interrupting himself: while he's speaking, the barge-in
monitor hears his own TTS through the mic and can self-trigger.  Two layers:

  1. ECHO-AWARE GATING (always available, no extra dependency):
     The playback paths push the TTS output ENVELOPE (loudness over time) into
     a small ring buffer.  The monitor asks "how loud was I ~AEC_LATENCY_MS
     ago?" and raises the barge-in gate accordingly, so WIVEXX's own loud
     passages stop tripping the interrupt — while a genuinely louder human
     voice still gets through.

  2. TRUE AEC (optional, needs `speexdsp`):
     When real time-aligned reference SAMPLES are available (the sounddevice
     PCM playback path) and `AEC_ENABLED` is on, a speex EchoCanceller
     subtracts the reference from the mic frame before the RMS/VAD check.
     Falls back to layer 1 when speexdsp or samples are absent.

Everything is best-effort and wrapped: any failure leaves the original
barge-in behaviour untouched.

Public API (used by wivexx_api.py):
    set_logger(logger) ; configure(latency_ms=, gating=, factor=, enabled=)
    reset()
    push_output(samples, rate)        # feed TTS samples being played
    push_envelope(rms16, dur_s)       # feed a known loudness for a duration
    echo_estimate(now=None) -> float  # int16-rms estimate of current echo
    adjusted_gate(base_gate) -> float # gate raised by the echo estimate
    true_aec_available() -> bool
    cancel(mic_i16, rate) -> np.ndarray  # cleaned frame (or mic unchanged)
"""

import threading
import time
from typing import Optional

import numpy as np

_logger = None
_lock = threading.Lock()

# envelope ring buffer: list of (t_start, t_end, rms16)
_env: list = []
_ENV_KEEP = 6.0                  # seconds of history to retain

# aligned reference samples (only fed by the PCM path): (t_start, rate, np.int16)
_ref_lock = threading.Lock()
_ref_clips: list = []
_REF_KEEP = 6.0

# tunables (overridden from config)
_latency_ms = 180.0
_gating = True
_factor = 0.8                    # how much of the echo estimate adds to the gate
_enabled = False                 # true-AEC master switch

_canceller = None
_canceller_frame = 0


def set_logger(logger) -> None:
    global _logger
    _logger = logger


def configure(latency_ms=None, gating=None, factor=None, enabled=None) -> None:
    global _latency_ms, _gating, _factor, _enabled
    if latency_ms is not None:
        _latency_ms = float(latency_ms)
    if gating is not None:
        _gating = bool(gating)
    if factor is not None:
        _factor = float(factor)
    if enabled is not None:
        _enabled = bool(enabled)


def reset() -> None:
    with _lock:
        _env.clear()
    with _ref_lock:
        _ref_clips.clear()


# ── envelope feed + estimate ─────────────────────────────────────────────────
def _rms16(samples: np.ndarray) -> float:
    a = np.asarray(samples)
    if a.size == 0:
        return 0.0
    if a.dtype.kind in "iu":
        a = a.astype(np.float32)
    else:                                  # float [-1,1] → int16 scale
        a = a.astype(np.float32) * 32768.0
    return float(np.sqrt(np.mean(a * a)))


def push_envelope(rms16: float, dur_s: float = 0.05) -> None:
    """Record that WIVEXX emitted audio at *rms16* loudness for *dur_s* now."""
    now = time.time()
    with _lock:
        _env.append((now, now + max(0.0, dur_s), float(rms16)))
        cut = now - _ENV_KEEP
        while _env and _env[0][1] < cut:
            _env.pop(0)


def push_output(samples, rate: int) -> None:
    """Feed a chunk of TTS samples being played.  Stores the loudness envelope
    (always) and the raw samples (for true AEC, when this is the PCM path)."""
    try:
        a = np.asarray(samples)
        if a.ndim > 1:                     # stereo → mono
            a = a.mean(axis=1)
        dur = len(a) / float(rate or 16000)
        # break long clips into ~50 ms envelope cells for time resolution
        cell = max(1, int((rate or 16000) * 0.05))
        now = time.time()
        with _lock:
            for i in range(0, len(a), cell):
                seg = a[i:i + cell]
                t0 = now + (i / float(rate or 16000))
                _env.append((t0, t0 + len(seg) / float(rate or 16000),
                             _rms16(seg)))
            cut = now - _ENV_KEEP
            while _env and _env[0][1] < cut:
                _env.pop(0)
        with _ref_lock:
            _ref_clips.append((now, int(rate or 16000),
                               a.astype(np.float32)))
            rcut = now - _REF_KEEP
            while _ref_clips and (_ref_clips[0][0]
                                  + len(_ref_clips[0][2]) / _ref_clips[0][1]) < rcut:
                _ref_clips.pop(0)
    except Exception as e:
        if _logger:
            _logger.debug("aec push_output failed: %s", e)


def echo_estimate(now: Optional[float] = None) -> float:
    """int16-RMS of the TTS output that should be hitting the mic right now,
    i.e. what was played ~latency_ms ago."""
    if not _env:
        return 0.0
    t = (now or time.time()) - _latency_ms / 1000.0
    best = 0.0
    with _lock:
        for (t0, t1, rms) in reversed(_env):
            if t0 - 0.06 <= t <= t1 + 0.06:
                best = max(best, rms)
        if best == 0.0 and _env:
            # nothing aligned exactly — fall back to the most recent cell
            best = _env[-1][2] * 0.6
    return best


def adjusted_gate(base_gate: float) -> float:
    """Barge-in gate raised by the current echo estimate (if gating is on)."""
    if not _gating:
        return base_gate
    try:
        return base_gate + _factor * echo_estimate()
    except Exception:
        return base_gate


# ── optional true AEC via speexdsp ───────────────────────────────────────────
def true_aec_available() -> bool:
    if not _enabled:
        return False
    try:
        import speexdsp  # noqa: F401
        return True
    except Exception:
        return False


def _aligned_reference(n: int, rate: int, now: Optional[float] = None):
    """Pull *n* reference samples (mono float) that were played ~latency ago,
    resampled to *rate*.  None when no overlapping clip is buffered."""
    t = (now or time.time()) - _latency_ms / 1000.0
    with _ref_lock:
        for (t0, r, a) in reversed(_ref_clips):
            t1 = t0 + len(a) / r
            if t0 <= t <= t1:
                start = int((t - t0) * r)
                seg = a[start:start + int(n * r / rate)]
                if len(seg) == 0:
                    return None
                if r != rate:
                    idx = (np.arange(n) * (r / rate)).astype(int)
                    idx = idx[idx < len(seg)]
                    seg = seg[idx] if len(idx) else seg
                out = np.zeros(n, dtype=np.float32)
                out[:len(seg)] = seg[:n]
                return out
    return None


def cancel(mic_i16: np.ndarray, rate: int) -> np.ndarray:
    """Echo-cancelled mic frame using speexdsp when available + aligned ref,
    else the mic frame unchanged."""
    global _canceller, _canceller_frame
    if not true_aec_available():
        return mic_i16
    try:
        from speexdsp import EchoCanceller
        n = len(mic_i16)
        ref = _aligned_reference(n, rate)
        if ref is None:
            return mic_i16
        if _canceller is None or _canceller_frame != n:
            _canceller = EchoCanceller.create(n, n * 10, rate)
            _canceller_frame = n
        ref_i16 = np.clip(ref * 32768.0, -32768, 32767).astype(np.int16)
        out = _canceller.process(mic_i16.astype(np.int16).tobytes(),
                                 ref_i16.tobytes())
        return np.frombuffer(out, dtype=np.int16)
    except Exception as e:
        if _logger:
            _logger.debug("true AEC failed, passing mic through: %s", e)
        return mic_i16
