#include <sys/types.h>
#include <dirent.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <libgen.h>

#include "parse.h"
#include "pagemap.h"

static struct maps_list *alloc_maplist(void)
{
	struct maps_list *map;

	map = malloc(sizeof *map);
	if (map == NULL)
		goto err;

	memset(map, 0, sizeof(*map));
	INIT_LIST_HEAD(&map->list);
err:
	return map;
}

static struct maps *alloc_map(void)
{
	struct maps *map;

	map = malloc(sizeof *map);
	if (map == NULL)
		goto err;

	memset(map, 0, sizeof(*map));
	INIT_LIST_HEAD(&map->list);
err:
	return map;
}

static struct maps *parse_maps(FILE *file, int pid, int tid)
{
	struct maps *the_map = NULL;
	char line[1024];
	int ret;

	while (fgets(line, sizeof(line), file)) {
		struct maps *map = alloc_map();
		unsigned long start, end;
		char name[1024];

		if (map == NULL)
			return 0;

		if (the_map == NULL)
			the_map = map;

		ret = sscanf(line, "%lx-%lx %*s %*s %*s %*s %s",
			     &start, &end, name);

		if (ret < 2) {
			printf("Error reading input: %s\n", line);
			break;
		}

		map->start = start;
		map->end = end;
		map->size = end - start;
		map->pid = pid;
		map->tid = tid;

		if (ret >= 3)
			strncpy(map->name, name, sizeof(map->name));

		list_add_tail(&map->list, &the_map->list);
	}

	return the_map;
}

static void clear_pageframe(struct pageframe *pf)
{
	memset(pf, 0, sizeof(*pf));
}

static struct pageframe *alloc_pageframe(void)
{
	struct pageframe *pageframe;

	pageframe = malloc(sizeof *pageframe);
	if (pageframe == NULL)
		goto err;

	clear_pageframe(pageframe);
err:
	return pageframe;
}

#define BITRANGE(first, last) (((2ll << (last - first)) - 1) << first)

static void pageframe_to_struct(unsigned long long p, struct pageframe *pf)
{
	/* Refer Documentation/vm/pagemap.txt for the format */
	pf->page_present = !!(BITRANGE(63, 63) & p);
	pf->page_swapped = !!(BITRANGE(62, 62) & p);
	pf->page_shift	 =   (BITRANGE(55, 60) & p) >> 55;
	pf->pfn		 =   (BITRANGE(0, 54) & p);
	pf->swap_type	 =   (BITRANGE(0, 4) & p);
	pf->swap_offset	 =   (BITRANGE(5, 54) & p) >> 5;
#if 0
	printf("pfn: %lx shift: %d present: %d swapped %d\n",
		pf->pfn, pf->page_shift, pf->page_present, pf->page_swapped);
#endif
}

static int compare_pageframe(struct bintree *at, struct bintree *bt)
{
	struct pageframe *a, *b;
	a = tree_to_pageframe(at);
	b = tree_to_pageframe(bt);

	return a->pfn - b->pfn;
}

struct bintree_ops pageframe_ops = {
	.compare = compare_pageframe,
};

static int read_cmdline(int pid, int tid, char *cmdline, size_t len)
{
	FILE *file;
	char path[512];
	int ret;

	snprintf(path, sizeof(path), "/proc/%d/task/%d/cmdline", pid, tid);
	file = fopen(path, "rb");

	if (!file)
		return -1;

	ret = fread(cmdline, 1, len, file);
	if (ret > 0)
		cmdline[ret - 1] = 0;
	fclose(file);

	return ret > 0 ? 0 : -1;
}

static char *get_name_by_pid(int pid)
{
	static int last_pid;
	static char cmdline[128];
	static char *bname;

	if (last_pid == pid)
		return bname;

	if (read_cmdline(pid, pid, cmdline, sizeof(cmdline))) {
		bname = NULL;
		return NULL;
	}

	bname = basename(cmdline);

	last_pid = pid;
	return bname;
}

static int check_parse_opts(struct parse_opts *opts, struct pageframe *pf,
		struct maps *map)
{
	if (opts->parse_mask & PARSE_PID) {
		if (opts->pid == map->pid)
			return 1;
	}

	if (opts->parse_mask & PARSE_MAP_NAME) {
		if (!strcmp(opts->name, map->name))
			return 1;
	}

	if (opts->parse_mask & PARSE_PROCESS_NAME) {
		if (!strcmp(opts->name, get_name_by_pid(map->pid)))
			return 1;
	}

	return 0;
}

/* Read data from the /proc/pid/pagemap file */
static int parse_pageframe(FILE *file, struct pageframe *pf_tree,
			struct maps *maps, struct parse_opts *opts)
{
	struct maps *map;
	struct maps_list *tmp;
	struct pageframe *match, *pageframe = NULL;
	long start, len, i;
	unsigned long long pf[10240];
	int ret, error;

	if (maps == NULL)
		return 0;

