import Foundation

#if canImport(JLAudioUnitKit)
import JLAudioUnitKit
#endif

/// Streaming OPUS decoder + encoder built on top of `JLAudioUnitKit`.
///
/// Mirrors the Android `audio/OpusStreamCodec.kt` (decoder side). The codec
/// pushes 16 kHz / 16-bit / mono / 20 ms PCM frames through the supplied
/// callback, matching the Android contract used by `RcspTranslationRuntime`.
///
/// Important: `JLOpusDecoder.opusDecoderInputData(_:)` accepts **one OPUS
/// packet at a time**. The headset advertises packets at 40 bytes for 16k/20ms;
/// pass the byte stream directly — the decoder does not buffer. If the
/// upstream feed combines multiple packets in one chunk, we slice along
/// `packetSize` to keep the decoder happy.
final class OpusStreamDecoder: NSObject {

  let sampleRate: Int
  let channels: Int
  let packetSize: Int

  private let onPcm: (Data) -> Void
  private let onError: (Int, String?) -> Void

  #if canImport(JLAudioUnitKit)
  private var decoder: JLOpusDecoder?
  #endif
  private var started = false
  private let lock = NSLock()

  init(
    channels: Int = 1,
    packetSize: Int = 40,
    sampleRate: Int = 16000,
    onPcm: @escaping (Data) -> Void,
    onError: @escaping (Int, String?) -> Void = { _, _ in }
  ) {
    self.channels = channels
    self.packetSize = packetSize
    self.sampleRate = sampleRate
    self.onPcm = onPcm
    self.onError = onError
  }

  func start() {
    lock.lock(); defer { lock.unlock() }
    if started { return }
    #if canImport(JLAudioUnitKit)
    let format = JLOpusFormat.defaultFormats()
    format.sampleRate = Int32(sampleRate)
    format.channels = Int32(channels)
    format.frameDuration = 20
    // The headset's stream is the JieLi-flavored OPUS (no header). The
    // demo's `defaultFormats` already disables the header for the headset
    // direction, so we keep that default.
    format.dataSize = Int32(packetSize)
    decoder = JLOpusDecoder(decoder: format, delegate: self)
    #endif
    started = true
  }

  func stop() {
    lock.lock(); defer { lock.unlock() }
    if !started { return }
    started = false
    #if canImport(JLAudioUnitKit)
    decoder?.opusOnRelease()
    decoder = nil
    #endif
  }

  /// Feed one or more OPUS packets concatenated. The decoder is fed one
  /// `packetSize`-sized chunk at a time when the input length is a multiple
  /// of `packetSize`; otherwise the entire buffer is forwarded as-is and
  /// the decoder does its own resync.
  func feedEncoded(_ data: Data) {
    #if canImport(JLAudioUnitKit)
    guard started, let dec = decoder else { return }
    if packetSize > 0 && data.count > packetSize && data.count % packetSize == 0 {
      var offset = 0
      while offset < data.count {
        let end = min(offset + packetSize, data.count)
        dec.opusDecoderInputData(data.subdata(in: offset..<end))
        offset = end
      }
    } else {
      dec.opusDecoderInputData(data)
    }
    #endif
  }
}

#if canImport(JLAudioUnitKit)
extension OpusStreamDecoder: JLOpusDecoderDelegate {
  func opusDecoder(_ decoder: JLOpusDecoder, data: Data?, error: Error?) {
    if let err = error {
      onError((err as NSError).code, err.localizedDescription)
      return
    }
    guard let pcm = data, !pcm.isEmpty else { return }
    onPcm(pcm)
  }
}
#endif

/// Streaming OPUS encoder used to repackage TTS PCM for the headset's
/// downlink path. Buffers PCM chunks and emits encoded packets through the
/// callback.
final class OpusStreamEncoder: NSObject {

  let sampleRate: Int
  let channels: Int
  private let onEncoded: (Data) -> Void
  private let onError: (Int, String?) -> Void

  #if canImport(JLAudioUnitKit)
  private var encoder: JLOpusEncoder?
  #endif
  private var started = false
  private let lock = NSLock()

  init(
    channels: Int = 1,
    sampleRate: Int = 16000,
    onEncoded: @escaping (Data) -> Void,
    onError: @escaping (Int, String?) -> Void = { _, _ in }
  ) {
    self.sampleRate = sampleRate
    self.channels = channels
    self.onEncoded = onEncoded
    self.onError = onError
  }

  func start() {
    lock.lock(); defer { lock.unlock() }
    if started { return }
    #if canImport(JLAudioUnitKit)
    // Use the JieLi default config (no OGG header — same as the headset
    // expects, matches the demo's `JLOpusEncodeConfig.default()`).
    let cfg = JLOpusEncodeConfig.default()
    cfg.sampleRate = Int32(sampleRate)
    cfg.channels = Int32(channels)
    cfg.frameDuration = 20
    encoder = JLOpusEncoder(format: cfg, delegate: self)
    #endif
    started = true
  }

  func stop() {
    lock.lock(); defer { lock.unlock() }
    if !started { return }
    started = false
    #if canImport(JLAudioUnitKit)
    encoder?.opusOnRelease()
    encoder = nil
    #endif
  }

  /// Feed one or more raw PCM bytes (16-bit LE, `channels` interleaved when >1).
  func feedPcm(_ pcm: Data) {
    #if canImport(JLAudioUnitKit)
    guard started, let enc = encoder else { return }
    enc.opusEncode(pcm)
    #endif
  }
}

#if canImport(JLAudioUnitKit)
extension OpusStreamEncoder: JLOpusEncoderDelegate {
  func opusEncoder(_ encoder: JLOpusEncoder, data: Data?, error: Error?) {
    if let err = error {
      onError((err as NSError).code, err.localizedDescription)
      return
    }
    guard let opus = data, !opus.isEmpty else { return }
    onEncoded(opus)
  }
}
#endif