/*
 * fincore - count pages of file contents in core
 *
 * Copyright (C) 2017 Red Hat, Inc. All rights reserved.
 * Written by Masatake YAMATO <yamato@redhat.com>
 *
 * 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, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it would 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, see <https://gnu.org/licenses/>.
 */

#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>
#include <getopt.h>
#include <stdio.h>
#include <string.h>

#ifdef HAVE_FTS_OPEN
#include <fts.h>
#endif

#include "c.h"
#include "cctype.h"
#include "nls.h"
#include "closestream.h"
#include "xalloc.h"
#include "strutils.h"
#include "blkdev.h"

#include "libsmartcols.h"

/* For large files, mmap is called in iterative way.
   Window is the unit of vma prepared in each mmap
   calling.

   Window size depends on page size.
   e.g. 128MB on x86_64. ( = N_PAGES_IN_WINDOW * 4096 ). */
#define N_PAGES_IN_WINDOW ((size_t)(32 * 1024))

#ifndef HAVE_CACHESTAT

#ifndef SYS_cachestat
#if defined (__alpha__)
#define SYS_cachestat 561
#else
#define SYS_cachestat 451
#endif
#endif

struct cachestat_range {
	uint64_t off;
	uint64_t len;
};

struct cachestat {
	uint64_t nr_cache;
	uint64_t nr_dirty;
	uint64_t nr_writeback;
	uint64_t nr_evicted;
	uint64_t nr_recently_evicted;
};

static inline int cachestat(unsigned int fd,
			    const struct cachestat_range *cstat_range,
			    struct cachestat *cstat, unsigned int flags)
{
	return syscall(SYS_cachestat, fd, cstat_range, cstat, flags);
}

#endif // HAVE_CACHESTAT

struct colinfo {
	const char * const name;
	double whint;
	int flags;
	const char *help;
	unsigned int pages : 1;
};

enum {
	COL_PAGES,
	COL_SIZE,
	COL_FILE,
	COL_RES,
	COL_DIRTY_PAGES,
	COL_DIRTY,
	COL_WRITEBACK_PAGES,
	COL_WRITEBACK,
	COL_EVICTED_PAGES,
	COL_EVICTED,
	COL_RECENTLY_EVICTED_PAGES,
	COL_RECENTLY_EVICTED,
};

static const struct colinfo infos[] = {
	[COL_PAGES]                  = { "PAGES",                  1, SCOLS_FL_RIGHT, N_("file data resident in memory in pages"), 1},
	[COL_RES]                    = { "RES",                    5, SCOLS_FL_RIGHT, N_("file data resident in memory in bytes")},
	[COL_SIZE]                   = { "SIZE",                   5, SCOLS_FL_RIGHT, N_("size of the file")},
	[COL_FILE]                   = { "FILE",                   4, 0, N_("file name")},
	[COL_DIRTY_PAGES]            = { "DIRTY_PAGES",            1, SCOLS_FL_RIGHT, N_("number of dirty pages"), 1},
	[COL_DIRTY]                  = { "DIRTY",                  5, SCOLS_FL_RIGHT, N_("number of dirty bytes")},
	[COL_WRITEBACK_PAGES]        = { "WRITEBACK_PAGES",        1, SCOLS_FL_RIGHT, N_("number of pages marked for writeback"), 1},
	[COL_WRITEBACK]              = { "WRITEBACK",              5, SCOLS_FL_RIGHT, N_("number of bytes marked for writeback")},
	[COL_EVICTED_PAGES]          = { "EVICTED_PAGES",          1, SCOLS_FL_RIGHT, N_("number of evicted pages"), 1},
	[COL_EVICTED]                = { "EVICTED",                5, SCOLS_FL_RIGHT, N_("number of evicted bytes")},
	[COL_RECENTLY_EVICTED_PAGES] = { "RECENTLY_EVICTED_PAGES", 1, SCOLS_FL_RIGHT, N_("number of recently evicted pages"), 1},
	[COL_RECENTLY_EVICTED]       = { "RECENTLY_EVICTED",       5, SCOLS_FL_RIGHT, N_("number of recently evicted bytes")},
};

