#include <unistd.h>
#include <fstream>
#include <iostream>
#include <string>

#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <errno.h>

#include <signal.h>
#include <stdio.h>
#include <string.h>
#include <sys/time.h>
#include <unistd.h>

#include "bcc_version.h"
#include "BPF.h"

/* name of test */
#define TESTNAME "setitimer"
/* kernel function that runs on the isolated CPU */
#define PROBEFUNC "nohz_full_kick_func"

#define handle_error_en(en, msg) \
        do { errno = en; perror(msg); exit(0); } while (0)

void setaffinity(pthread_t thread, int isolcpu)
{
    int s,j;
    cpu_set_t cpuset;

    CPU_ZERO(&cpuset);
    CPU_SET(isolcpu, &cpuset);

    s = pthread_setaffinity_np(thread, sizeof(cpu_set_t), &cpuset);
    if (s != 0)
        handle_error_en(s, "pthread_setaffinity_np");
}

static int count = 0;

void timer_handler (int signum)
{
	if ((++count % 10) == 0) {
		printf ("timer expired %d times\n", count);
	}
}

void *busythread(void *ptr)
{
	int isolcpu = (long int) ptr;

	setaffinity(pthread_self(), isolcpu);
      
	/* Busy spin */
	while (1);

	return ptr;
}

struct itimerval timer;
pthread_t thread1;

/* setup_test runs before attaching the BPF probe */
int setup_test(long int isolatedcpu)
{
    struct sigaction sa;
    // create the thread objs
    // start the threads
    pthread_create(&thread1, NULL, *busythread, (void *)isolatedcpu);
    sleep(5);

    /* Install timer_handler as the signal handler for SIGVTALRM. */
     memset (&sa, 0, sizeof (sa));
     sa.sa_handler = &timer_handler;
     sigaction (SIGVTALRM, &sa, NULL);

     /* Configure the timer to expire after 250 msec... */
     timer.it_value.tv_sec = 0;
     timer.it_value.tv_usec = 250000;
     /* ... and every 250 msec after that. */
     timer.it_interval.tv_sec = 0;
     timer.it_interval.tv_usec = 250000;

     return 0;
}

/* start_test runs after attaching the BPF probe */
void start_test(void)
{
     /*
      * ITIMER_VIRTUAL This timer counts down against the user-mode CPU time
      * consumed by the process. (The measurement includes CPU time consumed
      * by all threads in the process). At each expiration, a SIGVTALRM
      * signal is generated.
      */
     setitimer (ITIMER_VIRTUAL, &timer, NULL);
}

/* BPF tracing */

const std::string BPF_PROGRAM = R"(
#include <uapi/linux/ptrace.h>
#include <linux/sched.h>    // for TASK_COMM_LEN

BPF_ARRAY(count, u32, 2);

int test_probe_fn(struct pt_regs *ctx) {

	int key = 0, cpu;
	u32 *isolcpu, *val;

	isolcpu = count.lookup(&key);
	if (!isolcpu) {
		return 0;
	}

	key = 1;
	val = count.lookup(&key);
	if (!val) {
		return 0;
	}
	cpu = bpf_get_smp_processor_id();

	if (cpu == *isolcpu) {
		(*val)++;
	}

	return 0;
}
)";

int main(int argc, char *argv[])
{
	long int isolcpu;
	ebpf::BPF bpf;
	auto init_res = bpf.init(BPF_PROGRAM);
	int success = 0;

	if (init_res.code() != 0) {
		std::cerr << init_res.msg() << std::endl;
		return 1;
	}

	if (argc < 2) {
		printf("%s isolated-cpu\n", argv[0]);
		return 1;
	}

        isolcpu = strtoul(argv[1], NULL, 10);

	if (setup_test(isolcpu))
		return 1;

	auto attach_res = bpf.attach_kprobe(PROBEFUNC, "test_probe_fn");

	if (attach_res.code() != 0) {
		std::cerr << attach_res.msg() << std::endl;
		return 1;
	}

	auto count_table = bpf.get_array_table<uint32_t>("count");
	auto update_res = count_table.update_value(0, isolcpu);
	if (update_res.code() != 0) {
 		std::cerr << update_res.msg() << std::endl;
		return 1;
	}

	printf("Starting %s test (isolated CPU %d)\n", TESTNAME, isolcpu);
	printf("Testing with thread pinned to isolated CPU\n");
	start_test();
	sleep(5);

	uint32_t key = 1, val = 0;
	attach_res = count_table.get_value(key, val);

	if (attach_res.code() != 0) {
		std::cerr << attach_res.msg() << std::endl;
		return 1;
	}

	printf("Kernel function %s called %d times on isolated CPU %d, expecting > 0.\n",
		PROBEFUNC, val, isolcpu);
	if (val > 0) {
		success += 1;
	}

	printf("Testing with thread not pinned to isolated CPU\n");
	setaffinity(thread1, 0);

	/* Zero the counter after moving thread to CPU 0 */
	update_res = count_table.update_value(1, 0);
	if (update_res.code() != 0) {
 		std::cerr << update_res.msg() << std::endl;
		return 1;
	}

	sleep(5);

	attach_res = count_table.get_value(key, val);
	if (attach_res.code() != 0) {
		std::cerr << attach_res.msg() << std::endl;
		return 1;
	}

	printf("Kernel function %s called %d times on isolated CPU %d, expecting 0.\n",
		PROBEFUNC, val, isolcpu);
	if (val == 0) {
		success += 1;
	}

	return success == 2 ? 0 : 1;
}