#include "limiter.h"
#include "simd.h"
#include <cassert>
#include <catch2/catch.hpp>
#include <cmath>
#include <emscripten/atomic.h>
#include <emscripten/console.h>
#include <random>
#include <speex/speex_resampler.h>

static float timeToAlpha(float seconds, float sampleRate) {
  float timeSamples = seconds * sampleRate;

  return 1.0f - std::exp(-1.0f / timeSamples);
}

static void slidingWindowMax(const float *input, float *output, int *dq,
                             int inputLength, int windowLength,
                             int outputLength) {
  int dq_head = 0, dq_tail = 0;

  for (int i = 0; i < inputLength; ++i) {
    if (dq_head != dq_tail && dq[dq_head] <= i - windowLength) {
      dq_head++;
    }
    while (dq_head != dq_tail && input[dq[dq_tail - 1]] <= input[i]) {
      dq_tail--;
    }
    dq[dq_tail++] = i;

    int peak_idx = i - windowLength + 1;
    if (peak_idx >= 0 && peak_idx < outputLength) {
      output[peak_idx] = input[dq[dq_head]];
    }
  }
}

Limiter::Limiter(int channelCount, int sampleRate, float lookaheadSeconds,
                 float attackSeconds, float releaseSeconds, float preGainDb,
                 bool bypass, float stereoLink)
    : lookaheadFrames(lookaheadSeconds * (double)sampleRate *
                      (double)upsampleFactor),
      desiredLookaheadSeconds(lookaheadSeconds),
      desiredAttackSeconds(attackSeconds),
      desiredReleaseSeconds(releaseSeconds), preGainDb(preGainDb),
      bypass(bypass), stereoLink(stereoLink), channelCount(channelCount),
      sampleRate(sampleRate), delayFrames(0), upsampler(nullptr),
      downsampler(nullptr), workBuffer(channelCount, workBufferSize),
      attackSeconds(attackSeconds), releaseSeconds(releaseSeconds),
      attackAlpha(timeToAlpha(attackSeconds, sampleRate * upsampleFactor)),
      releaseAlpha(timeToAlpha(releaseSeconds, sampleRate * upsampleFactor)),
      channelCurrentGainEnvelope(channelCount, 1.f),
      linkedCurrentGainEnvelope(1.f),
      limiterWorkBuffer(channelCount, workBufferSize),
      slidingWindowWorkBuffer(0, 0) {
  assert(channelCount == 1 || channelCount == 2);

  int err = 0;
  upsampler =
      speex_resampler_init(channelCount, sampleRate,
                           sampleRate * upsampleFactor, resampleQuality, &err);

  assert(err == 0);
  assert(upsampler != nullptr);

  downsampler = speex_resampler_init(channelCount, sampleRate * upsampleFactor,
                                     sampleRate, resampleQuality, &err);
  assert(err == 0);
  assert(downsampler != nullptr);

  delayFrames = speex_resampler_get_input_latency(upsampler) +
                speex_resampler_get_output_latency(downsampler) +
                lookaheadFrames / upsampleFactor;

  lookaheadDelayLine =
      std::make_unique<DelayLine>(channelCount, lookaheadFrames);
  deque.resize(workBufferSize + lookaheadFrames + 1);
  slidingWindowWorkBuffer =
      BufferF32(channelCount, workBufferSize + lookaheadFrames);
}

Limiter::~Limiter() {
  if (upsampler != nullptr) {
    speex_resampler_destroy(upsampler);
  }
  if (downsampler != nullptr) {
    speex_resampler_destroy(downsampler);
  }
}