static int columns[ARRAY_SIZE(infos) * 2] = {-1};
static size_t ncolumns;

struct fincore_control {
	const size_t pagesize;

	struct libscols_table *tb;		/* output */

	unsigned int bytes : 1,
		     noheadings : 1,
		     raw : 1,
		     json : 1,
		     recursive : 1,
		     total : 1,
		     cachestat : 1;

};

struct fincore_state {
	const char * name;
	long long unsigned int file_size;

	struct cachestat cstat;
	struct {
		unsigned int dirty : 1,
			     writeback : 1,
			     evicted : 1,
			     recently_evicted : 1;
	} cstat_fields;
};

static struct fincore_state total = {
	.name = "TOTAL",
};

static int column_name_to_id(const char *name, size_t namesz)
{
	size_t i;

	for (i = 0; i < ARRAY_SIZE(infos); i++) {
		const char *cn = infos[i].name;

		if (!c_strncasecmp(name, cn, namesz) && !*(cn + namesz))
			return i;
	}
	warnx(_("unknown column: %s"), name);
	return -1;
}

static int get_column_id(int num)
{
	assert(num >= 0);
	assert((size_t) num < ncolumns);
	assert(columns[num] < (int) ARRAY_SIZE(infos));
	return columns[num];
}

static const struct colinfo *get_column_info(int num)
{
	return &infos[ get_column_id(num) ];
}

static int get_cstat_value(const struct fincore_state *st, int column_id,
			   uint64_t *value)
{
	switch(column_id) {
	case COL_PAGES:
	case COL_RES:
		*value = st->cstat.nr_cache;
		return 1;
	case COL_DIRTY_PAGES:
	case COL_DIRTY:
		if (!st->cstat_fields.dirty)
			break;
		*value = st->cstat.nr_dirty;
		return 1;
	case COL_WRITEBACK_PAGES:
	case COL_WRITEBACK:
		if (!st->cstat_fields.writeback)
			break;
		*value = st->cstat.nr_writeback;
		return 1;
	case COL_EVICTED_PAGES:
	case COL_EVICTED:
		if (!st->cstat_fields.evicted)
			break;
		*value = st->cstat.nr_evicted;
		return 1;
	case COL_RECENTLY_EVICTED_PAGES:
	case COL_RECENTLY_EVICTED:
		if (!st->cstat_fields.recently_evicted)
			break;
		*value = st->cstat.nr_recently_evicted;
		return 1;
	default:
		assert(0);
	}
	return 0;
}

static int add_output_data(struct fincore_control *ctl,
			   struct fincore_state *st)
{
	size_t i;
	char *tmp;
	uint64_t value = 0;
	struct libscols_line *ln;

	assert(ctl);
	assert(ctl->tb);

	ln = scols_table_new_line(ctl->tb, NULL);
	if (!ln)
		err(EXIT_FAILURE, _("failed to allocate output line"));

