#pragma once
#include "fft/pffft.h"
#include "simd.h"
#include <cassert>
#include <emscripten/val.h>
#include <optional>
#include <random>
#include <type_traits>

// A Buffer is a simple wrapper around a list of float arrays (one per audio
// channel) that provides a few convenience methods for working with audio data.

// Buffer<T> can own and modify data
// Buffer<const T> can only view data and never takes ownership

EMSCRIPTEN_DECLARE_VAL_TYPE(Float32Array);
EMSCRIPTEN_DECLARE_VAL_TYPE(Float64Array);
EMSCRIPTEN_DECLARE_VAL_TYPE(ChannelArrayF32);
EMSCRIPTEN_DECLARE_VAL_TYPE(ChannelArrayF64);

template <typename T> struct EmscriptenFloatTraits {};

template <> struct EmscriptenFloatTraits<float> {
  using SampleArrayValType = Float32Array;
  using ChannelArrayValType = ChannelArrayF32;
  static constexpr const char *sampleArrayJSTypeName = "Float32Array";
  static constexpr const char *moduleHeapAccessor = "GROWABLE_HEAP_F32";
};

template <> struct EmscriptenFloatTraits<double> {
  using SampleArrayValType = Float64Array;
  using ChannelArrayValType = ChannelArrayF64;
  static constexpr const char *sampleArrayJSTypeName = "Float64Array";
  static constexpr const char *moduleHeapAccessor = "GROWABLE_HEAP_F64";
};