void Limiter::updateParameters() {
  const auto desiredLookaheadSeconds =
      emscripten_atomic_load_f32(&this->desiredLookaheadSeconds);
  const auto desiredAttackSeconds =
      emscripten_atomic_load_f32(&this->desiredAttackSeconds);
  const auto desiredReleaseSeconds =
      emscripten_atomic_load_f32(&this->desiredReleaseSeconds);

  const int desiredLookaheadFrames =
      desiredLookaheadSeconds * (double)sampleRate * (double)upsampleFactor;

  if (desiredAttackSeconds != attackSeconds) {
    attackSeconds = desiredAttackSeconds;
    attackAlpha = timeToAlpha(attackSeconds, sampleRate * upsampleFactor);
  }
  if (desiredReleaseSeconds != releaseSeconds) {
    releaseSeconds = desiredReleaseSeconds;
    releaseAlpha = timeToAlpha(releaseSeconds, sampleRate * upsampleFactor);
  }
  if (desiredLookaheadFrames != lookaheadFrames) {
    lookaheadFrames = desiredLookaheadFrames;

    delayFrames = speex_resampler_get_input_latency(upsampler) +
                  speex_resampler_get_output_latency(downsampler) +
                  lookaheadFrames / upsampleFactor;

    lookaheadDelayLine =
        std::make_unique<DelayLine>(channelCount, lookaheadFrames);
    deque.resize(workBufferSize + lookaheadFrames + 1);
    slidingWindowWorkBuffer =
        BufferF32(channelCount, workBufferSize + lookaheadFrames);
  }
}