	for (i = 0; i < ncolumns; i++) {
		int rc = 0;
		int column_id = get_column_id(i);
		int format_value = 0;

		switch(column_id) {
		case COL_FILE:
			rc = scols_line_set_data(ln, i, st->name);
			break;
		case COL_SIZE:
			if (ctl->bytes)
				xasprintf(&tmp, "%jd", (intmax_t) st->file_size);
			else
				tmp = size_to_human_string(SIZE_SUFFIX_1LETTER, st->file_size);
			rc = scols_line_refer_data(ln, i, tmp);
			break;
		case COL_PAGES:
		case COL_RES:
		case COL_DIRTY_PAGES:
		case COL_DIRTY:
		case COL_WRITEBACK_PAGES:
		case COL_WRITEBACK:
		case COL_EVICTED:
		case COL_EVICTED_PAGES:
		case COL_RECENTLY_EVICTED:
		case COL_RECENTLY_EVICTED_PAGES:
			format_value = get_cstat_value(st, column_id, &value);
			break;
		default:
			return -EINVAL;
		}

		if (format_value) {
			if (get_column_info(i)->pages) {
				xasprintf(&tmp, "%ju", (uintmax_t) value);
			} else {
				value *= ctl->pagesize;
				if (ctl->bytes)
					xasprintf(&tmp, "%ju", (uintmax_t) value);
				else
					tmp = size_to_human_string(SIZE_SUFFIX_1LETTER, value);
			}
			rc = scols_line_refer_data(ln, i, tmp);
		}

		if (rc)
			err(EXIT_FAILURE, _("failed to add output data"));
	}

	return 0;
}

static int do_mincore(struct fincore_control *ctl,
		      void *window, const size_t len,
		      struct fincore_state *st)
{
	static unsigned char vec[N_PAGES_IN_WINDOW];
	int n = (len / ctl->pagesize) + ((len % ctl->pagesize)? 1: 0);

	if (mincore (window, len, vec) < 0) {
		warn(_("failed to do mincore: %s"), st->name);
		return -errno;
	}

	while (n > 0)
	{
		if (vec[--n] & 0x1)
		{
			vec[n] = 0;
			st->cstat.nr_cache++;
			if (ctl->total)
				total.cstat.nr_cache++;
		}
	}

	return 0;
}

static int mincore_fd (struct fincore_control *ctl,
		       int fd,
		       struct fincore_state *st)
{
	size_t window_size = N_PAGES_IN_WINDOW * ctl->pagesize;
	long long unsigned int file_offset, len;
	int rc = 0;

	for (file_offset = 0; file_offset < st->file_size; file_offset += len) {
		void  *window = NULL;

		len = st->file_size - file_offset;
		if (len >= window_size)
			len = window_size;

		/* PROT_NONE is enough for Linux, but qemu-user wants PROT_READ */
		window = mmap(window, len, PROT_READ, MAP_PRIVATE, fd, file_offset);
		if (window == MAP_FAILED) {
			rc = -EINVAL;
			warn(_("failed to do mmap: %s"), st->name);
			break;
		}

		rc = do_mincore(ctl, window, len, st);
		munmap (window, len);
		if (rc)
			break;

	}

	return rc;
}

static int fincore_fd (struct fincore_control *ctl,
		       int fd,
		       struct fincore_state *st)
{
	int rc;
	const struct cachestat_range cstat_range = { 0 };

	rc = cachestat(fd, &cstat_range, &st->cstat, 0);
	if (!rc) {
		st->cstat_fields.dirty = 1;
		st->cstat_fields.writeback = 1;
		st->cstat_fields.evicted = 1;
		st->cstat_fields.recently_evicted = 1;

		if (ctl->total) {
			total.cstat.nr_cache += st->cstat.nr_cache;
			total.cstat.nr_dirty += st->cstat.nr_dirty;
			total.cstat.nr_writeback += st->cstat.nr_writeback;
			total.cstat.nr_evicted += st->cstat.nr_evicted;
			total.cstat.nr_recently_evicted += st->cstat.nr_recently_evicted;
		}

		return 0;
	}

	if (errno != ENOSYS || ctl->cachestat)
		warn(_("failed to do cachestat: %s"), st->name);

	if (errno == EPERM || ctl->cachestat)
		return -errno;

	return mincore_fd(ctl, fd, st);
}

/*
 * Returns: <0 on error, 0 success, 1 ignore.
 */
static int fincore_name(struct fincore_control *ctl,
			const char *filename,
			const char *showname,
			struct stat *statp)
{
	int fd;
	int rc = 0;
	struct stat _sb, *sb = statp ?: &_sb;
	struct fincore_state _st = { .name = filename }, *st = &_st;

