AMRWB.h

#pragma once
 
#include <string.h>
 
#include "AMRCodec.h"
extern "C" {
#include "wb/dec_if.h"
#include "wb/enc_if.h"
}
 
class AMRWB : public AMRCodec {
 public:
  enum class Mode {
    WB_6_60 = 0,  // 6.60 kbit/s
    WB_8_85,      // 8.85 kbit/s
    WB_12_65,     // 12.65 kbit/s
    WB_14_25,     // 14.25 kbit/s
    WB_15_85,     // 15.85 kbit/s
    WB_18_25,     // 18.25 kbit/s
    WB_19_85,     // 19.85 kbit/s
    WB_23_05,     // 23.05 kbit/s
    WB_23_85      // 23.85 kbit/s
  };
 
  AMRWB() = default;
 
  inline ~AMRWB() override {
    if (encoderState) {
      E_IF_exit(encoderState);
      encoderState = nullptr;
    }
 
    if (decoderState) {
      D_IF_exit(decoderState);
      decoderState = nullptr;
    }
  }
 
  inline void setMode(Mode mode) { currentMode = mode; }
 
  inline Mode getMode() const { return currentMode; }
 
  inline int encode(const int16_t* pcmSamples, size_t sampleCount,
                    uint8_t* amrData, size_t amrBufferSize) override {
    if (encoderState == nullptr) encoderState = E_IF_init();
 
    if (encoderState == nullptr || !amrData || amrBufferSize == 0) {
      return 0;
    }
 
    // AMR-WB works with 320 samples per frame (20ms at 16kHz)
    const size_t samplesPerFrame = 320;
    size_t bytesWritten = 0;
 
    for (size_t i = 0; i < sampleCount && bytesWritten < amrBufferSize;
         i += samplesPerFrame) {
      // Make sure we have a full frame
      if (i + samplesPerFrame > sampleCount) {
        break;
      }
 
      // Check if we have enough space for worst-case frame size (64 bytes)
      if (bytesWritten + 64 > amrBufferSize) {
        break;
      }
 
      // Encode the frame
      int frameBytes = E_IF_encode(encoderState, static_cast<int>(currentMode),
                                   const_cast<short*>(pcmSamples + i),
                                   amrData + bytesWritten,
                                   0  // forceSpeech
      );
 
      bytesWritten += frameBytes;
    }
 
    return bytesWritten;
  }
 
  inline int decode(const uint8_t* amrData, size_t amrSize, int16_t* pcmSamples,
                    size_t maxSampleCount) override {
    if (decoderState == nullptr) decoderState = D_IF_init();
 
    if (decoderState == nullptr || !amrData || amrSize == 0 || !pcmSamples ||
        maxSampleCount == 0) {
      return 0;
    }
 
    // AMR-WB produces 320 samples per frame (20ms at 16kHz)
    const size_t samplesPerFrame = getFrameSizeSamples();
    size_t totalSamplesDecoded = 0;
    size_t offset = 0;
 
    // Process the input buffer frame by frame
    while (offset < amrSize &&
           totalSamplesDecoded + samplesPerFrame <= maxSampleCount) {
      // Extract frame type from the frame header
      uint8_t frameType = (amrData[offset] >> 3) & 0x0F;
 
      // Find frame size based on frame type
      size_t frameSize = 1;
      if (frameSize <= 8) {
        frameSize = getEncodedFrameSizeBytes(frameType);
      }
 
      if (offset + frameSize > amrSize) {
        break;  // Not enough data for another frame
      }
 
      // Decode this frame
      D_IF_decode(decoderState, const_cast<unsigned char*>(amrData + offset),
                  pcmSamples + totalSamplesDecoded,
                  0  // BFI (bad frame indicator)
      );
 
      offset += frameSize;
      totalSamplesDecoded += samplesPerFrame;
    }
 
    return totalSamplesDecoded;
  }
 
  inline int getSampleRate() const override { return 16000; }
 
  int getFrameSizeSamples() override { return 320; }
 
  int getEncodedFrameSizeBytes() override {
    return getEncodedFrameSizeBytes(static_cast<int>(currentMode));
  }
 
 private:
  void* encoderState = nullptr;
  void* decoderState = nullptr;
  Mode currentMode = Mode::WB_23_85;
 
  int getEncodedFrameSizeBytes(int mode) {
    // Bytes per encoded frame for each mode
    const uint8_t frameSizes[] = {
        18,  // WB_6_60 (6.60 kbps)
        24,  // WB_8_85 (8.85 kbps)
        33,  // WB_12_65 (12.65 kbps)
        37,  // WB_14_25 (14.25 kbps)
        41,  // WB_15_85 (15.85 kbps)
        47,  // WB_18_25 (18.25 kbps)
        51,  // WB_19_85 (19.85 kbps)
        59,  // WB_23_05 (23.05 kbps)
        61   // WB_23_85 (23.85 kbps)
    };
    return frameSizes[mode];
  }
};