void Limiter::process(BufferF32 buf) {
  updateParameters();

  const auto bypass = emscripten_atomic_load_u32(&this->bypass) != 0;
  const auto stereoLink = emscripten_atomic_load_f32(&this->stereoLink);
  const auto preGainDb = emscripten_atomic_load_f32(&this->preGainDb);
  const auto preGain = std::pow(10.0f, preGainDb / 20.0f);

  const auto bufFrameCount = buf.getFrameCount();
  for (int i = 0; i < bufFrameCount; i += chunkSize) {
    auto thisChunk = buf.slice(i, std::min(i + chunkSize, (int)bufFrameCount));

    const auto upsampledLength = upsampleFactor * thisChunk.getFrameCount();
    auto workSlice = workBuffer.slice(0, upsampledLength);

    if (!bypass) {
      for (int channel = 0; channel < channelCount; channel++) {
        auto channelData = thisChunk.getChannelData(channel);

#pragma clang loop vectorize(enable)
        for (int j = 0; j < thisChunk.getFrameCount(); j++) {
          channelData[j] *= preGain;
        }
      }
    }

    // upsample
    for (int channel = 0; channel < channelCount; channel++) {
      spx_uint32_t in_len = thisChunk.getFrameCount();
      spx_uint32_t out_len = upsampledLength;
      speex_resampler_process_float(
          upsampler, channel, thisChunk.getChannelData(channel), &in_len,
          workBuffer.getChannelData(channel), &out_len);
      assert(in_len == thisChunk.getFrameCount());
      assert(out_len == upsampledLength);
    }

    // assemble [delayed, current] view
    for (int channel = 0; channel < channelCount; ++channel) {
      auto slidingWindowBufferData =
          slidingWindowWorkBuffer.getChannelData(channel);
      lookaheadDelayLine->read(slidingWindowBufferData, channel, 0,
                               lookaheadFrames);
      memcpy(slidingWindowBufferData + lookaheadFrames,
             workSlice.getChannelData(channel),
             upsampledLength * sizeof(float));
    }

    // make workSlice a prefix of [delayed, current]
    lookaheadDelayLine->pushpull(workSlice);

    // calculate target gains
    for (int channel = 0; channel < channelCount; channel++) {
      auto channelData = workSlice.getChannelData(channel);
      auto peaksInWorkslice = limiterWorkBuffer.getChannelData(channel);

      auto currentGainPrefilter = channelCurrentGainEnvelope[channel];
      auto slidingWindowBufferData =
          slidingWindowWorkBuffer.getChannelData(channel);

      vfastnanzero_array(slidingWindowBufferData, slidingWindowBufferData,
                         upsampledLength + lookaheadFrames);

      vfastabs_array(slidingWindowBufferData, slidingWindowBufferData,
                     upsampledLength + lookaheadFrames);

      slidingWindowMax(slidingWindowBufferData, peaksInWorkslice, deque.data(),
                       upsampledLength + lookaheadFrames, lookaheadFrames,
                       upsampledLength);
    }

    if (channelCount == 1) {
      // mono: no linking
      auto channelData = workSlice.getChannelData(0);
      auto peaksInWorkslice = limiterWorkBuffer.getChannelData(0);
      auto currentGainEnvelope = channelCurrentGainEnvelope[0];

      for (int i = 0; i < upsampledLength; i++) {
        float v_max = peaksInWorkslice[i];

        const float targetGain = 0.97f / std::max(0.97f, v_max);
        const auto alpha =
            targetGain < currentGainEnvelope ? attackAlpha : releaseAlpha;
        currentGainEnvelope =
            currentGainEnvelope + (targetGain - currentGainEnvelope) * alpha;
        if (!bypass) {
          channelData[i] *= currentGainEnvelope;
        }
      }

      channelCurrentGainEnvelope[0] = currentGainEnvelope;
    } else {
      // stereo: adjustable linking
      auto leftChannelData = workSlice.getChannelData(0);
      auto rightChannelData = workSlice.getChannelData(1);
      auto leftPeaksInWorkslice = limiterWorkBuffer.getChannelData(0);
      auto rightPeaksInWorkslice = limiterWorkBuffer.getChannelData(1);
      auto leftCurrentGainEnvelope = channelCurrentGainEnvelope[0];
      auto rightCurrentGainEnvelope = channelCurrentGainEnvelope[1];
      auto linkedCurrentGainEnvelope = this->linkedCurrentGainEnvelope;

      for (int i = 0; i < upsampledLength; i++) {
        const auto leftPeak = leftPeaksInWorkslice[i];
        const auto rightPeak = rightPeaksInWorkslice[i];
        const auto linkedPeak = std::max(leftPeak, rightPeak);
        const auto leftTargetGain = 0.97f / std::max(0.97f, leftPeak);
        const auto rightTargetGain = 0.97f / std::max(0.97f, rightPeak);
        const auto linkedTargetGain = 0.97f / std::max(0.97f, linkedPeak);
        const auto leftAlpha = leftTargetGain < leftCurrentGainEnvelope
                                   ? attackAlpha
                                   : releaseAlpha;
        const auto rightAlpha = rightTargetGain < rightCurrentGainEnvelope
                                    ? attackAlpha
                                    : releaseAlpha;
        const auto linkedAlpha = linkedTargetGain < linkedCurrentGainEnvelope
                                     ? attackAlpha
                                     : releaseAlpha;
        leftCurrentGainEnvelope =
            leftCurrentGainEnvelope +
            (leftTargetGain - leftCurrentGainEnvelope) * leftAlpha;
        rightCurrentGainEnvelope =
            rightCurrentGainEnvelope +
            (rightTargetGain - rightCurrentGainEnvelope) * rightAlpha;
        linkedCurrentGainEnvelope =
            linkedCurrentGainEnvelope +
            (linkedTargetGain - linkedCurrentGainEnvelope) * linkedAlpha;

        const auto commonGain = stereoLink * linkedCurrentGainEnvelope;
        const auto diffFactor = 1.f - stereoLink;
        const auto leftGain = commonGain + diffFactor * leftCurrentGainEnvelope;
        const auto rightGain =
            commonGain + diffFactor * rightCurrentGainEnvelope;
        if (!bypass) {
          leftChannelData[i] *= leftGain;
          rightChannelData[i] *= rightGain;
        }
      }

      channelCurrentGainEnvelope[0] = leftCurrentGainEnvelope;
      channelCurrentGainEnvelope[1] = rightCurrentGainEnvelope;
      this->linkedCurrentGainEnvelope = linkedCurrentGainEnvelope;
    }

    // downsample
    for (int channel = 0; channel < channelCount; channel++) {
      spx_uint32_t in_len = upsampledLength;
      spx_uint32_t out_len = thisChunk.getFrameCount();
      speex_resampler_process_float(
          downsampler, channel, workSlice.getChannelData(channel), &in_len,
          thisChunk.getChannelData(channel), &out_len);
      assert(in_len == upsampledLength);
      assert(out_len == thisChunk.getFrameCount());
    }

    // safety
    for (int channel = 0; channel < channelCount; channel++) {
      vfastclamp_nanzero_array(thisChunk.getChannelData(channel), -1.f, 1.f,
                               thisChunk.getChannelData(channel),
                               thisChunk.getFrameCount());
    }
  }
}

