#include "resampler.h"
#include "buffer.h"
#include "remotebuffer.h"
#include <catch2/catch.hpp>
#include <emscripten/bind.h>
#include <speex/speex_resampler.h>

Resampler::Resampler(
    std::shared_ptr<RandomAccessAudioReadable> underlyingReadable,
    int numerator, int denominator)
    : state(State::Reset), currentFrameOffset(0),
      currentUnderlyingFrameOffset(0), numerator(numerator),
      denominator(denominator), excessOutputOffset(0), underlyingChunkOffset(0),
      excessOutputLength(0) {
  assert(underlyingReadable != nullptr);
  const auto underlyingSampleRate = underlyingReadable->getSampleRate();
  const auto underlyingChannelCount = underlyingReadable->getChannelCount();
  underlyingChunk = BufferF32(underlyingChannelCount, 256);
  excessOutput = BufferF32(underlyingChannelCount, 256);
  this->underlyingReadable = std::move(underlyingReadable);
  finishInit(underlyingChannelCount);
}

Resampler::Resampler(int channelCount, int numerator, int denominator)
    : state(State::Reset), currentFrameOffset(0),
      currentUnderlyingFrameOffset(0), numerator(numerator),
      denominator(denominator), excessOutputOffset(0), underlyingChunkOffset(0),
      excessOutputLength(0) {
  underlyingChunk = BufferF32(channelCount, 256);
  excessOutput = BufferF32(channelCount, 256);
  int err = 0;
  finishInit(channelCount);
}

Resampler::~Resampler() { speex_resampler_destroy(resamplerState); }

void Resampler::finishInit(int underlyingChannelCount) {
  assert(numerator > 0);
  assert(denominator > 0);
  const auto underlyingSampleRate = (underlyingReadable != nullptr)
                                        ? underlyingReadable->getSampleRate()
                                        : 48000;
  int err = 0;
  outputRate = (int)std::round((double)underlyingSampleRate *
                               (double)denominator / (double)numerator);
  resamplerState =
      speex_resampler_init_frac(underlyingChannelCount, numerator, denominator,
                                underlyingSampleRate, outputRate, 9, &err);

  assert(err == 0);
  assert(resamplerState != nullptr);
}

void Resampler::setUnderlying(
    std::shared_ptr<RandomAccessAudioReadable> underlyingReadable,
    int numerator, int denominator) {
  assert(underlyingReadable == nullptr ||
         underlyingReadable->getChannelCount() ==
             underlyingChunk.getChannelCount());
  this->underlyingReadable = std::move(underlyingReadable);
  this->numerator = numerator;
  this->denominator = denominator;
  setSampleRateRatio(numerator, denominator);
}

int Resampler::getSampleRate() const { return outputRate; }

int Resampler::getChannelCount() const {
  return underlyingChunk.getChannelCount();
}

int Resampler::getFrameCount() const {
  if (underlyingReadable != nullptr) {
    return (int)std::ceil((double)underlyingReadable->getFrameCount() *
                          (double)numerator / (double)denominator);
  } else {
    return 0;
  }
}

int Resampler::getNumerator() const { return numerator; }

int Resampler::getDenominator() const { return denominator; }

void Resampler::setSampleRateRatio(int newNumerator, int newDenominator) {
  assert(newNumerator > 0);
  assert(newDenominator > 0);
  numerator = newNumerator;
  denominator = newDenominator;
  if (underlyingReadable == nullptr) {
    return;
  }
  const auto inputRate = underlyingReadable->getSampleRate();
  outputRate = (int)std::round((double)inputRate * (double)denominator /
                               (double)numerator);
  speex_resampler_set_rate_frac(resamplerState, numerator, denominator,
                                inputRate, outputRate);
  state = State::Reset;
}

void Resampler::read(int frameOffset, std::shared_ptr<BufferF32> output) {
  read(frameOffset, output->slice(0, output->getFrameCount()));
}

