#include <asm/unistd.h>
#include <immintrin.h>
#include <linux/hw_breakpoint.h>
#include <linux/perf_event.h>
#include <sched.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/syscall.h>
#include <unistd.h>

#include <cmath>

// --------------------------------------------------------------------
// Pin to CPU

static void pin_to_cpu(uint32_t cpu) {
  cpu_set_t cpu_mask;
  CPU_ZERO(&cpu_mask);
  CPU_SET(0, &cpu_mask);
  if (sched_setaffinity(cpu, sizeof(cpu_mask), &cpu_mask) != 0) {
    fprintf(stderr, "Could not set CPU affinity\n");
    exit(EXIT_FAILURE);
  }
}

// --------------------------------------------------------------------
// perf instrumentation -- a mixture of man 3 perf_event_open and
// <https://stackoverflow.com/a/42092180>

static long perf_event_open(struct perf_event_attr *hw_event, pid_t pid,
                            int cpu, int group_fd, unsigned long flags) {
  int ret;
  ret = syscall(__NR_perf_event_open, hw_event, pid, cpu, group_fd, flags);
  return ret;
}

static void setup_perf_event(struct perf_event_attr *evt, int *fd, uint64_t *id,
                             uint32_t evt_type, uint64_t evt_config,
                             int group_fd) {
  memset(evt, 0, sizeof(struct perf_event_attr));
  evt->type = evt_type;
  evt->size = sizeof(struct perf_event_attr);
  evt->config = evt_config;
  evt->disabled = 1;
  evt->exclude_kernel = 1;
  evt->exclude_hv = 1;
  evt->read_format = PERF_FORMAT_GROUP | PERF_FORMAT_ID;

  *fd = perf_event_open(evt, 0, -1, group_fd, 0);
  if (*fd == -1) {
    fprintf(stderr, "Error opening leader %llx\n", evt->config);
    exit(EXIT_FAILURE);
  }

  ioctl(*fd, PERF_EVENT_IOC_ID, id);
}

static struct perf_event_attr perf_cycles_evt;
static int perf_cycles_fd;
static uint64_t perf_cycles_id;

static struct perf_event_attr perf_clock_evt;
static int perf_clock_fd;
static uint64_t perf_clock_id;

static struct perf_event_attr perf_instrs_evt;
static int perf_instrs_fd;
static uint64_t perf_instrs_id;

static struct perf_event_attr perf_cache_misses_evt;
static int perf_cache_misses_fd;
static uint64_t perf_cache_misses_id;

static struct perf_event_attr perf_cache_references_evt;
static int perf_cache_references_fd;
static uint64_t perf_cache_references_id;

static struct perf_event_attr perf_branch_misses_evt;
static int perf_branch_misses_fd;
static uint64_t perf_branch_misses_id;

static struct perf_event_attr perf_branch_instructions_evt;
static int perf_branch_instructions_fd;
static uint64_t perf_branch_instructions_id;

static void perf_init(void) {
  // Cycles
  setup_perf_event(&perf_cycles_evt, &perf_cycles_fd, &perf_cycles_id,
                   PERF_TYPE_HARDWARE, PERF_COUNT_HW_CPU_CYCLES, -1);
  // Clock
  setup_perf_event(&perf_clock_evt, &perf_clock_fd, &perf_clock_id,
                   PERF_TYPE_SOFTWARE, PERF_COUNT_SW_TASK_CLOCK,
                   perf_cycles_fd);
  // Instructions
  setup_perf_event(&perf_instrs_evt, &perf_instrs_fd, &perf_instrs_id,
                   PERF_TYPE_HARDWARE, PERF_COUNT_HW_INSTRUCTIONS,
                   perf_cycles_fd);
  // Cache misses
  setup_perf_event(&perf_cache_misses_evt, &perf_cache_misses_fd,
                   &perf_cache_misses_id, PERF_TYPE_HARDWARE,
                   PERF_COUNT_HW_CACHE_MISSES, perf_cycles_fd);
  // Cache references
  setup_perf_event(&perf_cache_references_evt, &perf_cache_references_fd,
                   &perf_cache_references_id, PERF_TYPE_HARDWARE,
                   PERF_COUNT_HW_CACHE_REFERENCES, perf_cycles_fd);
  // Branch misses
  setup_perf_event(&perf_branch_misses_evt, &perf_branch_misses_fd,
                   &perf_branch_misses_id, PERF_TYPE_HARDWARE,
                   PERF_COUNT_HW_BRANCH_MISSES, perf_cycles_fd);
  // Branch instructions
  setup_perf_event(&perf_branch_instructions_evt, &perf_branch_instructions_fd,
                   &perf_branch_instructions_id, PERF_TYPE_HARDWARE,
                   PERF_COUNT_HW_BRANCH_INSTRUCTIONS, perf_cycles_fd);
}