	if ((fd = open(filename, O_RDONLY)) < 0) {
		warn(_("failed to open: %s"), showname);
		return -errno;
	}

	if (!statp) {
		if (fstat (fd, sb) < 0) {
			warn(_("failed to do fstat: %s"), showname);
			close (fd);
			return -errno;
		}
	}

	if (S_ISBLK(sb->st_mode)) {
		rc = blkdev_get_size(fd, &st->file_size);
		if (rc)
			warn(_("failed ioctl to get size: %s"), showname);
	} else if (S_ISREG(sb->st_mode)) {
		st->file_size = sb->st_size;
	} else {
		rc = 1;			/* ignore things like symlinks
					 * and directories*/
	}

	if (!rc)
		rc = fincore_fd(ctl, fd, st);

	close (fd);

	if (!rc) {
		st->name = showname;
		if (ctl->total)
			total.file_size += st->file_size;
		rc = add_output_data(ctl, st);
	}

	return rc;
}

static void __attribute__((__noreturn__)) usage(void)
{
	FILE *out = stdout;
	size_t i;

	fputs(USAGE_HEADER, out);
	fprintf(out, _(" %s [options] file...\n"), program_invocation_short_name);

	fputs(USAGE_OPTIONS, out);
	fputs(_(" -J, --json            use JSON output format\n"), out);
	fputs(_(" -b, --bytes           print sizes in bytes rather than in human readable format\n"), out);
	fputs(_(" -c, --total           produce a grand total\n"), out);
	fputs(_(" -n, --noheadings      don't print headings\n"), out);
	fputs(_(" -o, --output <list>   output columns\n"), out);
	fputs(_("     --output-all      output all columns\n"), out);
	fputs(_(" -r, --raw             use raw output format\n"), out);
	fputs(_(" -R, --recursive       recursively check all files in directories\n"), out);
	fputs(_(" -C, --cachestat       force usage of cachestat syscall\n"), out);

	fputs(USAGE_SEPARATOR, out);
	fprintf(out, USAGE_HELP_OPTIONS(23));

	fputs(USAGE_COLUMNS, out);

	for (i = 0; i < ARRAY_SIZE(infos); i++)
		fprintf(out, " %22s  %s\n", infos[i].name, _(infos[i].help));

	fprintf(out, USAGE_MAN_TAIL("fincore(1)"));

	exit(EXIT_SUCCESS);
}

