#include "mipmapworker.h"
#include "util.h"
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <emscripten/bind.h>
#include <emscripten/console.h>
#include <limits>

// Atomics.waitAsync polyfill for browsers that don't support it
// 7% of desktop browsers as of 2025-11 according to caniuse.com

// clang-format off
EM_JS(emscripten::EM_VAL, _async_wait_do_polyfill, (emscripten::EM_VAL heapHandle, uintptr_t address, uint32_t expectedValue, double timeout), {
  const heap = Emval.toValue(heapHandle);
  const start = performance.now();
  const promise = (async () => {
    while (performance.now() - start < timeout) {
      if (Atomics.load(heap, address) !== expectedValue) {
        return "ok";
      }
      await new Promise(resolve => setTimeout(resolve, 10));
    }
    return "timed-out";
  })();
  return Emval.toHandle({
    async: true,
    value: promise
  });
});
// clang-format on

AsyncWaitPolyfill::AsyncWaitPolyfill() {
  atomics = emscripten::val::global("Atomics");
  getGrowableHeapI32 = emscripten::val::module_property("GROWABLE_HEAP_I32");
  hasWaitAsync = atomics.hasOwnProperty("waitAsync");
}

AsyncWaitResult AsyncWaitPolyfill::waitAsync(uintptr_t address,
                                             uint32_t expectedValue,
                                             double timeout) const {
  if (hasWaitAsync) {
    return atomics.call<AsyncWaitResult>("waitAsync", getGrowableHeapI32(),
                                         address / sizeof(int32_t),
                                         expectedValue, timeout);
  } else {
    return emscripten::val::take_ownership(
               _async_wait_do_polyfill(getGrowableHeapI32().as_handle(),
                                       address / sizeof(int32_t), expectedValue,
                                       timeout))
        .as<AsyncWaitResult>();
  }
}

void MipMapWorker::mipMapWorkerStart(int payloadPtr) {
  auto payload = reinterpret_cast<MipMapWorker *>(payloadPtr);
  payload->workerEntrypoint();
}

MipMapWorker::MipMapWorker() {
  emscripten_lock_init(&bufferCacheLock);
  emscripten_lock_init(&lock);
  emscripten_condvar_init(&condvar);

  emscripten_wasm_worker_t worker =
      emscripten_malloc_wasm_worker(/*stackSize: */ 3670016);
  emscripten_wasm_worker_post_function_vi(worker, mipMapWorkerStart,
                                          reinterpret_cast<int>(this));
}

MipMapWorker::~MipMapWorker() {}

// TODO these references should track the ones used by playback controller
// so we can deallocate old clips
std::shared_ptr<FfmpegAudioBuffer>
MipMapWorker::getBufferCached(const std::string &uuid) {
  {
    emscripten_lock_raii guard(bufferCacheLock);
    auto it = bufferCache.find(uuid);
    if (it != bufferCache.end()) {
      return it->second;
    }
  }
  // TODO ideally this would happen while holding the lock
  // to prevent duplicate creation, but for now assertions
  // firing in FfmpegAudioBuffer::createFromUuid can damage
  // the lock if it's held, which could deadlock the UI thread
  // don't take that risk
  auto buffer = FfmpegAudioBuffer::createFromUuid(uuid);
  if (buffer) {
    emscripten_lock_raii guard(bufferCacheLock);
    bufferCache[uuid] = buffer;
  }
  return buffer;
}

std::shared_ptr<MipMapRequest> MipMapWorker::waitForWork() {
  emscripten_lock_raii guard(lock);

  for (;;) {
    if (!workQueue.empty()) {
      auto request = std::move(workQueue.front());
      workQueue.pop_front();
      return request;
    }

    emscripten_condvar_waitinf(&condvar, &lock);
  }
}

void MipMapWorker::workerEntrypoint() {
  for (;;) {
    auto request = waitForWork();

    request->run();
  }
}

std::shared_ptr<MipMapRequest>
MipMapWorker::createRequest(const std::string &uuid, int mipMapLevel, int start,
                            int end) {
  assert(end >= start);
  assert(mipMapLevel >= 0);

  // TODO this should really happen on the worker
  auto buffer = getBufferCached(uuid);
  if (!buffer) {
    return nullptr;
  }

  // Reject requests beyond available duration to prevent accumulation of
  // polling pages
  const auto samplesPerBucket = 1UL << mipMapLevel;
  const auto startSample = start * samplesPerBucket;
  const auto availableDuration = buffer->getAvailableDuration();
  const auto sampleRate = buffer->getSampleRate();
  const auto availableSamples = availableDuration * sampleRate;

  if (startSample > availableSamples) {
    return nullptr;
  }

  emscripten_lock_raii guard(lock);

  if (workQueue.size() >= QUEUE_LIMIT) {
    return nullptr;
  }

  auto request = std::make_shared<MipMapRequest>(std::move(buffer), mipMapLevel,
                                                 start, end);
  workQueue.push_back(request);
  emscripten_condvar_signal(&condvar, EMSCRIPTEN_NOTIFY_ALL_WAITERS);

  return request;
}

std::vector<int16_t> MipMapRequest::makeSharedMinMaxArray(int blockCount) {
  std::vector<int16_t> result(blockCount * 2);
  for (auto i = 0; i < blockCount; i++) {
    result[i * 2] = std::numeric_limits<int16_t>::max();
    result[i * 2 + 1] = std::numeric_limits<int16_t>::min();
  }
  return result;
}