void Resampler::read(int frameOffset, BufferF32 output) {
  if (underlyingReadable == nullptr) {
    output.fill(0);
    return;
  }
  assert(output.getChannelCount() == underlyingChunk.getChannelCount());
  if (state == State::Reset || frameOffset != currentFrameOffset) {
    seekTo(frameOffset);
  }
  // now currentFrameOffset == frameOffset and state == ContinuousPlayback
  int outputOffset = 0;
  while (outputOffset < (int)output.getFrameCount()) {
    if (excessOutputOffset < excessOutputLength) {
      auto framesToCopy = std::min((int)output.getFrameCount() - outputOffset,
                                   excessOutputLength - excessOutputOffset);
      output.set(outputOffset,
                 excessOutput.slice(excessOutputOffset,
                                    excessOutputOffset + framesToCopy));
      excessOutputOffset += framesToCopy;
      outputOffset += framesToCopy;
      continue;
    }

    if (underlyingChunkOffset < underlyingChunk.getFrameCount()) {
      auto underlyingSlice = underlyingChunk.slice(
          underlyingChunkOffset, underlyingChunk.getFrameCount());
      int samplesConsumed = 0;
      int samplesProduced = 0;
      for (int channel = 0; channel < output.getChannelCount(); channel++) {
        uint32_t in_len = underlyingSlice.getFrameCount();
        uint32_t out_len = excessOutput.getFrameCount();
        speex_resampler_process_float(
            resamplerState, channel, underlyingSlice.getChannelData(channel),
            &in_len, excessOutput.getChannelData(channel), &out_len);
        if (channel == 0) {
          samplesConsumed = in_len;
          samplesProduced = out_len;
        }
      }
      underlyingChunkOffset += samplesConsumed;
      excessOutputOffset = 0;
      excessOutputLength = samplesProduced;

      assert(samplesConsumed > 0);

      continue;
    } else {
      underlyingReadable->readZeroPadded(currentUnderlyingFrameOffset,
                                         underlyingChunk);
      underlyingChunkOffset = 0;
      currentUnderlyingFrameOffset += underlyingChunk.getFrameCount();
      continue;
    }
  }
  currentFrameOffset += outputOffset;
}

void Resampler::seekTo(int frameOffset) {
  const auto inputLatency = speex_resampler_get_input_latency(resamplerState);
  currentFrameOffset = frameOffset;
  currentUnderlyingFrameOffset =
      (int)std::round((double)currentFrameOffset * (double)numerator /
                      (double)denominator) -
      inputLatency;
  speex_resampler_skip_zeros(resamplerState);
  state = State::ContinuousPlayback;
  excessOutputOffset = 0;
  excessOutputLength = 0;
  underlyingChunkOffset = underlyingChunk.getFrameCount();
}

EMSCRIPTEN_BINDINGS(resampler) {
  using namespace emscripten;
  class_<Resampler, base<RandomAccessAudioReadable>>("Resampler")
      .smart_ptr_constructor<std::shared_ptr<Resampler>,
                             std::shared_ptr<RandomAccessAudioReadable> &&,
                             int &&, int &&>("Resampler",
                                             &std::make_shared<Resampler>)
      .property("numerator", &Resampler::getNumerator)
      .property("denominator", &Resampler::getDenominator)
      .function("setSampleRateRatio(numerator, denominator)",
                &Resampler::setSampleRateRatio);
}

TEST_CASE("resampler memory tests", "[resampler]") {
  auto inputBuf = std::make_shared<BufferF32>(2, 100000);
  inputBuf->noise();
  auto inputRd = std::make_shared<BufferAsRemoteAudioBuffer>(44100, inputBuf);

  auto outputBuf = std::make_shared<BufferF32>(2, 1024);

  auto resampler = Resampler(inputRd, 12345, 23456);

  REQUIRE(resampler.getChannelCount() == 2);

  for (int i = 0; i < resampler.getFrameCount();
       i += outputBuf->getFrameCount()) {
    resampler.read(i, outputBuf);
  }

  auto resampler2 = Resampler(2, 12345, 23456);
  resampler2.read(0, outputBuf);
  REQUIRE(outputBuf->peak() == 0);
}