int main(int argc, char ** argv)
{
	int c;
	size_t i;
	int rc = EXIT_SUCCESS;
	char *outarg = NULL;

	struct fincore_control ctl = {
		.pagesize = getpagesize()
	};

	enum {
		OPT_OUTPUT_ALL = CHAR_MAX + 1
	};
	static const struct option longopts[] = {
		{ "bytes",      no_argument, NULL, 'b' },
		{ "total",      no_argument, NULL, 'c' },
		{ "noheadings", no_argument, NULL, 'n' },
		{ "output",     required_argument, NULL, 'o' },
		{ "output-all",	no_argument,       NULL, OPT_OUTPUT_ALL },
		{ "version",    no_argument, NULL, 'V' },
		{ "help",	no_argument, NULL, 'h' },
		{ "json",       no_argument, NULL, 'J' },
		{ "raw",        no_argument, NULL, 'r' },
		{ "recursive",	no_argument, NULL, 'R' },
		{ "cachestat",  no_argument, NULL, 'C' },
		{ NULL, 0, NULL, 0 },
	};

	setlocale(LC_ALL, "");
	bindtextdomain(PACKAGE, LOCALEDIR);
	textdomain(PACKAGE);
	close_stdout_atexit();

	while ((c = getopt_long (argc, argv, "bcCno:JrRVh", longopts, NULL)) != -1) {
		switch (c) {
		case 'b':
			ctl.bytes = 1;
			break;
		case 'c':
			ctl.total = 1;
			break;
		case 'C':
#ifndef HAVE_CACHESTAT
			errx(EXIT_FAILURE, _("cachestat option is not supported"));
#endif
			ctl.cachestat = 1;
			break;
		case 'n':
			ctl.noheadings = 1;
			break;
		case 'o':
			outarg = optarg;
			break;
		case OPT_OUTPUT_ALL:
			for (ncolumns = 0; ncolumns < ARRAY_SIZE(infos); ncolumns++)
				columns[ncolumns] = ncolumns;
			break;
		case 'J':
			ctl.json = 1;
			break;
		case 'r':
			ctl.raw = 1;
			break;
		case 'R':
#ifndef HAVE_FTS_OPEN
			errx(EXIT_FAILURE, _("recursive option is not supported"));
#endif
			ctl.recursive = 1;
			break;
		case 'V':
			print_version(EXIT_SUCCESS);
		case 'h':
			usage();
		default:
			errtryhelp(EXIT_FAILURE);
		}
	}

	if (optind == argc) {
		warnx(_("no file specified"));
		errtryhelp(EXIT_FAILURE);
	}

	if (!ncolumns) {
		columns[ncolumns++] = COL_RES;
		columns[ncolumns++] = COL_PAGES;
		columns[ncolumns++] = COL_SIZE;
		columns[ncolumns++] = COL_FILE;
	}

	if (outarg && string_add_to_idarray(outarg, columns, ARRAY_SIZE(columns),
					 &ncolumns, column_name_to_id) < 0)
		return EXIT_FAILURE;

	scols_init_debug(0);
	ctl.tb = scols_new_table();
	if (!ctl.tb)
		err(EXIT_FAILURE, _("failed to allocate output table"));

	scols_table_enable_noheadings(ctl.tb, ctl.noheadings);
	scols_table_enable_raw(ctl.tb, ctl.raw);
	scols_table_enable_json(ctl.tb, ctl.json);
	if (ctl.json)
		scols_table_set_name(ctl.tb, "fincore");

	for (i = 0; i < ncolumns; i++) {
		const struct colinfo *col = get_column_info(i);
		struct libscols_column *cl;

		cl = scols_table_new_column(ctl.tb, col->name, col->whint, col->flags);
		if (!cl)
			err(EXIT_FAILURE, _("failed to allocate output column"));

		if (ctl.json) {
			int id = get_column_id(i);

			switch (id) {
			case COL_FILE:
				scols_column_set_json_type(cl, SCOLS_JSON_STRING);
				break;
			case COL_SIZE:
			case COL_RES:
				if (!ctl.bytes)
					break;
				FALLTHROUGH;
			default:
				scols_column_set_json_type(cl, SCOLS_JSON_NUMBER);
				break;
			}
		}
	}

	if (ctl.recursive) {
#ifdef HAVE_FTS_OPEN
		FTS *fts = fts_open(argv + optind, FTS_PHYSICAL, NULL);
		FTSENT *ent;

		if (!fts) {
			warn(_("failed to iterate tree"));
			rc = EXIT_FAILURE;
		} else {
			while ((ent = fts_read(fts)) != NULL) {
				if (ent->fts_info == FTS_F || ent->fts_info == FTS_DEFAULT) {
					/* fts changes directory when iterating,
					 * so we need to use .fts_accpath to access
					 * the file named .fts_path */
					rc |= fincore_name(&ctl, ent->fts_accpath, ent->fts_path, ent->fts_statp);
				}
			}
			fts_close(fts);
		}
#endif
	} else {
		for(; optind < argc; optind++)
			rc |= fincore_name(&ctl, argv[optind], argv[optind], NULL);
	}

	if (ctl.total)
		rc |= add_output_data(&ctl, &total);

	scols_print_table(ctl.tb);
	scols_unref_table(ctl.tb);

	return rc ? EXIT_FAILURE : EXIT_SUCCESS;
}
