import Foundation
import LiveKitWebRTC
import RNNoise
import os

#if canImport(DeepFilter)
  import DeepFilter
#endif

// Noise suppression for the microphone, run inside WebRTC's capture
// processing.
//
// WebRTC's audio processing runs high-pass filter → echo canceller →
// noise suppressor → gain control, then hands each 10 ms frame to a
// "capture post-processing" delegate before encoding. Its own noise
// suppressor is the classic statistical one: good on steady hiss, poor on
// wind, typing and a room's echo, which are what people actually hear.
// So WebRTC's is switched off, and a neural suppressor runs in the
// delegate instead, after the echo canceller has taken out the far side.
//
// The delegate is skipped by WebRTC while the microphone is muted (its
// output is unused), and runs on WebRTC's audio thread: nothing here may
// block, allocate per frame, or wait for the main actor.

/// A mono, frame-at-a-time neural noise suppressor.
///
/// `process` works in place on one frame of `frameLength` samples in
/// WebRTC's float scale — ±32768 for full scale, int16's range as floats —
/// and returns how sure it is that the frame is speech, 0 to 1. Each
/// implementation owns its own conversion to whatever scale its model
/// wants.
public protocol NoiseSuppressor: AnyObject {
  /// Shown in the audio readout.
  var name: String { get }
  /// Samples per frame at 48 kHz, the only rate either model runs at.
  var frameLength: Int { get }
  func process(_ frame: UnsafeMutablePointer<Float>) -> Float
}

/// How far either suppressor may take noise down, in dB. Unlimited, a
/// room goes to digital silence between words, which a listener hears as
/// the line dropping, and a model can hollow out a voice with the noise
/// around it. 30 dB takes a room to near silence and leaves speech its own
/// texture.
public let noiseAttenuationLimitDb: Float = 30

/// xiph's RNNoise: a small recurrent network, BSD-3, 480-sample frames at
/// 48 kHz, and — unusually — in WebRTC's own ±32768 scale already, so the
/// samples go in and come out untouched by any conversion.
///
/// RNNoise has no attenuation limit of its own, so one is made the usual
/// way: the untouched input is mixed back in at the limit.
public final class RNNoiseSuppressor: NoiseSuppressor {
  public let name = "RNNoise"
  public let frameLength: Int
  private let state: OpaquePointer
  private let scratch: UnsafeMutablePointer<Float>
  private let floor = pow(10, -noiseAttenuationLimitDb / 20)

  public init() {
    state = rnnoise_create(nil)
    frameLength = Int(rnnoise_get_frame_size())
    scratch = .allocate(capacity: frameLength)
  }
  deinit {
    rnnoise_destroy(state)
    scratch.deallocate()
  }
  public func process(_ frame: UnsafeMutablePointer<Float>) -> Float {
    scratch.update(from: frame, count: frameLength)
    let speech = rnnoise_process_frame(state, frame, scratch)
    for i in 0..<frameLength { frame[i] += floor * scratch[i] }
    return speech
  }
}

#if canImport(DeepFilter)
  /// DeepFilterNet 3 (Rikorose/DeepFilterNet's libDF, MIT or Apache-2.0):
  /// a larger model than RNNoise, markedly better on non-stationary noise
  /// (typing, wind, voices in the next room). 480-sample hop at 48 kHz,
  /// mono, in ±1 floats — so samples are scaled by 1/32768 on the way in
  /// and back on the way out. It adds about 30 ms of latency (two frames
  /// of lookahead and the overlap of its transform).
  public final class DeepFilterSuppressor: NoiseSuppressor {
    public let name = "DeepFilterNet 3"
    public let frameLength: Int
    private let state: OpaquePointer
    private let input: UnsafeMutablePointer<Float>
    private let output: UnsafeMutablePointer<Float>

    /// The model, where the app carries one: `CALLS_DEEPFILTER_MODEL`
    /// first (tests and the interop harness), then the bundle.
    public static var model: URL? {
      if let path = ProcessInfo.processInfo.environment["CALLS_DEEPFILTER_MODEL"] {
        return URL(fileURLWithPath: path)
      }
      return Bundle.main.url(forResource: "DeepFilterNet3_onnx", withExtension: "tar.gz")
    }

    /// Nil when there is no model file. libDF aborts the process on a
    /// model it cannot read, so the file is checked here first.
    public init?(model: URL, attenuationLimit: Float = noiseAttenuationLimitDb) {
      guard FileManager.default.isReadableFile(atPath: model.path) else { return nil }
      guard let state = df_create(model.path, attenuationLimit, nil) else { return nil }
      self.state = state
      frameLength = Int(df_get_frame_length(state))
      input = .allocate(capacity: frameLength)
      output = .allocate(capacity: frameLength)
    }
    deinit {
      df_free(state)
      input.deallocate()
      output.deallocate()
    }

    /// Returns the model's local SNR estimate mapped to 0...1: at or below
    /// 5 dB reads as noise, at or above 15 dB as speech.
    public func process(_ frame: UnsafeMutablePointer<Float>) -> Float {
      for i in 0..<frameLength { input[i] = frame[i] / 32768 }
      let snr = df_process_frame(state, input, output)
      for i in 0..<frameLength { frame[i] = output[i] * 32768 }
      return min(max((snr - 5) / 10, 0), 1)
    }
  }
