/* 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 <bpf/bpf.h>
#include <bpf/libbpf.h>
#include <err.h>
#include <fcntl.h>
#include <linux/perf_event.h>
#include <signal.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/syscall.h>
#include <sys/sysinfo.h>
#include <time.h>
#include <unistd.h>

#include "main.h"
#include "skel.h"

#define SAMPLING_FREQ 9999

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

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

static bool poll = true;

static void handle_sigint(int signo, siginfo_t *info, void *ctx) { poll = false; }

static int handle_ring_buffer(void *ctx, void *data, size_t data_sz) {
  data_t *sysrapl_data = (data_t *)data;

  printf("%llu %llu %llu %llu %llu %llu\n",
            sysrapl_data->time_ns,
            sysrapl_data->rapl_cores, sysrapl_data->rapl_pkg,
            sysrapl_data->rapl_ram, sysrapl_data->rapl_gpu,
            sysrapl_data->rapl_psys);

  return 0;
}

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_TOTAL_TIME_ENABLED | PERF_FORMAT_TOTAL_TIME_RUNNING;
  attr->disabled = 1;

  return 0;
}

static int attach_perf_event_sampling(const struct bpf_program *prog,
                                      int frequency, struct bpf_link *link) {
  struct perf_event_attr attr;
  int fd = 0;

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

  attr.type = PERF_TYPE_SOFTWARE;
  attr.config = PERF_COUNT_SW_CPU_CLOCK;
  attr.freq = 1;
  attr.sample_period = frequency;

  fd = syscall(SYS_perf_event_open, &attr, -1, 0, -1, 0);
  if (fd < 0)
    return -1;

  link = bpf_program__attach_perf_event(prog, fd);
  if (link == NULL)
    return -1;

  return 0;
}

int main(void) {
  struct bpf_link *perf_event_link = NULL;
  struct bpf_map *rapl_event_maps[RAPL_EVENT_COUNT] = {NULL, NULL, NULL, NULL};
  struct ring_buffer *rb = NULL;
  struct bpf *skel = NULL;
  bool *switches[RAPL_EVENT_COUNT] = {NULL, NULL, NULL, NULL, NULL};
  int perf_event_fds[RAPL_EVENT_COUNT] = {-1, -1, -1, -1, -1};
  int perf_event_group_fd = -1;
  int err;
  // We'll be sampling RAPL counters only on CPU 0 and then extrapolate the
  // values to the captured syscalls.
  const int cpu = 0;

  struct sigaction act = { 0 };
  act.sa_sigaction = &handle_sigint;
  act.sa_flags = 0;
  sigaction(SIGINT, &act, NULL);

  // Start the bootstrapping of a eBPF program using the 4 loading steps
  // See <bootstrap url> for more info.
  skel = bpf__open();

  // Assign eBPF maps to the matching RAPL domain
  rapl_event_maps[0] = skel->maps.rapl_cores;
  switches[0] = &skel->rodata->has_rapl_cores;

  rapl_event_maps[1] = skel->maps.rapl_pkg;
  switches[1] = &skel->rodata->has_rapl_pkg;

  rapl_event_maps[2] = skel->maps.rapl_ram;
  switches[2] = &skel->rodata->has_rapl_ram;

  rapl_event_maps[3] = skel->maps.rapl_gpu;
  switches[3] = &skel->rodata->has_rapl_gpu;

  rapl_event_maps[4] = skel->maps.rapl_psys;
  switches[4] = &skel->rodata->has_rapl_psys;

  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];

    *switches[i] = true;
  }

  bpf__load(skel);

  rb = ring_buffer__new(bpf_map__fd(skel->maps.ringbuffer), handle_ring_buffer,
                        NULL, NULL);

  attach_perf_event_sampling(skel->progs.collect_rapl_info, SAMPLING_FREQ,
                                 perf_event_link);

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

    bpf_map__update_elem(rapl_event_maps[i], &cpu, sizeof(int),
                               &perf_event_fds[i], sizeof(int), BPF_ANY);
  }

  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);
  }

  bpf__attach(skel);

  // Poll data from ring buffer providing data from the eBPF program.
  while (poll) {
    err = ring_buffer__poll(rb, 100);
    if (err == -EINTR)
      break;
  }

  ring_buffer__free(rb);
  bpf__destroy(skel);

  return EXIT_SUCCESS;
}