static void perf_close(void) {
  close(perf_clock_fd);
  close(perf_cycles_fd);
  close(perf_instrs_fd);
  close(perf_cache_misses_fd);
  close(perf_cache_references_fd);
}

static void disable_perf_count(void) {
  ioctl(perf_cycles_fd, PERF_EVENT_IOC_DISABLE, PERF_IOC_FLAG_GROUP);
}

static void enable_perf_count(void) {
  ioctl(perf_cycles_fd, PERF_EVENT_IOC_ENABLE, PERF_IOC_FLAG_GROUP);
}

static void reset_perf_count(void) {
  ioctl(perf_cycles_fd, PERF_EVENT_IOC_RESET, PERF_IOC_FLAG_GROUP);
}

struct perf_read_value {
  uint64_t value;
  uint64_t id;
};

struct perf_read_format {
  uint64_t nr;
  struct perf_read_value values[];
};

static char perf_read_buf[4096];

struct perf_count {
  uint64_t cycles;
  double seconds;
  uint64_t instructions;
  uint64_t cache_misses;
  uint64_t cache_references;
  uint64_t branch_misses;
  uint64_t branch_instructions;

  perf_count()
      : cycles(0),
        seconds(0.0),
        instructions(0),
        cache_misses(0),
        cache_references(0),
        branch_misses(0),
        branch_instructions(0) {}
};

static void read_perf_count(struct perf_count *count) {
  if (!read(perf_cycles_fd, perf_read_buf, sizeof(perf_read_buf))) {
    fprintf(stderr, "Could not read cycles from perf\n");
    exit(EXIT_FAILURE);
  }
  struct perf_read_format *rf = (struct perf_read_format *)perf_read_buf;
  if (rf->nr != 7) {
    fprintf(stderr, "Bad number of perf events\n");
    exit(EXIT_FAILURE);
  }
  for (int i = 0; i < static_cast<int>(rf->nr); i++) {
    struct perf_read_value *value = &rf->values[i];
    if (value->id == perf_cycles_id) {
      count->cycles = value->value;
    } else if (value->id == perf_clock_id) {
      count->seconds = ((double)(value->value / 1000ull)) / 1000000.0;
    } else if (value->id == perf_instrs_id) {
      count->instructions = value->value;
    } else if (value->id == perf_cache_misses_id) {
      count->cache_misses = value->value;
    } else if (value->id == perf_cache_references_id) {
      count->cache_references = value->value;
    } else if (value->id == perf_branch_misses_id) {
      count->branch_misses = value->value;
    } else if (value->id == perf_branch_instructions_id) {
      count->branch_instructions = value->value;
    } else {
      fprintf(stderr, "Spurious value in perf read (%ld)\n", value->id);
      exit(EXIT_FAILURE);
    }
  }
}

template <typename T>
perf_count get_perf_stats(T&& lambda) {
  struct perf_count counts;
  static auto once = []() {
    // init perf once in program
    perf_init();
    return 0;
  }();
  (void)once;

  disable_perf_count();
  reset_perf_count();
  enable_perf_count();

  lambda();

  disable_perf_count();
  read_perf_count(&counts);

  return counts;
}

void do_something() {
    // stuff
}

int main() {
  pin_to_cpu(0);

  perf_count counts = get_perf_stats([]() { do_something(); });

  printf("cycles= %lu\n", counts.cycles);
  printf("seconds= %f\n", counts.seconds);
  printf("counts.instructions= %lu\n", counts.instructions);
  printf("cache_misses= %lu\n", counts.cache_misses);
  printf("cache_references= %lu\n", counts.cache_references);
  printf("branch_misses= %lu\n", counts.branch_misses);
  printf("branch_instructions= %lu\n", counts.branch_instructions);
}