#endif

/// A gentle gain after the suppressor, instead of WebRTC's gain control.
///
/// WebRTC's gain control runs *before* the capture post-processing hook,
/// so left on it would raise the room along with a quiet voice and then
/// hand the suppressor louder noise to remove — the pumping a denoised
/// call is known for. It is switched off while a suppressor runs, and this
/// takes its place, after the noise is gone.
///
/// Conservative on purpose: it only ever lifts (0 to +12 dB), it only
/// learns from frames the suppressor calls speech, it rises slowly
/// (about 1 dB a second of speech) and falls fast, and a per-frame peak
/// limit keeps it from ever clipping. A voice already at a sensible level
/// is left alone.
public struct SpeechGain {
  /// Speech level aimed for, dBFS RMS. -26 is the usual target for
  /// conversational speech (ITU-T P.56 practice).
  public static let targetDb: Float = -26
  public static let maxGainDb: Float = 12
  public static let voiced: Float = 0.6
  public private(set) var gainDb: Float = 0

  public init() {}

  public mutating func apply(_ frame: UnsafeMutablePointer<Float>, count: Int, speech: Float) {
    var energy: Float = 0
    var peak: Float = 0
    for i in 0..<count {
      energy += frame[i] * frame[i]
      peak = max(peak, abs(frame[i]))
    }
    let rms = (energy / Float(max(count, 1))).squareRoot() / 32768
    if speech >= SpeechGain.voiced, rms > 0.001 {  // above -60 dBFS
      let levelDb = 20 * log10(rms)
      let wanted = min(max(SpeechGain.targetDb - levelDb, 0), SpeechGain.maxGainDb)
      // Up by 0.01 dB a frame (1 dB/s), down by up to 1 dB a frame.
      gainDb += min(max(wanted - gainDb, -1), 0.01)
    }
    guard gainDb > 0.01 else { return }
    var gain = pow(10, gainDb / 20)
    let limit: Float = 0.9 * 32768
    if peak * gain > limit { gain = max(1, limit / peak) }
    for i in 0..<count { frame[i] *= gain }
  }
}

/// What the suppressor is doing, for the audio readout.
public enum SuppressionState: Equatable, Sendable {
  /// Toggle off, or nothing has asked yet.
  case off
  /// Processing frames; `cpu` is the share of real time spent in it on
  /// the audio thread, 0 to 1 — wall clock, so preemption counts too.
  case running(engine: String, cpu: Double)
  /// Initialised and waiting — muted (WebRTC skips processing then), or
  /// no audio has flowed yet.
  case idle(engine: String)
  /// The capture runs at a rate the models do not (a Bluetooth headset's
  /// 16 kHz, say): frames pass through untouched.
  case unsupportedRate(Int)
  /// Apple's Voice Isolation is on; two suppressors in a row would only
  /// take more of the voice with them.
  case standingAside
}

/// The frame adapter between WebRTC's capture hook and a suppressor.
///
/// WebRTC hands over 10 ms at a time, 480 samples at 48 kHz, which is one
/// frame of either model; a buffer that is a whole number of frames is
/// processed frame by frame, and anything else passes through untouched.
/// So does every rate but 48 kHz, and every channel count but one.
///
/// The suppressor is made once and kept across WebRTC's
/// initialise/release cycles, which happen whenever the processing is
/// reconfigured: a model can take a moment to load, and the audio thread
/// is no place to do it. `release` stops processing; it does not free.
public final class CaptureDenoiser: NSObject, RTCAudioCustomProcessingDelegate, @unchecked Sendable {
  public static let supportedRate = 48_000

  private struct State {
    var enabled = true
    var standAside = false
    var rate = 0
    var channels = 0
    var initialised = false
    var frames: UInt64 = 0
    /// Every buffer WebRTC has handed over, processed or not: proof that
    /// the microphone is capturing.
    var buffers: UInt64 = 0
    /// Buffers since the last initialise; the first is logged.
    var buffersSinceInit: UInt64 = 0
    var busySeconds: Double = 0
    var lastFrame: Double = -.infinity
    var cpu: Double = 0
  }
  private let lock = OSAllocatedUnfairLock(initialState: State())
  /// Read under the lock once per buffer, so it can be replaced mid-call.
  private var current: NoiseSuppressor
  private var gain = SpeechGain()
  private let clock: () -> Double
  /// Called (on the audio thread) when WebRTC reinitialises at a rate the
  /// suppressor can or cannot take, so the processing options can follow.
  public var rateChanged: (@Sendable (Int) -> Void)?

