#include <spslr.h>

#include <dlfcn.h>
#include <stdio.h>
#include <stdlib.h>

#include "task_struct.h"
#include "export.h"

int second_pid();
const char* second_comm();
int subsecond_pid();
const char* subsecond_comm();

struct task_struct global = { .pid = 42, .comm = "main_global", .arrfun = {
	{ .a = 1, .b = 2, .c = 3, .d = 4.0 },
	{ .a = 5, .b = 6, .c = 7, .d = 8.0 },
	{ .a = 9, .b = 10, .c = 11, .d = 12.0 },
	{ .a = 13, .b = 14, .c = 15, .d = 16.0 },
	{ .a = 17, .b = 18, .c = 19, .d = 20.0 }
}, .tasks = LIST_HEAD_SELF(global.tasks) };

EXPORT_SYMBOL(global);

static int fetch_module_spslr_symbols(void* handle, struct spslr_module* mod) {
	if (!handle || !mod)
		return -1;

	mod->ipin_cnt = dlsym(handle, SPSLR_MODULE_SYM_IPIN_CNT);
	if (!mod->ipin_cnt)
		return -1;

	mod->ipins = dlsym(handle, SPSLR_MODULE_SYM_IPINS);
	if (!mod->ipins)
		return -1;

	mod->ipin_op_cnt = dlsym(handle, SPSLR_MODULE_SYM_IPIN_OP_CNT);
	if (!mod->ipin_op_cnt)
		return -1;

	mod->ipin_ops = dlsym(handle, SPSLR_MODULE_SYM_IPIN_OPS);
	if (!mod->ipin_ops)
		return -1;

	mod->dpin_cnt = dlsym(handle, SPSLR_MODULE_SYM_DPIN_CNT);
	if (!mod->dpin_cnt)
		return -1;

	mod->dpins = dlsym(handle, SPSLR_MODULE_SYM_DPINS);
	if (!mod->dpins)
		return -1;

	return 0;
}

static int do_module_test_access_pid(const char *path, const struct task_struct *t) {
	typedef int (*module_test_access_fn)(const struct task_struct *t);

	void *handle = dlopen(path, RTLD_NOW | RTLD_LOCAL);
	if (!handle) {
		fprintf(stderr, "dlopen failed for %s: %s\n", path, dlerror());
		return -1;
	}

	dlerror();

	struct spslr_module mod;
	if (fetch_module_spslr_symbols(handle, &mod) < 0) {
		fprintf(stderr, "failed to fetch spslr symbols in test module\n");
		dlclose(handle);
		return -1;
	}

	spslr_patch_module(&mod);

	module_test_access_fn fn = (module_test_access_fn)dlsym(handle, "module_test_access_pid");

	const char *err = dlerror();
	if (err) {
		fprintf(stderr, "dlsym failed: %s\n", err);
		dlclose(handle);
		return -1;
	}

	int ret = fn(t);

	dlclose(handle);
	return ret;
}

static void print_layout() {
    // TODO -> Make builtin __spslr_initial_offsetof(type, field) that is not patched
    printf("Current task_struct layout:\n");
    printf("    stuck0 (int)                 : %2llu -> %2llu\n", 0, offsetof(struct task_struct, stuck0));
    printf("    pid (int)                    : %2llu -> %2llu\n", 8, offsetof(struct task_struct, pid));
    printf("    comm (const char*)           : %2llu -> %2llu\n", 16, offsetof(struct task_struct, comm));
    printf("    arrfun (struct ArrayFun [5]) : %2llu -> %2llu\n", 24, offsetof(struct task_struct, arrfun));
    printf("    tasks (struct list_head)     : %2llu -> %2llu\n", 144, offsetof(struct task_struct, tasks));
    printf("    stuck1 (int)                 : %2llu -> %2llu\n", 160, offsetof(struct task_struct, stuck1));
    printf("    stuck2 (int)                 : %2llu -> %2llu\n", 166, offsetof(struct task_struct, stuck2));

    printf("Current ArrayFun layout:\n");
    printf("    a (int)    : %2llu -> %2llu\n", 0, offsetof(struct ArrayFun, a));
    printf("    b (int)    : %2llu -> %2llu\n", 4, offsetof(struct ArrayFun, b));
    printf("    c (int)    : %2llu -> %2llu\n", 8, offsetof(struct ArrayFun, c));
    printf("    d (double) : %2llu -> %2llu\n", 16, offsetof(struct ArrayFun, d));
}

void arr_test() {
	printf("Global arrfun[3].b: %d (should be 14)\n", global.arrfun[3].b);
}

int main(void)
{
    spslr_init();
    spslr_selfpatch();

    print_layout();

    struct list_head task_list;
    INIT_LIST_HEAD(&task_list);

    struct task_struct t1 = { .pid = 1, .comm = "init" };
    struct task_struct t2 = { .pid = 2, .comm = "kthreadd" };
    struct task_struct t3 = { .pid = 3, .comm = "worker" };

    INIT_LIST_HEAD(&t1.tasks);
    INIT_LIST_HEAD(&t2.tasks);
    INIT_LIST_HEAD(&t3.tasks);

    list_add_tail(&t1.tasks, &task_list);
    list_add_tail(&t2.tasks, &task_list);
    list_add_tail(&t3.tasks, &task_list);

    printf("Task list:\n");

    struct list_head *pos;
    list_for_each(pos, &task_list) {
        struct task_struct *task = list_entry(pos, struct task_struct, tasks);
        printf("  pid=%d, comm=%s\n", task->pid, task->comm);
    }

    size_t myOffset = ((size_t)&((struct task_struct*)0)->pid);
    printf("DIY offsetof(task_struct, pid) yiels %2llu\n", myOffset);

    printf("Global: pid=%d comm=\"%s\"\n", global.pid, global.comm);
    printf("Second global: pid=%d comm=\"%s\"\n", second_pid(), second_comm());
    printf("Subsecond global: pid=%d comm=\"%s\"\n", subsecond_pid(), subsecond_comm());

    arr_test();

    int module_pid = do_module_test_access_pid("./subject/libspslr_module.so", &global);
    printf("module_test_access returned %d (should be 42)\n", module_pid);

    spslr_cleanup();
    return 0;
}