#include "appendonlybuffer.h"
#include "../util.h"
#include <catch2/catch.hpp>
#include <cstdlib>
#include <cstring>
#include <emscripten/atomic.h>
#include <emscripten/emscripten.h>
#include <vector>

AppendOnlyBuffer::AppendOnlyBuffer() { emscripten_lock_init(&appenderLock); }

AppendOnlyBuffer::~AppendOnlyBuffer() {
  const auto numPages = emscripten_atomic_load_u32(&atomicPageCount);
  for (int32_t i = 0; i < numPages; i++) {
    free(pages[i]);
  }
}

size_t AppendOnlyBuffer::size() const {
  return emscripten_atomic_load_u32(&atomicSize);
}

void AppendOnlyBuffer::append(size_t count, const uint8_t *data) {
  emscripten_lock_raii guard(appenderLock);
  uint32_t pageCount = emscripten_atomic_load_u32(&atomicPageCount);
  uint32_t totalSize = emscripten_atomic_load_u32(&atomicSize);

  while (count > 0) {
    const auto bytesFree = AOB_PAGE_SIZE * pageCount - totalSize;
    if (bytesFree == 0) {
      // need a new page
      if (pageCount >= AOB_PAGE_MAX_COUNT) {
        break;
      }
      pages[pageCount++] = (uint8_t *)malloc(AOB_PAGE_SIZE);
      continue;
    }
    const auto pageIndex = totalSize / AOB_PAGE_SIZE;
    const auto writeOffset = totalSize % AOB_PAGE_SIZE;
    const auto pageToUse = std::min(bytesFree, count);
    memcpy(pages[pageIndex] + writeOffset, data, pageToUse);
    data += pageToUse;
    count -= pageToUse;
    totalSize += pageToUse;
  }
  emscripten_atomic_store_u32(&atomicPageCount, pageCount);
  emscripten_atomic_store_u32(&atomicSize, totalSize);
}

size_t AppendOnlyBuffer::read(size_t offset, size_t count, uint8_t *out) {
  const auto currentSize = emscripten_atomic_load_u32(&atomicSize);
  if (offset >= currentSize) {
    memset(out, 0, count);
    return 0;
  }
  const auto actualReadSize = std::min(currentSize - offset, count);
  auto actualCount = actualReadSize;

  while (actualCount > 0) {
    const auto pageIndex = offset / AOB_PAGE_SIZE;
    const auto pageOffset = offset % AOB_PAGE_SIZE;
    const auto pageToRead = std::min(AOB_PAGE_SIZE - pageOffset, actualCount);
    memcpy(out, pages[pageIndex] + pageOffset, pageToRead);
    out += pageToRead;
    actualCount -= pageToRead;
    offset += pageToRead;
  }

  if (actualReadSize < count) {
    memset(out, 0, count - actualReadSize);
  }

  return actualReadSize;
}