template <typename T> class Buffer final {
public:
  using SampleArrayValType =
      EmscriptenFloatTraits<std::remove_const_t<T>>::SampleArrayValType;
  using ChannelArrayValType =
      EmscriptenFloatTraits<std::remove_const_t<T>>::ChannelArrayValType;
  using sampleType = T;

private:
  template <typename U> friend class Buffer;

  static constexpr const size_t MAX_CHANNEL_COUNT = 32;

  std::conditional_t<std::is_const_v<T>, const T *, T *>
      _data[MAX_CHANNEL_COUNT];
  int _channelCount;
  size_t _frameCount;
  bool _ownsData;

  template <typename U>
  Buffer(int channelCount, size_t frameCount, U *const *data) noexcept
      : _channelCount(channelCount), _frameCount(frameCount), _ownsData(false) {
    for (int i = 0; i < channelCount; i++) {
      _data[i] = data[i];
    }
  }

  Buffer(int channelCount) noexcept
      : _channelCount(channelCount), _frameCount(0), _ownsData(false) {}

  emscripten::val getModuleHeapFloatBuffer() const {
    auto moduleHeapF =
        val::module_property(EmscriptenFloatTraits<T>::moduleHeapAccessor)();
    return moduleHeapF["buffer"];
  }

  template <typename BufferType>
  static BufferType slice_impl(typename BufferType::sampleType *const *data,
                               int channelCount, size_t frameCount,
                               int startFrame, int endFrame) {
    if (startFrame < 0) {
      startFrame = frameCount + startFrame;
    }
    if (endFrame < 0) {
      endFrame = frameCount + endFrame;
    }
    assert(startFrame >= 0 && startFrame < frameCount);
    assert(endFrame >= 0 && endFrame <= frameCount);
    if (startFrame >= endFrame) {
      return BufferType(channelCount);
    }
    typename BufferType::sampleType *view[MAX_CHANNEL_COUNT];
    for (int i = 0; i < channelCount; i++) {
      view[i] = data[i] + startFrame;
    }
    return BufferType(channelCount, endFrame - startFrame, view);
  }

public:
  using val = emscripten::val;

  static Buffer fromVLA(int channelCount, size_t frameCount, T *data) noexcept {
    T *ptrs[channelCount];
    for (int i = 0; i < channelCount; i++) {
      ptrs[i] = data + i * frameCount;
    }
    return Buffer(channelCount, frameCount, ptrs);
  }

  static std::shared_ptr<Buffer> fromArray(const ChannelArrayValType &a) {
    assert(a.isArray());
    auto channelCount = a["length"].template as<size_t>();
    std::optional<size_t> frameCount;
    for (int i = 0; i < channelCount; i++) {
      const auto &channel = a[i];
      assert(channel.instanceof(
          val::global(EmscriptenFloatTraits<T>::sampleArrayJSTypeName)));
      auto channelSampleCount = channel["length"].template as<size_t>();
      if (frameCount.has_value()) {
        assert(channelSampleCount == frameCount.value());
      } else {
        frameCount = channelSampleCount;
      }
    }
    auto buffer =
        std::make_shared<Buffer<T>>(channelCount, frameCount.value_or(0));
    for (int i = 0; i < channelCount; i++) {
      const auto &channel = a[i];
      buffer->js_set(i, channel.template as<SampleArrayValType>());
    }
    return buffer;
  }

  template <typename U = T> operator Buffer<const U>() const noexcept {
    static_assert(std::is_same_v<U, T>, "Type mismatch in conversion operator");
    return Buffer<const U>(_channelCount, _frameCount, _data);
  }

  Buffer &operator=(const Buffer &other) noexcept {
    if (this != &other) {
      if constexpr (!std::is_const_v<T>) {
        if (_ownsData) {
          for (int i = 0; i < _channelCount; i++) {
            pffft_aligned_free(_data[i]);
          }
        }
      }

      _ownsData = false;
      _channelCount = other._channelCount;
      _frameCount = other._frameCount;
      for (int i = 0; i < _channelCount; i++) {
        _data[i] = other._data[i];
      }
    }
    return *this;
  }

  Buffer(const Buffer &other) noexcept
      : _channelCount(other._channelCount), _frameCount(other._frameCount),
        _ownsData(false) {
    for (int i = 0; i < _channelCount; i++) {
      _data[i] = other._data[i];
    }
  }

  Buffer() noexcept : _channelCount(0), _frameCount(0), _ownsData(false) {}

  template <typename U = T, typename = std::enable_if_t<!std::is_const_v<U>>>
  Buffer(int channelCount, size_t frameCount)
      : _channelCount(channelCount), _frameCount(frameCount), _ownsData(true) {
    assert(channelCount >= 0 && channelCount <= MAX_CHANNEL_COUNT);
    assert(frameCount >= 0);
    if (frameCount == 0) {
      for (int i = 0; i < channelCount; i++) {
        _data[i] = nullptr;
      }
    } else {
      for (int i = 0; i < channelCount; i++) {
        _data[i] = (T *)pffft_aligned_malloc(frameCount * sizeof(T));
        memset(_data[i], 0, frameCount * sizeof(T));
      }
    }
  }

  ~Buffer() noexcept {
    if constexpr (!std::is_const_v<T>) {
      if (_ownsData && _frameCount > 0) {
        for (int i = 0; i < _channelCount; i++) {
          pffft_aligned_free(_data[i]);
        }
      }
    }
  }

  template <typename U = T,
            typename std::enable_if<std::is_const<U>::value, int>::type = 0>
  Buffer(Buffer &&other) noexcept = delete;

  template <typename U = T,
            typename std::enable_if<std::is_const<U>::value, int>::type = 0>
  Buffer &operator=(Buffer &&other) noexcept = delete;

  template <typename U = T, typename = std::enable_if_t<!std::is_const_v<U>>>
  Buffer(Buffer &&other) noexcept
      : _channelCount(other._channelCount), _frameCount(other._frameCount),
        _ownsData(other._ownsData) {
    for (int i = 0; i < _channelCount; i++) {
      _data[i] = other._data[i];
      other._data[i] = nullptr;
    }
    other._ownsData = false;
    other._channelCount = 0;
    other._frameCount = 0;
  }

  template <typename U = T, typename = std::enable_if_t<!std::is_const_v<U>>>
  Buffer &operator=(Buffer &&other) noexcept {
    if (this != &other) {
      if (_ownsData && _frameCount > 0) {
        for (int i = 0; i < _channelCount; i++) {
          pffft_aligned_free(_data[i]);
        }
      }
      _channelCount = other._channelCount;
      _frameCount = other._frameCount;
      for (int i = 0; i < _channelCount; i++) {
        _data[i] = other._data[i];
        other._data[i] = nullptr;
      }
      _ownsData = other._ownsData;
      other._ownsData = false;
      other._channelCount = 0;
      other._frameCount = 0;
    }
    return *this;
  }

  Buffer slice(int startFrame, int endFrame) {
    return slice_impl<Buffer>(_data, _channelCount, _frameCount, startFrame,
                              endFrame);
  }

  Buffer<const T> slice(int startFrame, int endFrame) const {
    return slice_impl<Buffer<const T>>(_data, _channelCount, _frameCount,
                                       startFrame, endFrame);
  }

  Buffer slice(int startFrame, int endFrame, int channelCount) {
    assert(channelCount >= 0 && channelCount <= _channelCount);
    return slice_impl<Buffer>(_data, channelCount, _frameCount, startFrame,
                              endFrame);
  }

  Buffer<const T> slice(int startFrame, int endFrame, int channelCount) const {
    assert(channelCount >= 0 && channelCount <= _channelCount);
    return slice_impl<Buffer<const T>>(_data, channelCount, _frameCount,
                                       startFrame, endFrame);
  }

  Buffer slice(int startFrame) { return slice(startFrame, _frameCount); }

  Buffer<const T> slice(int startFrame) const {
    return slice(startFrame, _frameCount);
  }

  Buffer sliceChannel(int wantChannel) {
    assert(wantChannel >= 0 && wantChannel < _channelCount);
    T *view[1];
    view[0] = _data[wantChannel];
    return Buffer(1, _frameCount, view);
  }

  Buffer<const T> sliceChannel(int wantChannel) const {
    assert(wantChannel >= 0 && wantChannel < _channelCount);
    const T *view[1];
    view[0] = _data[wantChannel];
    return Buffer<const T>(1, _frameCount, view);
  }

  void cloneChannel(int sourceChannel, int destChannel) {
    assert(sourceChannel >= 0 && sourceChannel < _channelCount);
    assert(destChannel >= 0 && destChannel < _channelCount);
    if (_frameCount == 0) {
      return;
    }
    memcpy(_data[destChannel], _data[sourceChannel], _frameCount * sizeof(T));
  }

  void set(int channelIndex, int offset, size_t frameCount, const T *data) {
    assert(channelIndex >= 0 && channelIndex < _channelCount);
    assert(offset >= 0 && offset + frameCount <= _frameCount);
    if (_frameCount == 0) {
      return;
    }
    memcpy(_data[channelIndex] + offset, data, frameCount * sizeof(T));
  }

  void set(int offset, const Buffer<const T> &source) {
    assert(offset >= 0 && offset + source.getFrameCount() <= _frameCount);
    assert(source._channelCount == _channelCount);
    if (_frameCount == 0) {
      return;
    }
    for (int i = 0; i < _channelCount; i++) {
      const auto sourceData = source._data[i];
      memcpy(_data[i] + offset, sourceData, source.getFrameCount() * sizeof(T));
    }
  }

  void fill(int offset, size_t frameCount, T value) {
    assert(offset >= 0 && offset + frameCount <= _frameCount);
    for (int i = 0; i < _channelCount; i++) {
      const auto channelData = _data[i];
#pragma clang loop vectorize(enable)
      for (size_t j = 0; j < frameCount; j++) {
        channelData[j + offset] = value;
      }
    }
  }

  void fill(T value) noexcept {
    for (int i = 0; i < _channelCount; i++) {
      const auto channelData = _data[i];
#pragma clang loop vectorize(enable)
      for (size_t j = 0; j < _frameCount; j++) {
        channelData[j] = value;
      }
    }
  }

  void noise() {
    static thread_local std::mt19937 generator;
    std::uniform_real_distribution<T> dist(-1, 1);
    for (int i = 0; i < _channelCount; i++) {
      const auto channelData = _data[i];
      for (size_t j = 0; j < _frameCount; j++) {
        channelData[j] = dist(generator);
      }
    }
  }

  bool hasNaN() const noexcept {
    for (int i = 0; i < _channelCount; i++) {
      const auto channelData = _data[i];
      // #pragma clang loop vectorize(enable)
      //  TODO this does not want to auto-vectorize
      for (size_t j = 0; j < _frameCount; j++) {
        if (channelData[j] != channelData[j]) {
          return true;
        }
      }
    }
    return false;
  }

  T peak() const noexcept {
    if (_channelCount == 0 || _frameCount == 0) {
      return 0;
    }
    std::remove_const_t<T> peak = _data[0][0];
    for (int i = 0; i < _channelCount; i++) {
      const auto channelData = _data[i];
      if constexpr (std::is_same_v<std::remove_const_t<T>, float>) {
        peak =
            std::max(vfastmax_abs_elem_array(channelData, _frameCount), peak);
      } else {
        for (size_t j = 0; j < _frameCount; j++) {
          peak = std::max(peak, std::fabs(channelData[j]));
        }
      }
    }
    return peak;
  }

  void mixDownInPlace() noexcept {
    if (_channelCount <= 1) {
      return;
    }
    T invNumChannels = static_cast<T>(1.0) / _channelCount;
    for (int i = 0; i < _channelCount; i++) {
      const auto channelData = _data[i];
#pragma clang loop vectorize(enable)
      for (size_t j = 0; j < _frameCount; j++) {
        channelData[j] *= invNumChannels;
      }
    }
    const auto channelZeroData = _data[0];
    for (int i = 1; i < _channelCount; i++) {
      const auto channelData = _data[i];
#pragma clang loop vectorize(enable)
      for (size_t j = 0; j < _frameCount; j++) {
        channelZeroData[j] += channelData[j];
      }
    }
  }

  void sumWith(const Buffer<T> &other) {
    assert(_channelCount == other._channelCount);
    assert(_frameCount == other._frameCount);
    for (int channel = 0; channel < _channelCount; channel++) {
      auto mySamples = _data[channel];
      const auto otherSamples = other._data[channel];
#pragma clang loop vectorize(enable)
      for (int i = 0; i < _frameCount; i++) {
        mySamples[i] += otherSamples[i];
      }
    }
  }

  void interleaveTo(Buffer<T> &other) const {
    if (_channelCount == 1) {
      other.set(0, *this);
      return;
    }
    assert(_channelCount == 2);
    assert(other._channelCount == 1);
    assert(2 * _frameCount == other._frameCount);
    auto otherData = other._data[0];
    auto myLeftData = _data[0];
    auto myRightData = _data[1];
#pragma clang loop vectorize(enable)
    for (int i = 0; i < _frameCount; i++) {
      otherData[2 * i] = myLeftData[i];
      otherData[2 * i + 1] = myRightData[i];
    }
  }

  void interleaveTo(const std::shared_ptr<Buffer<T>> &other) const {
    if (other != nullptr) {
      interleaveTo(*other);
    }
  }

  void deinterleaveFrom(const Buffer<const T> &other) {
    if (_channelCount == 1) {
      set(0, other);
      return;
    }
    assert(_channelCount == 2);
    assert(other._channelCount == 1);
    assert(2 * _frameCount == other._frameCount);
    auto otherData = other._data[0];
    auto myLeftData = _data[0];
    auto myRightData = _data[1];
#pragma clang loop vectorize(enable)
    for (int i = 0; i < _frameCount; i++) {
      myLeftData[i] = otherData[2 * i];
      myRightData[i] = otherData[2 * i + 1];
    }
  }

  void deinterleaveFrom(const std::shared_ptr<Buffer<T>> &other) {
    if (other != nullptr) {
      deinterleaveFrom(*other);
    }
  }

  void js_set(int channelIndex, const SampleArrayValType &source) {
    assert(source.instanceof(
        val::global(EmscriptenFloatTraits<T>::sampleArrayJSTypeName)));
    assert(channelIndex >= 0 && channelIndex < _channelCount);

    auto thisData = val::global(EmscriptenFloatTraits<T>::sampleArrayJSTypeName)
                        .new_(getModuleHeapFloatBuffer(),
                              reinterpret_cast<uintptr_t>(_data[channelIndex]),
                              _frameCount);
    thisData.template call<void>("set", source.template as<val>());
  }

  void js_set(int channelIndex, int offset, const SampleArrayValType &source) {
    assert(source.instanceof(
        val::global(EmscriptenFloatTraits<T>::sampleArrayJSTypeName)));
    assert(channelIndex >= 0 && channelIndex < _channelCount);
    assert(offset >= 0 && offset < _frameCount);

    auto thisData = val::global(EmscriptenFloatTraits<T>::sampleArrayJSTypeName)
                        .new_(getModuleHeapFloatBuffer(),
                              reinterpret_cast<uintptr_t>(_data[channelIndex]),
                              _frameCount);
    thisData.template call<void>("set", source.template as<val>(), offset);
  }

  void js_set(const ChannelArrayValType &source) {
    assert(source.isArray());
    auto channelCount = source["length"].template as<size_t>();
    assert(channelCount == _channelCount);
    for (int i = 0; i < channelCount; i++) {
      js_set(i, source[i].template as<SampleArrayValType>());
    }
  }

  SampleArrayValType js_view(int channelIndex) const {
    assert(channelIndex >= 0 && channelIndex < _channelCount);

    return val::global(EmscriptenFloatTraits<T>::sampleArrayJSTypeName)
        .new_(getModuleHeapFloatBuffer(),
              reinterpret_cast<uintptr_t>(_data[channelIndex]), _frameCount)
        .template as<SampleArrayValType>();
  }

  ChannelArrayValType js_view() const {
    auto result = val::array();
    for (int i = 0; i < _channelCount; i++) {
      result.call<void>("push", js_view(i));
    }
    return result.as<ChannelArrayValType>();
  }

  void js_setInto(int channelIndex, SampleArrayValType dest) const {
    assert(dest.instanceof(
        val::global(EmscriptenFloatTraits<T>::sampleArrayJSTypeName)));
    assert(channelIndex >= 0 && channelIndex < _channelCount);

    auto thisData = val::global(EmscriptenFloatTraits<T>::sampleArrayJSTypeName)
                        .new_(getModuleHeapFloatBuffer(),
                              reinterpret_cast<uintptr_t>(_data[channelIndex]),
                              _frameCount);
    dest.template call<void>("set", thisData);
  }

  void js_setInto(ChannelArrayValType dest) const {
    assert(dest.isArray());
    auto channelCount = dest["length"].template as<size_t>();
    assert(channelCount == _channelCount);
    for (int i = 0; i < channelCount; i++) {
      js_setInto(i, dest[i].template as<SampleArrayValType>());
    }
  }

  size_t getFrameCount() const noexcept { return _frameCount; }

  int getChannelCount() const noexcept { return _channelCount; }

  const T *getChannelData(int channelIndex) const {
    assert(channelIndex >= 0 && channelIndex < _channelCount);
    return _data[channelIndex];
  }

  T *getChannelData(int channelIndex) {
    assert(channelIndex >= 0 && channelIndex < _channelCount);
    return _data[channelIndex];
  }

  T *operator[](size_t i) { return getChannelData(i); }
  const T *operator[](size_t i) const { return getChannelData(i); }

  bool isSameStorage(const Buffer<T> &other) const {
    if (_channelCount != other._channelCount) {
      return false;
    }
    for (int i = 0; i < _channelCount; i++) {
      if (_data[i] != other._data[i]) {
        return false;
      }
    }
    return true;
  }
};

using BufferF32 = Buffer<float>;
using BufferF64 = Buffer<double>;