  public init(suppressor: NoiseSuppressor, clock: @escaping () -> Double = CaptureDenoiser.now) {
    current = suppressor
    self.clock = clock
  }

  /// Swaps the model, between buffers. A call starts on RNNoise and moves
  /// to DeepFilterNet once its model has loaded off the main thread
  /// (a third of a second on an M1 Max; the audio thread is no place for
  /// it, and neither is joining).
  public func replace(with suppressor: NoiseSuppressor) {
    lock.withLock { _ in current = suppressor }
  }

  private var suppressor: NoiseSuppressor { lock.withLock { _ in current } }

  public static func now() -> Double { ProcessInfo.processInfo.systemUptime }

  public var engineName: String { suppressor.name }

  /// The toggle. Off passes frames through untouched.
  public var enabled: Bool {
    get { lock.withLock { $0.enabled } }
    set { lock.withLock { $0.enabled = newValue } }
  }
  /// Set while Apple's Voice Isolation is doing the job.
  public var standAside: Bool {
    get { lock.withLock { $0.standAside } }
    set { lock.withLock { $0.standAside = newValue } }
  }
  public var framesProcessed: UInt64 { lock.withLock { $0.frames } }

  public var state: SuppressionState {
    let now = clock()
    return lock.withLock { s in
      if !s.enabled { return .off }
      if s.standAside { return .standingAside }
      if s.initialised, s.rate != CaptureDenoiser.supportedRate { return .unsupportedRate(s.rate) }
      if s.initialised, now - s.lastFrame < 1 { return .running(engine: current.name, cpu: s.cpu) }
      return .idle(engine: current.name)
    }
  }

  // MARK: RTCAudioCustomProcessingDelegate

  public func audioProcessingInitialize(sampleRate sampleRateHz: Int, channels: Int) {
    lock.withLock { s in
      s.rate = sampleRateHz
      s.channels = channels
      s.initialised = true
      s.buffersSinceInit = 0
    }
    AudioLog.info(
      "capture hook initialised at \(sampleRateHz) Hz, \(channels) ch; noise reduction \(sampleRateHz == CaptureDenoiser.supportedRate ? "can run" : "passes it through (meters still read it)")")
    // Every time, not only on a change. WebRTC initialises its processing
    // at 16 kHz when the module is made, long before any microphone opens,
    // and that first call can arrive before anyone is listening; a
    // Bluetooth headset that really does capture at 16 kHz then looked
    // like "no change", and the plan never heard about it — leaving
    // WebRTC's own suppressor off while the readout said it was standing
    // in. The listener decides what is a change.
    rateChanged?(sampleRateHz)
  }

  public var buffersSeen: UInt64 { lock.withLock { $0.buffers } }

  /// Counts one buffer; true for the first since the last initialise.
  /// Apart from WebRTC's buffer type so the tests can drive it.
  @discardableResult
  public func countBuffer() -> Bool {
    lock.withLock { s in
      s.buffers += 1
      s.buffersSinceInit += 1
      return s.buffersSinceInit == 1
    }
  }

  /// The rate and channels WebRTC last initialised the hook with (0
  /// before it has): for the diagnostics.
  public var format: (rate: Int, channels: Int, initialised: Bool) {
    lock.withLock { s in (s.rate, s.channels, s.initialised) }
  }

  public func audioProcessingProcess(audioBuffer: RTCAudioBuffer) {
    if countBuffer() {
      AudioLog.info(
        "capture hook: first buffer since initialise, \(audioBuffer.frames) frames x \(audioBuffer.channels) ch")
    }
    guard audioBuffer.channels >= 1 else { return }
    handle(audioBuffer.rawBuffer(forChannel: 0), frames: audioBuffer.frames, channels: audioBuffer.channels)
  }

  /// Everything the hook does with one buffer, apart from WebRTC's buffer
  /// type so the tests drive this exact path: the meter before, the
  /// suppressor (mono at 48 kHz only), then the meter after and the sink —
  /// at every rate and channel count, so a Bluetooth headset's 16 or 24
  /// kHz, which the suppressor passes through, is metered and recorded
  /// all the same.
  public func handle(_ samples: UnsafeMutablePointer<Float>, frames: Int, channels: Int) {
    inputLevel.record(samples, frames: frames, at: clock())
    if channels == 1 { process(samples, frames: frames) }
    observe(samples, frames: frames)
  }