void Limiter::setParameters(float lookaheadSeconds, float attackSeconds,
                            float releaseSeconds, float preGainDb, bool bypass,
                            float stereoLink) {
  assert(std::isfinite(lookaheadSeconds));
  assert(std::isfinite(attackSeconds));
  assert(std::isfinite(releaseSeconds));
  assert(!std::isnan(preGainDb));
  assert(std::isfinite(stereoLink));
  assert(lookaheadSeconds >= 0.0f);
  assert(lookaheadSeconds < 30.0f);
  assert(attackSeconds >= 0.0f);
  assert(releaseSeconds >= 0.0f);
  assert(stereoLink >= 0.0f && stereoLink <= 1.0f);

  emscripten_atomic_store_f32(&desiredLookaheadSeconds, lookaheadSeconds);
  emscripten_atomic_store_f32(&desiredAttackSeconds, attackSeconds);
  emscripten_atomic_store_f32(&desiredReleaseSeconds, releaseSeconds);

  emscripten_atomic_store_f32(&this->preGainDb, preGainDb);
  emscripten_atomic_store_u32(&this->bypass, bypass ? 1 : 0);
  emscripten_atomic_store_f32(&this->stereoLink, stereoLink);
}

// Naive brute-force sliding window maximum for testing
void naiveSlidingWindowMax(const float *input, float *output, int inputLength,
                           int windowLength) {
  for (int i = 0; i < inputLength - windowLength + 1; i++) {
    float maxVal = input[i];
    for (int j = i + 1; j < i + windowLength; j++) {
      if (input[j] > maxVal) {
        maxVal = input[j];
      }
    }
    output[i] = maxVal;
  }
}