TEST_CASE("AppendOnlyBuffer single thread", "[appendonlybuffer]") {
  AppendOnlyBuffer buffer;

  SECTION("Empty buffer") {
    REQUIRE(buffer.size() == 0);

    uint8_t out[10];
    REQUIRE(buffer.read(0, 10, out) == 0);
  }

  SECTION("Single append and read") {
    const char *data = "hello";
    buffer.append(5, reinterpret_cast<const uint8_t *>(data));

    REQUIRE(buffer.size() == 5);

    uint8_t out[10];
    REQUIRE(buffer.read(0, 5, out) == 5);
    REQUIRE(memcmp(out, data, 5) == 0);
  }

  SECTION("Multiple appends in order") {
    const char *data1 = "hello";
    const char *data2 = "world";

    buffer.append(5, reinterpret_cast<const uint8_t *>(data1));
    buffer.append(5, reinterpret_cast<const uint8_t *>(data2));

    REQUIRE(buffer.size() == 10);

    uint8_t out[10];
    REQUIRE(buffer.read(0, 10, out) == 10);
    REQUIRE(memcmp(out, "helloworld", 10) == 0);
  }

  SECTION("Partial reads") {
    const char *data = "hello world";
    buffer.append(11, reinterpret_cast<const uint8_t *>(data));

    uint8_t out[5];
    REQUIRE(buffer.read(0, 5, out) == 5);
    REQUIRE(memcmp(out, "hello", 5) == 0);

    REQUIRE(buffer.read(6, 5, out) == 5);
    REQUIRE(memcmp(out, "world", 5) == 0);
  }

  SECTION("Read beyond buffer") {
    const char *data = "hello";
    buffer.append(5, reinterpret_cast<const uint8_t *>(data));

    uint8_t out[10];
    memset(out, 0xFF, 10); // Fill with sentinel values

    REQUIRE(buffer.read(0, 10, out) == 5); // Only 5 bytes available
    REQUIRE(memcmp(out, "hello", 5) == 0);
    // Remaining bytes should be zeroed
    for (int i = 5; i < 10; i++) {
      REQUIRE(out[i] == 0);
    }
  }

  SECTION("Read from offset beyond buffer") {
    const char *data = "hello";
    buffer.append(5, reinterpret_cast<const uint8_t *>(data));

    uint8_t out[5];
    memset(out, 0xFF, 5);

    REQUIRE(buffer.read(10, 5, out) == 0); // Offset beyond buffer
    // All bytes should be zeroed
    for (int i = 0; i < 5; i++) {
      REQUIRE(out[i] == 0);
    }
  }

  SECTION("Zero-length operations") {
    buffer.append(0, nullptr);
    REQUIRE(buffer.size() == 0);

    uint8_t out[1];
    REQUIRE(buffer.read(0, 0, out) == 0);
  }

  SECTION("Large data capacity growth") {
    const size_t largeDataSize = 1024 * 1024 * 3 + 123;
    std::vector<uint8_t> largeData(largeDataSize);
    for (size_t i = 0; i < largeData.size(); i++) {
      largeData[i] = static_cast<uint8_t>(i % 256);
    }

    buffer.append(largeData.size(), largeData.data());
    REQUIRE(buffer.size() == largeDataSize);

    for (int i = 0; i < 2; i++) {
      std::vector<uint8_t> readData(largeDataSize);
      REQUIRE(buffer.read(0, largeDataSize, readData.data()) == largeDataSize);
      REQUIRE(readData == largeData);
    }
  }
}

static emscripten_wasm_worker_t fakeRealtimeThread;

struct AppendOnlyBufferWorkerPayload {
  uint32_t sema{};
  AppendOnlyBuffer buffer;
  uint32_t appendCount{};
  uint32_t readCount{};
  std::vector<uint8_t> expectedData;
};

static constexpr const int NUM_APPEND_THREADS = 1;
static constexpr const int NUM_READ_THREADS = 2;
static constexpr const int NUM_APPENDS_PER_THREAD = 1000;
static constexpr const int APPEND_SIZE = 32;

static void append_worker_start(int payloadPtr) {
  auto payload = reinterpret_cast<AppendOnlyBufferWorkerPayload *>(payloadPtr);

  std::mt19937 gen(42); // Fixed seed for reproducible data
  std::uniform_int_distribution<> byteDis(0, 255);

  for (int i = 0; i < NUM_APPENDS_PER_THREAD; i++) {
    // Generate predictable data based on append number
    uint8_t data[APPEND_SIZE];
    for (int j = 0; j < APPEND_SIZE; j++) {
      data[j] = static_cast<uint8_t>((i * APPEND_SIZE + j) % 256);
    }

    // Append to buffer (single append thread ensures serialization)
    payload->buffer.append(APPEND_SIZE, data);

    // Small delay to allow concurrent reads
    if (i % 50 == 0) {
      emscripten_sleep(1);
    }
  }

  emscripten_atomic_add_u32(&payload->appendCount, NUM_APPENDS_PER_THREAD);
  emscripten_atomic_add_u32(&payload->sema, 1);
  emscripten_atomic_notify(&payload->sema, EMSCRIPTEN_NOTIFY_ALL_WAITERS);
}

