#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 "settimeofday"
/* kernel function that runs on the isolated CPU */
#define PROBEFUNC "retrigger_next_event"

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

	pthread_create(&thread1, NULL, *busythread, (void *)isolatedcpu);
	sleep(5);

	return 0;
}


/* start_test runs after attaching the BPF probe */
void start_test(void)
{
	int ret;
	struct timeval tv;

	ret = gettimeofday(&tv, NULL);
	if (ret == -1) {
		perror("gettimeofday");
		exit(1);
	}

	tv.tv_sec += 1;

	ret = settimeofday(&tv, NULL);
	if (ret == -1) {
		perror("settimeofday");
		exit(1);
	}
}

/* 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 == 1 ? 0 : 1;
}