  /// The same buffers measured on the way *in*, before the suppressor:
  /// what the microphone delivers after WebRTC's echo canceller and
  /// high-pass filter. "Is the microphone picking anything up" is asked
  /// of this one, because a quiet room the suppressor has taken down by
  /// 30 dB can look dead afterwards and is not.
  public let inputLevel = CaptureLevel()

  /// What leaves the hook, measured: the level of every buffer after the
  /// suppressor and gain have had it (or passed it through), which is the
  /// audio that is encoded and sent. Read by the microphone meter.
  public let level = CaptureLevel()

  /// Something that wants the processed audio itself — the audio test's
  /// recording. Called on the audio thread with WebRTC's ±32768 floats and
  /// the rate WebRTC is processing at; it must copy and return, never
  /// block. Nil almost always.
  public var sink: (@Sendable (UnsafePointer<Float>, Int, Int) -> Void)? {
    get { lock.withLock { _ in sinkBox } }
    set { lock.withLock { _ in sinkBox = newValue } }
  }
  private var sinkBox: (@Sendable (UnsafePointer<Float>, Int, Int) -> Void)?

  /// The measuring half of the hook, apart from WebRTC's buffer type so
  /// the tests can drive it: after `process`, whatever the channel count
  /// or rate, so a stream the suppressor passes through is metered too.
  public func observe(_ samples: UnsafePointer<Float>, frames: Int) {
    guard frames > 0 else { return }
    level.record(samples, frames: frames, at: clock())
    let (sink, rate) = lock.withLock { s in (sinkBox, s.rate) }
    sink?(samples, frames, rate)
  }

  public func audioProcessingRelease() {
    lock.withLock { $0.initialised = false }
  }

  // MARK: The adapter itself, apart from WebRTC's buffer type

  /// Processes one buffer in place. Public for the tests, which cannot
  /// make an RTCAudioBuffer.
  public func process(_ samples: UnsafeMutablePointer<Float>, frames: Int) {
    let (run, suppressor) = lock.withLock { s in
      (
        s.enabled && !s.standAside && s.initialised && s.channels == 1
          && s.rate == CaptureDenoiser.supportedRate, current
      )
    }
    let length = suppressor.frameLength
    guard run, frames > 0, frames % length == 0 else { return }
    let started = clock()
    var offset = 0
    while offset < frames {
      let frame = samples + offset
      let speech = suppressor.process(frame)
      gain.apply(frame, count: length, speech: speech)
      offset += length
    }
    let finished = clock()
    let audio = Double(frames) / Double(CaptureDenoiser.supportedRate)
    lock.withLock { s in
      s.frames += UInt64(frames / length)
      s.lastFrame = finished
      // A smoothed share of one core: time spent against audio covered.
      let share = (finished - started) / audio
      s.cpu = s.frames <= UInt64(frames / length) ? share : s.cpu * 0.98 + share * 0.02
    }
  }
}

/// The level of the processed microphone, kept for a meter to read.
///
/// Written on the audio thread once per buffer (10 ms), read on the main
/// actor a few times a second. What a read gets is the loudest buffer
/// since the last read — a peak hold — so a meter sampled at 10 Hz does
/// not miss a syllable that fell between two samples. Levels are 0...1 of
/// full scale (WebRTC's ±32768).
public final class CaptureLevel: @unchecked Sendable {
  public struct Reading: Equatable, Sendable {
    /// Loudest buffer RMS since the last read, 0...1.
    public var rms: Float
    /// Loudest single sample since the last read, 0...1.
    public var peak: Float
    /// Buffers measured since the last read. Zero means nothing flowed:
    /// muted (WebRTC skips the hook then) or not capturing.
    public var buffers: Int
    /// When the last buffer was measured, on the clock the adapter was
    /// given (uptime seconds by default).
    public var at: Double
  }
  private let lock = OSAllocatedUnfairLock(initialState: Reading(rms: 0, peak: 0, buffers: 0, at: -.infinity))

  public init() {}

  public func record(_ samples: UnsafePointer<Float>, frames: Int, at time: Double) {
    var energy: Float = 0
    var peak: Float = 0
    for i in 0..<frames {
      let v = samples[i]
      energy += v * v
      peak = max(peak, abs(v))
    }
    let rms = (energy / Float(frames)).squareRoot() / 32768
    let top = min(peak / 32768, 1)
    lock.withLock { r in
      r.rms = max(r.rms, min(rms, 1))
      r.peak = max(r.peak, top)
      r.buffers += 1
      r.at = time
    }
  }

  /// The reading since the last call, and a fresh start for the next.
  public func take() -> Reading {
    lock.withLock { r in
      let out = r
      r.rms = 0
      r.peak = 0
      r.buffers = 0
      return out
    }
  }
}