static void read_worker_start(int payloadPtr) {
  auto payload = reinterpret_cast<AppendOnlyBufferWorkerPayload *>(payloadPtr);

  std::mt19937 gen(123); // Fixed seed for reproducible reads
  std::uniform_int_distribution<> offsetDis(0, 100);
  std::uniform_int_distribution<> sizeDis(1, 64);

  // Keep reading while appends are happening (single reader thread ensures
  // serialization)
  while (emscripten_atomic_load_u32(&payload->appendCount) <
         NUM_APPEND_THREADS * NUM_APPENDS_PER_THREAD) {
    size_t bufferSize = payload->buffer.size();
    if (bufferSize > 0) {
      size_t offset = offsetDis(gen) % bufferSize;
      size_t readSize =
          std::min(static_cast<size_t>(sizeDis(gen)), bufferSize - offset);

      uint8_t data[128];
      size_t bytesRead = payload->buffer.read(offset, readSize, data);

      if (bytesRead > 0) {
        emscripten_atomic_add_u32(&payload->readCount, 1);
      }
    }

    emscripten_sleep(2);
  }

  emscripten_atomic_add_u32(&payload->sema, 1);
  emscripten_atomic_notify(&payload->sema, EMSCRIPTEN_NOTIFY_ALL_WAITERS);
}

TEST_CASE("AppendOnlyBuffer concurrent", "[.][concurrent][noasan]") {
#if defined(__has_feature)
#if __has_feature(address_sanitizer)
  assert(false);
#endif
#endif

  AppendOnlyBufferWorkerPayload payload;

  // Create worker threads: one append thread and one read thread
  // This tests the correct concurrency model where:
  // - All appends are serialized (single append thread)
  // - All reads are serialized (single read thread)
  // - But appends and reads can happen concurrently with each other
  std::vector<emscripten_wasm_worker_t> appendWorkers;
  std::vector<emscripten_wasm_worker_t> readWorkers;

  appendWorkers.reserve(NUM_APPEND_THREADS);
  readWorkers.reserve(NUM_READ_THREADS);

  for (int i = 0; i < NUM_APPEND_THREADS; i++) {
    appendWorkers.push_back(emscripten_malloc_wasm_worker(3670016));
  }

  for (int i = 0; i < NUM_READ_THREADS; i++) {
    readWorkers.push_back(emscripten_malloc_wasm_worker(3670016));
  }

  fakeRealtimeThread = readWorkers[0];

  // Start both workers concurrently
  payload.sema = 0;

  for (int i = 0; i < NUM_APPEND_THREADS; i++) {
    emscripten_wasm_worker_post_function_vi(
        appendWorkers[i], append_worker_start, reinterpret_cast<int>(&payload));
  }

  for (int i = 0; i < NUM_READ_THREADS; i++) {
    emscripten_wasm_worker_post_function_vi(readWorkers[i], read_worker_start,
                                            reinterpret_cast<int>(&payload));
  }

  // Wait for both workers to complete
  const int totalWorkers = NUM_APPEND_THREADS + NUM_READ_THREADS;
  for (;;) {
    auto v = emscripten_atomic_load_u32(&payload.sema);
    if (v == totalWorkers) {
      break;
    }
    emscripten_atomic_wait_u32(&payload.sema, v,
                               ATOMICS_WAIT_DURATION_INFINITE);
  }

  // Verify results
  REQUIRE(payload.appendCount == NUM_APPEND_THREADS * NUM_APPENDS_PER_THREAD);
  REQUIRE(payload.buffer.size() ==
          NUM_APPEND_THREADS * NUM_APPENDS_PER_THREAD * APPEND_SIZE);
  REQUIRE(payload.readCount > 0); // Should have performed concurrent reads

  // Verify we can read the entire buffer
  size_t totalSize = payload.buffer.size();
  std::vector<uint8_t> finalData(totalSize);
  REQUIRE(payload.buffer.read(0, totalSize, finalData.data()) == totalSize);

  // Clean up workers
  for (auto worker : appendWorkers) {
    emscripten_terminate_wasm_worker(worker);
  }
  for (auto worker : readWorkers) {
    emscripten_terminate_wasm_worker(worker);
  }
}