	/* Go through the list of allocated memory areas */
	list_for_each_entry(map, &maps->list, list) {
		start = map->start >> (PAGE_SHIFT - 3);
		len = map->size >> (PAGE_SHIFT);

		ret = fseek(file, start, SEEK_SET);
		if (ret) {
			error = errno;
			fprintf(stderr, "Error seeking to %lx: %s\n", start,
				strerror(error));
			continue;
		}

		for (i = 0; i < len; i++) {
			if (!ret) {
				ret = fread(&pf, 1,
					MIN(sizeof(pf), (len - i) * 8), file);
			}
			if (ret < 0) {
				error = errno;
				continue;
			}
			if (!pageframe)
				pageframe = alloc_pageframe();
			ret -= sizeof(pf[0]);

			/* ignore unused pages */
			if (!pf[ret / sizeof(pf[0])])
				continue;

			pageframe_to_struct(pf[ret / sizeof(pf[0])], pageframe);

			/* ignore unused pages */
			if (!(pageframe->page_swapped ||
					pageframe->page_present))
				continue;

			if (check_parse_opts(opts, pageframe, map)) {
				match = tree_to_pageframe(
					bintree_add(&pf_tree->tree,
						&pageframe->tree,
						&pageframe_ops));
			} else {
				match = tree_to_pageframe(
					bintree_find(&pf_tree->tree,
						&pageframe->tree,
						&pageframe_ops));
			}

			if (match == NULL)
				continue;

			if (match == pageframe)
				pageframe = NULL;

			match->refcount++;
			/*
			 * Add a link from the physical page to this
			 * process's page map
			 */
			if (!match->ml) {
				match->ml = alloc_maplist();
				match->ml->map = map;
			} else {
				tmp = alloc_maplist();
				tmp->map = map;
				list_add(&tmp->list, &match->ml->list);
			}

			if (match->page_present)
				map->pages_present++;
			else if (match->page_swapped)
				map->pages_swapped++;
		}
	}

	return 0;
}

static int read_pageframe(int pid, int tid, struct pageframe *pageframe,
			struct process **process_list, struct parse_opts *opts)
{
	struct maps *maps;
	struct process *process;
	FILE *file;
	char path[512];

	process = malloc(sizeof(*process));
	memset(process, 0, sizeof(*process));
	INIT_LIST_HEAD(&process->list);

	if (*process_list == NULL)
		*process_list = process;

	process->pid = pid;
	process->tid = tid;

	list_add_tail(&process->list, &(*process_list)->list);

	snprintf(path, sizeof(path), "/proc/%d/task/%d/maps", pid, tid);
	file = fopen(path, "rb");

	if (!file)
		return 0;

	maps = parse_maps(file, pid, tid);
	fclose(file);
	process->maps = maps;

	snprintf(path, sizeof(path), "/proc/%d/task/%d/pagemap", pid, tid);
	file = fopen(path, "rb");

	if (!file)
		return 0;

	parse_pageframe(file, pageframe, maps, opts);
	fclose(file);

	if (read_cmdline(pid, tid, process->name, sizeof(process->name)))
		return 1;

	if (maps == NULL)
		return 1;

	list_for_each_entry(maps, &process->maps->list, list) {
		process->pages_present += maps->pages_present;
		process->pages_swapped += maps->pages_swapped;
	}

	return 1;
}

static int parse_pid(DIR **dir)
{
	struct dirent *dirent;
	int error;

restart:
	dirent = readdir(*dir);
	if (!dirent) {
		if (errno == 0) {
			closedir(*dir);
			*dir = NULL;
			return 0;
		}
		error = errno;
		printf("Failed to read /proc directory: %s\n", strerror(error));
		return -1;
	}

	if (dirent->d_name[0] < '0' || dirent->d_name[0] > '9')
		goto restart;

	return atoi(dirent->d_name);
}

static int opendir_check(DIR **dir, const char *path)
{
	int error;

	if (!*dir) {
		*dir = opendir(path);
		if (!dir) {
			error = errno;
			fprintf(stderr, "Failed to open %s directory: %s\n",
				path, strerror(error));
			return -1;
		}
	}

	return 0;
}

static int get_next_tid(int pid, DIR **dir)
{
	if (*dir == NULL) {
		char path[64];

		snprintf(path, sizeof(path), "/proc/%d/task/", pid);
		if (opendir_check(dir, path))
			return -1;
	}

	return parse_pid(dir);
}

static int get_next_pid(DIR **dir)
{
	if (opendir_check(dir, "/proc"))
		return -1;

	return parse_pid(dir);
}

static int get_next_pid_by_name(DIR **dir, char *name)
{
	int pid;
	char *pname;

	if (opendir_check(dir, "/proc"))
		return -1;

	while (1) {
		pid = parse_pid(dir);
		if (pid <= 0)
			break;

		pname = get_name_by_pid(pid);
		if (pname == NULL)
			continue;
		if (strcmp(pname, name))
			continue;

		return pid;
	}

	return 0;
}

static int read_pageframe_with_threads(int pid,
				struct pageframe *pageframe,
				struct process **process_list,
				struct parse_opts *opts)
{
	DIR *dir = NULL;
	int tid;
	int count = 0;

	while (1) {
		if (opts->with_threads)
			tid = get_next_tid(pid, &dir);
		else
			tid = pid;

		if (tid <= 0)
			return count;

		count += read_pageframe(pid, tid, pageframe, process_list,
					opts);

		if (!opts->with_threads)
			break;
	}

	return count;
}

int scan_all_pids(struct pageframe *pf, struct process **process_list,
		struct parse_opts *opts)
{
	DIR *dir = NULL;
	int pid;
	int count = 0;

	if (opts->parse_mask & PARSE_PROCESS_NAME) {
		while ((pid = get_next_pid_by_name(&dir, opts->name))) {
			count += read_pageframe_with_threads(pid, pf,
							process_list,
							opts);
		}
		dir = NULL;
	}

	if (opts->parse_mask & PARSE_PID)
		count = read_pageframe_with_threads(opts->pid, pf, process_list,
						opts);

	if ((count == 0) && !(opts->parse_mask & PARSE_MAP_NAME)) {
		printf("Failed to find any matching processes "
			"with given arguments\n");
		return -1;
	}

	if (opts->parse_mask & PARSE_DUMP)
		return 0;

	while (1) {
		pid = get_next_pid(&dir);
		if (pid <= 0)
			break;
		read_pageframe_with_threads(pid, pf, process_list, opts);
	}

	return 0;
}