TEST_CASE("slidingWindowMax", "[slidingwindowmax]") {
  SECTION("Compare with naive implementation on random data") {
    // Test with different window sizes and input lengths
    std::vector<std::pair<int, int>> testCases = {
        {10, 3},   // inputLength=10, windowLength=3
        {20, 5},   // inputLength=20, windowLength=5
        {100, 10}, // inputLength=100, windowLength=10
        {50, 1},   // windowLength=1 (edge case)
        {15, 15},  // windowLength equals inputLength (edge case)
    };

    for (auto [inputLength, windowLength] : testCases) {
      SECTION("inputLength=" + std::to_string(inputLength) +
              ", windowLength=" + std::to_string(windowLength)) {

        // Generate random input data
        std::vector<float> input(inputLength);
        std::mt19937 gen(42 + inputLength + windowLength); // Deterministic seed
        std::uniform_real_distribution<float> dist(-100.0f,
                                                   100.0f); // Range [-100, 100]
        for (int i = 0; i < inputLength; i++) {
          input[i] = dist(gen);
        }

        // Prepare output buffers
        int outputLength = inputLength - windowLength + 1;
        std::vector<float> efficientOutput(outputLength, 0.0f);
        std::vector<float> naiveOutput(outputLength, 0.0f);
        std::vector<int> deque(inputLength + 1);

        // Run both implementations
        slidingWindowMax(input.data(), efficientOutput.data(), deque.data(),
                         inputLength, windowLength, naiveOutput.size());
        naiveSlidingWindowMax(input.data(), naiveOutput.data(), inputLength,
                              windowLength);

        // Compare outputs
        for (int i = 0; i < outputLength; i++) {
          REQUIRE(std::abs(efficientOutput[i] - naiveOutput[i]) < 1e-6f);
        }
      }
    }
  }

  SECTION("Test with known input patterns") {
    SECTION("Monotonic increasing") {
      std::vector<float> input = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f};
      std::vector<float> expectedOutput = {3.0f, 4.0f, 5.0f}; // window size 3
      std::vector<float> actualOutput(3);
      std::vector<int> deque(6);

      slidingWindowMax(input.data(), actualOutput.data(), deque.data(), 5, 3,
                       actualOutput.size());

      for (int i = 0; i < 3; i++) {
        REQUIRE(actualOutput[i] == expectedOutput[i]);
      }
    }

    SECTION("Monotonic decreasing") {
      std::vector<float> input = {5.0f, 4.0f, 3.0f, 2.0f, 1.0f};
      std::vector<float> expectedOutput = {5.0f, 4.0f, 3.0f}; // window size 3
      std::vector<float> actualOutput(3);
      std::vector<int> deque(6);

      slidingWindowMax(input.data(), actualOutput.data(), deque.data(), 5, 3,
                       actualOutput.size());

      for (int i = 0; i < 3; i++) {
        REQUIRE(actualOutput[i] == expectedOutput[i]);
      }
    }

    SECTION("Peak in middle") {
      std::vector<float> input = {1.0f, 2.0f, 10.0f, 3.0f, 1.0f};
      std::vector<float> expectedOutput = {10.0f, 10.0f,
                                           10.0f}; // window size 3
      std::vector<float> actualOutput(3);
      std::vector<int> deque(6);

      slidingWindowMax(input.data(), actualOutput.data(), deque.data(), 5, 3,
                       actualOutput.size());

      for (int i = 0; i < 3; i++) {
        REQUIRE(actualOutput[i] == expectedOutput[i]);
      }
    }

    SECTION("All equal values") {
      std::vector<float> input = {5.0f, 5.0f, 5.0f, 5.0f, 5.0f};
      std::vector<float> expectedOutput = {5.0f, 5.0f, 5.0f}; // window size 3
      std::vector<float> actualOutput(3);
      std::vector<int> deque(6);

      slidingWindowMax(input.data(), actualOutput.data(), deque.data(), 5, 3,
                       actualOutput.size());

      for (int i = 0; i < 3; i++) {
        REQUIRE(actualOutput[i] == expectedOutput[i]);
      }
    }
  }

  SECTION("Edge cases") {
    SECTION("Window size 1") {
      std::vector<float> input = {3.0f, 1.0f, 4.0f, 1.0f, 5.0f};
      std::vector<float> expectedOutput = {3.0f, 1.0f, 4.0f, 1.0f, 5.0f};
      std::vector<float> actualOutput(5);
      std::vector<int> deque(6);

      slidingWindowMax(input.data(), actualOutput.data(), deque.data(), 5, 1,
                       actualOutput.size());

      for (int i = 0; i < 5; i++) {
        REQUIRE(actualOutput[i] == expectedOutput[i]);
      }
    }

    SECTION("Negative values") {
      std::vector<float> input = {-5.0f, -2.0f, -8.0f, -1.0f};
      std::vector<float> expectedOutput = {-2.0f, -1.0f}; // window size 3
      std::vector<float> actualOutput(2);
      std::vector<int> deque(5);

      slidingWindowMax(input.data(), actualOutput.data(), deque.data(), 4, 3,
                       actualOutput.size());

      for (int i = 0; i < 2; i++) {
        REQUIRE(actualOutput[i] == expectedOutput[i]);
      }
    }
  }
}

TEST_CASE("limiter", "[limiter]") {
  Limiter limiter(2, 48000, 0.01f, 0.01f, 0.01f, 1.0f, false, 0.8f);
  BufferF32 buf(2, 48000 * 10);
  buf.noise();
  for (int channel = 0; channel < buf.getChannelCount(); channel++) {
    for (int i = 0; i < buf.getFrameCount(); i++) {
      buf[channel][i] *= 2.f;
    }
  }

  limiter.process(buf);

  bool allUnderOne = true;
  for (int channel = 0; channel < buf.getChannelCount() && allUnderOne;
       channel++) {
    for (int i = 0; i < buf.getFrameCount(); i++) {
      if (std::fabs(buf[channel][i]) > 1.f) {
        allUnderOne = false;
        break;
      }
    }
  }
  REQUIRE(allUnderOne);

  // asan
  limiter.process(buf);
  limiter.setParameters(1.0f, 0.01f, 0.01f, 1.0f, false, 0.0f);
  limiter.process(buf);
  limiter.setParameters(0.01f, 1.0f, 0.01f, 1.0f, false, 0.0f);
  limiter.process(buf);
}