MipMapRequest::MipMapRequest(std::shared_ptr<FfmpegAudioBuffer> &&buffer,
                             int mipMapLevel, int start, int end)
    : buffer(std::move(buffer)), mipMapLevel(mipMapLevel), start(start),
      end(end), status(0), blocksProduced(0) {
  for (auto ch = 0; ch < 2; ch++) {
    outputs[ch] = makeSharedMinMaxArray(end - start);
  }
}

MipMapRequest::~MipMapRequest() {}

int MipMapRequest::getStatus() const {
  return emscripten_atomic_load_u32(&status);
}

int MipMapRequest::getBlocksProduced() const { return blocksProduced; }

std::optional<int> MipMapRequest::getContinueFrom() const {
  return continueFrom;
}

MipMapOutputBuffers MipMapRequest::getOutputBuffers() const {
  auto heapBufferI16 = emscripten::val::module_property("GROWABLE_HEAP_I16")();
  auto result = emscripten::val::array();
  auto typedArrayConstructor = emscripten::val::global("Int16Array");
  for (auto ch = 0; ch < 2; ch++) {
    auto typedArray = typedArrayConstructor.new_(outputs[0].size());

    const auto outHeapIndex =
        reinterpret_cast<uintptr_t>(outputs[ch].data()) / sizeof(int16_t);

    typedArray.call<void>("set", heapBufferI16.call<emscripten::val>(
                                     "subarray", outHeapIndex,
                                     outHeapIndex + outputs[ch].size()));

    result.call<void>("push", typedArray);
  }

  return result.as<MipMapOutputBuffers>();
}

void MipMapRequest::finalize(int finalStatus) {
  for (;;) {
    const auto currentStatus = emscripten_atomic_load_u32(&status);
    if (currentStatus != 0) {
      break;
    }
    if (currentStatus ==
        emscripten_atomic_cas_u32(&status, currentStatus, finalStatus)) {
      emscripten_atomic_notify(&status, EMSCRIPTEN_NOTIFY_ALL_WAITERS);
      break;
    }
  }
}

void MipMapRequest::cancel() { finalize(2); }

AsyncWaitResult MipMapRequest::waitAsync() const {
  return asyncWaitPolyfill.waitAsync(reinterpret_cast<uintptr_t>(&status), 0,
                                     30000.0);
}

void MipMapRequest::run() {
  if (getStatus() != 0) {
    return;
  }

  if (!buffer) {
    emscripten_console_warn("Buffer not found for MipMapRequest");
    finalize(2);
  }

  const auto bufferIsCompleted = buffer->isCompleted();

  const auto samplesPerBucket = 1UL << mipMapLevel;
  const auto rangeStart = start * samplesPerBucket;
  const auto rangeEnd = end * samplesPerBucket;

  const auto blockCount = std::max(0, end - start);

  constexpr const int WORK_BUF_SIZE = 4096;

  BufferF32 workBuf(2, WORK_BUF_SIZE);

  auto pos = rangeStart;
  while (pos < rangeEnd && getStatus() == 0) {
    buffer->read(pos, workBuf);

    auto framesHave = std::min(WORK_BUF_SIZE, buffer->getNextFrameOffset() -
                                                  static_cast<int>(pos));

    if (framesHave <= 0) {
      break;
    }

    for (auto ch = 0; ch < workBuf.getChannelCount(); ch++) {
      auto sampleData = workBuf.getChannelData(ch);
      auto &out = outputs[ch];
      for (auto j = 0; j < framesHave; j++) {
        const int bucketIndex = static_cast<int>(
            std::floor((pos - rangeStart + j) / samplesPerBucket));
        const int base = bucketIndex * 2;
        if (base + 1 >= out.size()) {
          break;
        }
        const int16_t s = static_cast<int16_t>(std::clamp(
            std::round(sampleData[j] * 32767.f),
            static_cast<float>(std::numeric_limits<int16_t>::min()),
            static_cast<float>(std::numeric_limits<int16_t>::max())));
        if (s < out[base]) {
          out[base] = s;
        }
        if (s > out[base + 1]) {
          out[base + 1] = s;
        }
      }
    }

    pos += framesHave;
  }

  blocksProduced = std::max(
      0, static_cast<int>(std::floor((pos - rangeStart) / samplesPerBucket)));

  if (pos < rangeEnd && !bufferIsCompleted) {
    continueFrom = start + blocksProduced;
  }

  finalize(1);
}

EMSCRIPTEN_BINDINGS(mipmaprequest) {
  using namespace emscripten;

  register_type<AsyncWaitResult>("{ async: boolean, value: Promise<string> }");
  register_type<MipMapOutputBuffers>("Int16Array[]");
  register_optional<int>();

  class_<MipMapRequest>("MipMapRequest")
      .smart_ptr<std::shared_ptr<MipMapRequest>>("MipMapRequest")
      .property("status", &MipMapRequest::getStatus)
      .property("blocksProduced", &MipMapRequest::getBlocksProduced)
      .property("continueFrom", &MipMapRequest::getContinueFrom)
      .function("getOutputBuffers", &MipMapRequest::getOutputBuffers)
      .function("cancel", &MipMapRequest::cancel)
      .function("waitAsync", &MipMapRequest::waitAsync);

  class_<MipMapWorker>("MipMapWorker")
      .smart_ptr<std::shared_ptr<MipMapWorker>>("MipMapWorker")
      .function("createRequest", &MipMapWorker::createRequest);
};