/* main.c
 *
 * Copyright 2023 Ondřej Míchal
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 *
 * SPDX-License-Identifier: GPL-3.0-or-later
 */

#include <inttypes.h>
#include <linux/perf_event.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/syscall.h>
#include <sys/sysinfo.h>
#include <sys/timerfd.h>
#include <unistd.h>

#define COLLECT_SLEEP_NS 100010 // 9999Hz
//#define COLLECT_SLEEP_NS 10101010 // 99Hz
#define RAPL_EVENT_COUNT 5

typedef struct data_t {
  __u64 time_ns;
  __u64 rapl_cores;
  __u64 rapl_pkg;
  __u64 rapl_ram;
  __u64 rapl_gpu;
  __u64 rapl_psys;
} data_t;

struct perf_event_domain {
  char *name;
  char *path;
  char *events[RAPL_EVENT_COUNT];
};

struct perf_data {
  uint64_t nr;
  struct {
    uint64_t value;
    uint64_t id;
  } values[5];
};

static struct perf_event_domain rapl_domain = {
    "power",
    "/sys/devices/power",
    {
        "energy-cores",
        "energy-pkg",
        "energy-ram",
        "energy-gpu",
        "energy-psys",
    },
};

static int get_rapl_perf_event(struct perf_event_domain *event_domain,
                               char *event_name, struct perf_event_attr *attr) {
  FILE *f = NULL;
  char line_buf[BUFSIZ];
  char domain_type_path[BUFSIZ];
  char event_path[BUFSIZ];
  char *event_type_str, *event_config_str;
  int event_type, event_config;

  memset(attr, 0, sizeof(struct perf_event_attr));

  if (snprintf(domain_type_path, BUFSIZ, "%s/type", event_domain->path) < 0)
    return -1;

  if (snprintf(event_path, BUFSIZ, "%s/events/%s", event_domain->path,
               event_name) < 0)
    return -1;

  f = fopen(domain_type_path, "r");
  if (f == NULL)
    return -1;

  if (fgets(line_buf, BUFSIZ, f) == NULL)
    return -1;

  fclose(f);

  event_type_str = strtok(line_buf, "\n");
  event_type = strtol(event_type_str, NULL, 10);

  f = fopen(event_path, "r");
  if (f == NULL)
    return -1;

  memset(line_buf, 0, BUFSIZ);
  if (fgets(line_buf, BUFSIZ, f) == NULL)
    return -1;

  fclose(f);

  strtok(line_buf, "=");
  event_config_str = strtok(NULL, "\n");
  event_config = strtol(event_config_str, NULL, 16);

  attr->size = sizeof(struct perf_event_attr);
  attr->type = event_type;
  attr->config = event_config;
  attr->sample_type = PERF_SAMPLE_IDENTIFIER;
  attr->read_format = PERF_FORMAT_GROUP | PERF_FORMAT_ID;
  attr->disabled = 1;

  return 0;
}

int main(void) {
  int perf_event_fds[RAPL_EVENT_COUNT] = {-1, -1, -1, -1, -1};
  int perf_event_group_fd = -1;
  // We'll be sampling RAPL counters only on CPU 0 and then extrapolate the
  // values to the captured syscalls.
  const int cpu = 0;
  uint64_t exp;

  for (int i = 0; i < RAPL_EVENT_COUNT; i++) {
    struct perf_event_attr attr;
    char *curr_event = NULL;

    curr_event = rapl_domain.events[i];

    if (get_rapl_perf_event(&rapl_domain, curr_event, &attr))
      continue;

    perf_event_fds[i] =
        syscall(SYS_perf_event_open, &attr, -1, cpu, perf_event_group_fd, 0);
    if (perf_event_fds[i] == -1)
      continue;

    ioctl(perf_event_fds[i], PERF_EVENT_IOC_RESET, 0);

    if (i == 0)
      perf_event_group_fd = perf_event_fds[i];
  }

  int timer_fd = timerfd_create(CLOCK_MONOTONIC, 0);

  struct timespec now;
  clock_gettime(CLOCK_MONOTONIC, &now);

  struct itimerspec timer_spec;
  timer_spec.it_value.tv_sec = now.tv_sec;
  timer_spec.it_value.tv_nsec = now.tv_nsec;
  timer_spec.it_interval.tv_sec = 0;
  timer_spec.it_interval.tv_nsec = COLLECT_SLEEP_NS;

  timerfd_settime(timer_fd, TFD_TIMER_ABSTIME, &timer_spec, NULL);

  for (int i = 0; i < RAPL_EVENT_COUNT; i++) {
    if (perf_event_fds[i] == -1)
      continue;

    ioctl(perf_event_fds[i], PERF_EVENT_IOC_ENABLE, 0);
  }

  // Poll data from ring buffer providing data from the eBPF program.
  while (1) {
    struct perf_data data;

    read(timer_fd, &exp, sizeof(uint64_t));

    clock_gettime(CLOCK_MONOTONIC, &now);

    read(perf_event_group_fd, &data, sizeof(struct perf_data));

    printf("%lu%lu %lu %lu %lu %lu %lu\n",
            now.tv_sec, now.tv_nsec,
            data.values[0].value, data.values[1].value,
            data.values[2].value, data.values[3].value,
            data.values[4].value);
    }

  return EXIT_SUCCESS;
}
