/*
 * Copyright (C) 2010 Timo Kokkonen <timo.t.kokkonen@iki.fi>
 *
 * 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 2 of the License.
 *
 * 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, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

#include <sys/types.h>
#include <sys/stat.h>
#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <unistd.h>

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

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

	map = calloc(sizeof *map, 1);
	if (map == NULL)
		goto err;

	INIT_LIST_HEAD(&map->list);
err:
	return map;
}

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

	map = calloc(sizeof(*map), 1);
	if (map == NULL)
		goto err;

	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;
		int skip;

		if (map == NULL)
			return 0;

		if (the_map == NULL)
			the_map = map;

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

		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;

		strncpy(map->name, line + skip, sizeof(map->name) - 1);

		/* zero out the newline */
		map->name[MAX(strlen(map->name) - 1, 0)] = '\0';

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

	return the_map;
}

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

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

	clear_pageframe(pageframe);
err:
	return pageframe;
}

static int pid_is_match(int pidn, struct list_head *pidlist)
{
	struct pidlist *pid;

	list_for_each_entry(pid, pidlist, list) {
		if (pid->pid == pidn)
			return 1;
	}

	return 0;
}

static int should_scan_process(struct parse_opts *opts, struct process *process)
{
	int match = 0;

	if (is_parse_option(opts, PARSE_PID) &&
		is_parse_option(opts, PARSE_MAP_NAME)) {
		if (pid_is_match(process->pid, &opts->pidlist))
			match = 1;
	} else {
		if (is_parse_option(opts, PARSE_PID)) {
			if (pid_is_match(process->pid, &opts->pidlist))
				match = 1;
		}

		if (is_parse_option(opts, PARSE_MAP_NAME))
			match = 1;
	}

	if (is_parse_option(opts, PARSE_NOADD_TREE))
		match = !match;

	return match;
}

static int should_scan_mapping(struct parse_opts *opts, struct maps *map)
{
	int match = 0;

	if (is_parse_option(opts, PARSE_MAP_NAME)) {
		if (strstr(map->name, opts->name))
			match = 1;

		if (is_parse_option(opts, PARSE_NOADD_TREE))
			match = !match;
	} else
		match = 1;

	return match;
}

static int should_add_to_tree(struct parse_opts *opts, struct pageframe *pf,
		struct maps *map)
{
	if (is_parse_option(opts, PARSE_NOADD_TREE))
		return 0;

	return 1;
}

/* Read data from the /proc/pid/pagemap file */
static int parse_pageframe(FILE *file, struct rb_root *root,
			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;

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

		if (!should_scan_mapping(opts, map))
			continue;

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

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

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

			pageframe->pf = (pf[ret / sizeof(pf[0])]);

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

			if (should_add_to_tree(opts, pageframe, map))
				match = pf_insert(root,	pageframe);
			else
				match = pf_search(root,	pageframe);

			if (match == NULL)
				continue;

			if (match == pageframe)
				pageframe = NULL;

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

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

	return 0;
}

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

	process = calloc(sizeof(*process), 1);
	INIT_LIST_HEAD(&process->list);

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

	if (!should_scan_process(opts, process))
		goto free;

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

	if (!file)
		goto free;

	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)
		goto free;

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

	if (read_cmdline(pid, tid, process->name, sizeof(process->name)))
		goto free;

	if (!is_parse_option(opts, PARSE_NOADD_TREE))
		process->is_initial_pid = 1;

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

	return 1;
free:
	free(process);

	/*
	 * Remove the pid from the list. It is no longer an
	 * interesting pid, since we can't access its data
	 */
	list_for_each_entry_safe(pidl, n, &opts->pidlist, list) {
		if (pidl->pid == pid) {
			list_del(&pidl->list);
			free(pidl);
			break;
		}
	}

	return 0;
}

static int read_pageframe_with_threads(int pid,
				struct rb_root *root,
				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, root, process_list, opts);

		if (!opts->with_threads)
			break;
	}

	return count;
}

int scan_all_pids(struct rb_root *root, struct process *process_list,
		struct parse_opts *opts)
{
	struct pidlist *pidlist, *n;
	DIR *dir = NULL;
	int pid;
	int count = 0;
	int len = 0, i;

	if (is_parse_option(opts, PARSE_PID)) {
		list_for_each_entry_safe(pidlist, n, &opts->pidlist, list) {
			count += read_pageframe_with_threads(pidlist->pid, root,
							process_list, opts);
		}
	}

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

	if (is_parse_option(opts, PARSE_DUMP))
		return 0;

	if (is_parse_option(opts, PARSE_MAP_NAME) &&
		!is_parse_option(opts, PARSE_PID)) {
		printf("Scanning page mappings for process: ");
		while (1) {
			pid = get_next_pid(&dir);
			if (pid <= 0)
				break;

			for (i = 0; i < len; i++)
				putchar('\b');
			len = printf("% 5d", pid);
			fflush(stdout);

			read_pageframe_with_threads(pid, root, process_list,
						opts);
		}

		for (i = 0; i < len; i++)
			putchar('\b');
		printf("Done  \n");
		len = 0;
	}
	/* Do not add new pages in the tree after the initial scan */
	opts->parse_mask |= PARSE_NOADD_TREE;

	printf("Scanning page mappings for process: ");
	while (1) {
		pid = get_next_pid(&dir);
		if (pid <= 0)
			break;

		for (i = 0; i < len; i++)
			putchar('\b');
		len = printf("% 5d", pid);
		fflush(stdout);

		read_pageframe_with_threads(pid, root, process_list, opts);
	}
	for (i = 0; i < len; i++)
		putchar('\b');
	printf("Done  \n");

	return 0;
}

int update_kpageflags(struct rb_root *root)
{
	struct pageframe *pf;
	int kpageflags_fd;
	int kpagecount_fd;
	int ret;

	kpageflags_fd = open("/proc/kpageflags", O_RDONLY);
	if (kpageflags_fd < 0)
		return -1;

	kpagecount_fd = open("/proc/kpagecount", O_RDONLY);
	if (kpagecount_fd < 0)
		return -1;

	pf = rb_to_pageframe(rb_first(root));

	while(pf) {
		long int pfnn = pfn(pf) * sizeof(pf->kpageflags);

		ret = lseek(kpageflags_fd, pfnn, SEEK_SET);
		if (ret < 0) {
			fprintf(stderr, "Error seeking to %lx: %m\n", pfnn);
			return -1;
		}

		ret = read(kpageflags_fd, &pf->kpageflags,
			sizeof(pf->kpageflags));
		if (ret < 0) {
			fprintf(stderr, "Error reading from %llx: %m\n",
				pf->pf * sizeof(pf->kpageflags));
			return -1;
		}

		ret = lseek(kpagecount_fd, pfnn, SEEK_SET);
		if (ret < 0) {
			fprintf(stderr, "Error seeking to %lx: %m\n", pfnn);
			return -1;
		}

		ret = read(kpagecount_fd, &pf->kpagecount,
			sizeof(pf->kpagecount));
		if (ret < 0) {
			fprintf(stderr, "Error reading from %llx: %m\n",
				pf->pf * sizeof(pf->kpagecount));
			return -1;
		}

		pf = rb_to_pageframe(rb_next(&pf->tree));
	}

	close(kpageflags_fd);
	close(kpagecount_fd);

	return 0;
}