#define _GNU_SOURCE
#define _FILE_OFFSET_BITS 64

#include <fcntl.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <time.h>
#include <unistd.h>

#define EPOCHS 30
#define ITERATIONS_PER_EPOCH 2000
#define BLOCK_SIZE 2048
#define PAGE_SIZE 4096

#define CHECK(cond)                                                            \
  do {                                                                         \
    if (!(cond)) {                                                             \
      fprintf(stderr, "CHECK failed at line %d: %s\n", __LINE__, #cond);       \
      perror("Error");                                                         \
      exit(1);                                                                 \
    }                                                                          \
  } while (0)

int main(int argc, char *argv[]) {
  int mode = -1;

  struct {
    const char *name;
    const char *desc;
  } modes[] = {{"pread", "pread system calls"},
               {"mmap", "memory mapped access"},
               {"prefetch", "mmap + pread prefetch"},
               {"optprefetch", "mmap + optimized pread prefetch"},
               {"fadvise", "mmap + posix_fadvise WILLNEED"},
               {"readahead", "mmap + readahead"},
               {"madvise", "mmap + madvise WILLNEED"}};

  if (argc != 3) {
    fprintf(stderr, "Usage: %s <mode> <file>\nModes:\n", argv[0]);
    for (int i = 0; i < 7; i++) {
      fprintf(stderr, "  %s - %s\n", modes[i].name, modes[i].desc);
    }
    exit(1);
  }

  for (int i = 0; i < 7; i++) {
    if (strcmp(argv[1], modes[i].name) == 0) {
      mode = i;
      break;
    }
  }

  CHECK(mode >= 0);
  printf("Using %s for random access.\n", modes[mode].desc);

  int fd = open(argv[2], O_RDONLY);
  CHECK(fd != -1);
  CHECK(posix_fadvise(fd, 0, 0, POSIX_FADV_DONTNEED) == 0);

  struct stat st;
  CHECK(fstat(fd, &st) != -1);

  size_t file_size = st.st_size;
  CHECK(file_size >= BLOCK_SIZE);

  void *mapped = mmap(NULL, file_size, PROT_READ, MAP_PRIVATE, fd, 0);
  CHECK(mapped != MAP_FAILED);

  void *buffer = malloc(BLOCK_SIZE);
  CHECK(buffer);

  CHECK(madvise(mapped, file_size, MADV_DONTNEED) == 0);
  CHECK(madvise(mapped, file_size, MADV_RANDOM) == 0);

  srand(time(NULL));

  double epoch_times[EPOCHS];

  for (int epoch = 0; epoch < EPOCHS; epoch++) {
    struct timespec start, end;
    CHECK(clock_gettime(CLOCK_MONOTONIC, &start) == 0);

    for (int i = 0; i < ITERATIONS_PER_EPOCH; i++) {
      size_t max_offset = file_size - BLOCK_SIZE;
      size_t offset = rand() % max_offset;

      uint64_t *data;
      switch (mode) {
      case 0: // pread
        CHECK(pread(fd, buffer, BLOCK_SIZE, offset) == BLOCK_SIZE);
        data = (uint64_t *)buffer;
        break;
      case 1: // mmap
        data = (uint64_t *)((char *)mapped + offset);
        break;
      case 2: // prefetch
        CHECK(pread(fd, buffer, BLOCK_SIZE, offset) == BLOCK_SIZE);
        data = (uint64_t *)((char *)mapped + offset);
        break;
      case 3: { // optprefetch
        size_t range_end = offset + BLOCK_SIZE - 1;
        size_t first_page = offset & ~(PAGE_SIZE - 1);
        size_t last_page = range_end & ~(PAGE_SIZE - 1);

        if (first_page != last_page) {
          size_t prefetch_start = first_page + PAGE_SIZE - 1;
          size_t prefetch_end = last_page + 1;
          size_t prefetch_size = prefetch_end - prefetch_start;
          CHECK(pread(fd, buffer, prefetch_size, prefetch_start) ==
                prefetch_size);
        }

        data = (uint64_t *)((char *)mapped + offset);
        break;
      }
      case 4: // fadvise
        CHECK(posix_fadvise(fd, offset, BLOCK_SIZE, POSIX_FADV_WILLNEED) == 0);
        data = (uint64_t *)((char *)mapped + offset);
        break;
      case 5: // readahead
        CHECK(readahead(fd, offset, BLOCK_SIZE) == 0);
        data = (uint64_t *)((char *)mapped + offset);
        break;
      case 6: { // madvise
        size_t page_start = offset & ~(PAGE_SIZE - 1);
        size_t page_end =
            (offset + BLOCK_SIZE + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
        CHECK(madvise((char *)mapped + page_start, page_end - page_start,
                      MADV_WILLNEED) == 0);
        data = (uint64_t *)((char *)mapped + offset);
        break;
      }
      }

      uint64_t sum = 0;
      for (int j = 0; j < BLOCK_SIZE / sizeof(uint64_t); j++)
        sum += data[j];

      __asm__ volatile("" : : "r"(sum) : "memory");
    }

    CHECK(clock_gettime(CLOCK_MONOTONIC, &end) == 0);

    double elapsed =
        (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9;
    epoch_times[epoch] = (elapsed / ITERATIONS_PER_EPOCH) * 1e6;
  }

  int compare_double(const void *a, const void *b) {
    double da = *(const double *)a;
    double db = *(const double *)b;
    return (da > db) - (da < db);
  }

  qsort(epoch_times, EPOCHS, sizeof(double), compare_double);

  close(fd);
  printf("Median time per random access: %.2f microseconds\n",
         epoch_times[EPOCHS / 2]);
}