/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2026, Netflix, Inc.
 *
 * This software was developed by Ali Mashtizadeh under the sponsorship from
 * Netflix, Inc.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 *
 */

#include <sys/param.h>
#include <sys/cdefs.h>
#include <sys/cpuset.h>
#include <sys/event.h>
#include <sys/queue.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/sysctl.h>
#include <sys/time.h>
#include <sys/ttycom.h>
#include <sys/user.h>
#include <sys/wait.h>

#include <assert.h>
#include <curses.h>
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <kvm.h>
#include <libgen.h>
#include <limits.h>
#include <locale.h>
#include <math.h>
#include <pmc.h>
#include <pmclog.h>
#include <regex.h>
#include <signal.h>
#include <stdalign.h>
#include <stdarg.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <stddef.h>
#include <string.h>
#include <sysexits.h>
#include <unistd.h>

#include <machine/cpufunc.h>

#include <libpmcstat.h>
#include "cmd_pmc.h"

#include <iostream>
#include <map>
#include <string>
#include <vector>

#include "display.hh"
#include "headers.hh"

#define DEFAULT_RATE 65536

class pmc_config
{
public:
	std::string	event;
	uint64_t	count;
	cpuset_t	cpumask;
	uint32_t	caps;
	pmc_config() : event(), count(0), ids()
	{
		CPU_ZERO(&cpumask);
	}
	pmc_config(const std::string &_event, uint64_t _count,
	    cpuset_t _cpumask)
	    : event(_event), count(_count), ids()
	{
		CPU_COPY(&_cpumask, &cpumask);
	}
	/*
	 * Read the capabilities and fills in the cpumask for a given event 
	 * counter.  The way that the hwpmc module currently works we do not 
	 * know which pmc class will back our counter until we look it up.  The 
	 * easiest way is to program the counter into the first CPU present in 
	 * the mask and then retrieve the capabilities.
	 */
	void getcaps() {
		int status;
		int testcpu;
		pmc_id_t id;

		testcpu = CPU_FFS(&cpumask) - 1;

		status = pmc_allocate(event.c_str(), PMC_MODE_SS, PMC_F_CALLCHAIN,
		    testcpu, &id, count);
		if (status < 0)
			err(EX_OSERR, "ERROR: Cannot allocate event '%s'", event.c_str());

		status = pmc_capabilities(id, &caps);
		if (status < 0)
			err(EX_OSERR, "ERROR: Cannot get pmc capabilities");

		pmc_release(id);

		if (caps & PMC_CAP_SYSWIDE) {
			CPU_ZERO(&cpumask);
			CPU_SET(0, &cpumask);
		}
		if (caps & PMC_CAP_DOMWIDE) {
			int domains;
			size_t len;

			CPU_ZERO(&cpumask);

			len = sizeof(domains);
			if (sysctlbyname("vm.ndomains", &domains, &len, NULL, 0) == -1)
				err(EX_OSERR, "ERROR: Cannot get number of domains");

			for (int i = 1; i < domains; i++) {
				cpuset_t dmask;

				CPU_ZERO(&dmask);
				status = cpuset_getaffinity(CPU_LEVEL_WHICH, CPU_WHICH_DOMAIN, i,
				    sizeof(dmask), &dmask);
				if (status < 0)
					err(EX_OSERR, "ERROR: Cannot get domain mask");
				CPU_SET(CPU_FFS(&dmask) - 1, &cpumask);
			}
		}
	}
	/*
	 * Allocate the PMC across all CPUs.
	 */
	void allocate() {
		int cpu;
		int status;
		pmc_id_t id;

		for (cpu = 0; cpu < CPU_SETSIZE; cpu++) {
			if (!CPU_ISSET(cpu, &cpumask))
				continue;

			status = pmc_allocate(event.c_str(), PMC_MODE_SS, PMC_F_CALLCHAIN,
			    cpu, &id, count);
			if (status < 0)
				err(EX_OSERR, "ERROR: Cannot allocate event '%s'",
				    event.c_str());

			status = pmc_set(id, count);
			if (status < 0)
				err(EX_OSERR, "ERROR: Cannot set sampling count for event '%s'",
				    event.c_str());

			ids.push_back(id);
		}
	}
	/*
	 * Release all PMC instances (one per CPU).
	 */
	void release() {
		for (auto i : ids) {
			if (pmc_release(i) < 0) {
				perror("pmc_start");
			}
		}
	}
	/*
	 * Start all PMC instances (one per CPU).
	 */
	void start() {
		for (auto i : ids) {
			if (pmc_start(i) < 0) {
				perror("pmc_start");
			}
		}
	}
	/*
	 * Stop all PMC instances (one per CPU).
	 */
	void stop() {
		for (auto i : ids) {
			if (pmc_stop(i) < 0) {
				perror("pmc_start");
			}
		}
	}
private:
	std::vector<pmc_id_t> ids;
};

static std::vector<pmc_config> pmcs = std::vector<pmc_config>();
static struct option longopts[] = {
	{ "rate",	required_argument,	NULL,	'n' },
	{ "counter",	required_argument,	NULL,	'c' },
	{ "time",	required_argument,	NULL,	't' },
	{ "study",	required_argument,	NULL,	's' },
	{ NULL,		0,			NULL,	0 }
};
static struct timespec start;
static int timelimit = 0;

static void
usage(void)
{
	printf("Usage: pmc record [options] [output.pmc]\n\n");
	printf("Record a study\n\n");
	printf("Options:\n");
	printf("\t-c              Sampling counter\n");
	printf("\t-r              Sample rate (default: %d)\n", DEFAULT_RATE);
	printf("\t-s              Specify a study\n");
	printf("\t-t              Run the study for specified amount of time\n");
	printf("\nStudies:\n");
#if defined(__i386__) || defined(__amd64__)
	printf("\tbranches        Study branch misprediction (Requires AMD IBS)\n");
	printf("\tc2c             Study cache to cache communications (Requires AMD IBS)\n");
#endif
	printf("\tdefault         Instruction sampling for general analysis\n");
	printf("\tflamegraph      Instruction sampling for general analysis\n");
#if defined(__i386__) || defined(__amd64__)
	printf("\tfrontend        Study frontend stalls (Requires AMD IBS)\n");
	printf("\tmemory          Study memory operations (Requires AMD IBS)\n");
#endif
}

int
writelog(int logfd, const void *buf, size_t len)
{
	int status;
	const char *cur = (const char *)buf;

	while (len != 0) {
		status = write(logfd, cur, len);
		if (status < 0) {
			if (errno == EINTR || errno == EAGAIN)
				continue;
			else
				return status;
		}
		if (status == 0)
			return status;

		cur += status;
		len -= status;
	}

	return 0;
}

/*
 * Write the PMC header.  Each of the header payloads have their own header 
 * containing the payloads size.  Thus the pmc tools can just skip over regions 
 * that they do not know how to handle.
 */
int
write_header(int logfd)
{
	int status;
	pmchdr_header hdr;

	hdr.magic = PMC_HEADER_MAGIC;
	hdr.version = PMC_HEADER_VERSION;
#if defined(__amd64__)
	hdr.arch = PMC_ARCH_AMD64;
#elif defined(__aarch64__)
	hdr.arch = PMC_ARCH_ARM64;
#elif defined(__powerpc64__)
	hdr.arch = PMC_ARCH_PPC64;
#elif defined(__riscv)
	hdr.arch = PMC_ARCH_RISCV64;
#elif defined(__arm__)
	hdr.arch = PMC_ARCH_ARM;
#else
	hdr.arch = 0;
#endif

	status = writelog(logfd, &hdr, sizeof(hdr));
	if (status < 0) {
		perror("writelog");
	}

	return status;
}

/*
 * Write out the sysinfo header.
 */
int
write_sysinfo(int logfd)
{
	int status;
	pmchdr_infohdr hdr;
	pmchdr_sysinfo sys;
	char val[64];
	size_t valsz;

	valsz = sizeof(val);
	status = sysctlbyname("hw.model", &val, &valsz, NULL, 0);
	if (status < 0) {
		perror("sysctlbyname");
		return status;
	}
	strlcpy(sys.cpumodel, val, sizeof(sys.cpumodel));

	valsz = sizeof(val);
	status = sysctlbyname("kern.osrelease", &val, &valsz, NULL, 0);
	if (status < 0) {
		perror("sysctlbyname");
		return status;
	}
	strlcpy(sys.osrelease, val, sizeof(sys.osrelease));

	valsz = sizeof(val);
	status = sysctlbyname("kern.build_id", &val, &valsz, NULL, 0);
	if (status < 0) {
		perror("sysctlbyname");
		return status;
	}
	strlcpy(sys.buildid, val, sizeof(sys.buildid));

	hdr.type = INFOHDR_TYPE_SYSINFO;
	hdr.length = sizeof(sys);
	status = writelog(logfd, &hdr, sizeof(hdr));
	if (status < 0) {
		err(EX_IOERR, "writelog");
	}

	status = writelog(logfd, &sys, sizeof(sys));
	if (status < 0) {
		err(EX_IOERR, "writelog");
	}

	return status;
};

#if defined(__i386__) || defined(__amd64__)
#define CPUID_ROOT_BASE		0x00000000
#define CPUID_ROOT_VM		0x40000000
#define CPUID_ROOT_EXT		0x80000000

/*
 * Write out the CPU Info block.  After the standard header that declares the 
 * size of the payload, we write out all the CPUID root and leafs for each of 
 * the three major roots: base, VM, and extended artribute space.  On all 
 * modern processors the root contains the maximum leaf as the first value 
 * allowing us to decode how many leafs there are.
 */
int
write_cpuinfo(int logfd)
{
	int status;
	u_int tmp[4];
	uint32_t base_max, vm_max, ext_max;
	uint32_t len;
	pmchdr_infohdr *hdr;
	uint32_t *buf, *off;

	/*
	 * Find the length of the main, VM, and extended CPUID leafs
	 */
	do_cpuid(CPUID_ROOT_BASE, tmp);
	base_max = tmp[0];
	len = base_max + 1;

	do_cpuid(CPUID_ROOT_VM, tmp);
	vm_max = tmp[0];
	if (vm_max != 0)
		len += vm_max - CPUID_ROOT_VM + 1;

	do_cpuid(CPUID_ROOT_EXT, tmp);
	ext_max = tmp[0];
	if (ext_max != 0)
		len += ext_max - CPUID_ROOT_EXT + 1;

	// 4 Registers per Leaf x 4 Bytes per Register
	len *= 4;

	buf = (uint32_t *)new uint32_t[len + sizeof(pmchdr_infohdr) / 4];
	hdr = (pmchdr_infohdr *)buf;
	hdr->type = INFOHDR_TYPE_CPUID;
	hdr->length = 4 * len;
	off = (uint32_t *)(hdr + 1);

	for (uint32_t i = 0; i <= base_max; i++) {
		do_cpuid(i, off);
		off += 4;
	}

	if (vm_max) {
		for (uint32_t i = CPUID_ROOT_VM; i <= vm_max; i++) {
			do_cpuid(i, off);
			off += 4;
		}
	}
	if (ext_max) {
		for (uint32_t i = CPUID_ROOT_EXT; i <= ext_max; i++) {
			do_cpuid(i, off);
			off += 4;
		}
	}

	status = writelog(logfd, buf, 4 * len + sizeof(*hdr));
	if (status < 0) {
		delete[] buf;
		err(EX_IOERR, "writelog");
	}

	delete[] buf;

	return 0;
}
#elif defined(__aarch64__)
int
write_cpuinfo(__unused int logfd)
{
	return 0;
}
#elif defined(__powerpc64__)
int
write_cpuinfo(__unused int logfd)
{
	return 0;
}
#else
int
write_cpuinfo(__unused int logfd)
{
	return 0;
}
#endif

int
write_footer(int logfd)
{
	int status;
	pmchdr_infohdr hdr;

	hdr.type = INFOHDR_TYPE_DONE;
	hdr.length = 0;

	status = writelog(logfd, &hdr, sizeof(hdr));
	if (status < 0) {
		perror("writelog");
	}

	return status;
}

#if defined(__i386__) || defined(__amd64__)
int
setup_study(const std::string &study, uint64_t rate, cpuset_t mask)
{
	u_int tmp[4];
	alignas(4) char vendor[16];
	std::string event;
	uint32_t ext_max;
	uint32_t ibs_features;
	bool is_amd, is_intel;

	do_cpuid(CPUID_ROOT_BASE, tmp);

	/* Find the brand */
	is_intel = false;
	is_amd = false;

	/* i386 complains without explicit alignment */
	((u_int *)&vendor)[0] = tmp[1];
	((u_int *)&vendor)[1] = tmp[3];
	((u_int *)&vendor)[2] = tmp[2];
	vendor[12] = 0;
	if (strncmp(vendor, INTEL_VENDOR_ID, 12) == 0)
		is_intel = true;
	if (strncmp(vendor, AMD_VENDOR_ID, 12) == 0)
		is_amd = true;
	if (strncmp(vendor, HYGON_VENDOR_ID, 12) == 0)
		is_amd = true;

	ibs_features = 0;
	do_cpuid(CPUID_ROOT_EXT, tmp);
	ext_max = tmp[0];
	if (is_amd && ext_max >= CPUID_IBSID) {
		do_cpuid(CPUID_IBSID, tmp);
		ibs_features = tmp[0];
	}

	if (study == "default" || study == "flamegraph") {
		pmcs.push_back(pmc_config("unhalted-cycles", rate, mask));
		return 0;
	}

	if (is_intel) {
		err(EX_SOFTWARE, "ERROR: Study is unsupported on Intel");
	}

	if (study == "frontend") {
		event = "ibs-fetch,randomize";
	} else if (study == "branches" || study == "memory") {
		event = "ibs-op";
		if (ibs_features & CPUID_IBSID_OPCNT)
			event += ",opcount";
	} else if (study == "c2c") {
		event = "ibs-op";
		if (ibs_features & CPUID_IBSID_IBSLOADLATENCYFILT)
			event += ",l3miss";
		if (ibs_features & CPUID_IBSID_OPCNT)
			event += ",opcount";
	} else {
		err(EX_USAGE, "ERROR: Study '%s' unknown", study.c_str());
	}

	pmcs.push_back(pmc_config(event, rate, mask));

	return 0;
}
#else
int
setup_study(const std::string &study, uint64_t rate, cpuset_t mask)
{
	if (study == "default" || study == "flamegraph") {
		pmcs.push_back(pmc_config("unhalted-cycles", rate, mask));
		return 0;
	}

	err(EX_SOFTWARE, "ERROR: Study is unsupported on this architecture");
}
#endif

/*
 * Compute the current runtime and print it to the screen, if it exceeds 
 * timelimit return 1 otherwise 0.
 */
int
update_status(int logfd)
{
	int status;
	long sec, nsec;
	struct timespec end;
	struct stat sb;
	std::string filesz;

	status = clock_gettime(CLOCK_MONOTONIC, &end);
	if (status < 0) {
		err(EX_OSERR, "Could not read current time");
	}

	sec = end.tv_sec - start.tv_sec;
	nsec = end.tv_nsec - start.tv_nsec;
	if (nsec < 0) {
		sec--;
		nsec += 1000000000L;
	}

	// Do this another way for sockets
	status = fstat(logfd, &sb);
	if (status < 0) {
		err(EX_OSERR, "Failed to fstat log file");
	}

	filesz = format_binprefix(sb.st_size);

	/*
	 * Use %-10s to ensure we print spaces after in case the size string 
	 * shrinks.
	 */
	printf("Elapsed: %ld.%03ld seconds, Log file size: %-10s\r",
	    sec, nsec / 1000000,
	    filesz.c_str());
	fflush(stdout);

	if (timelimit != 0 && sec >= timelimit) {
		printf("\nElapsed time completed\n");
		return 1;
	}

	return 0;
}

int
cmd_pmc_record(int argc, char **argv)
{
	struct kevent kev;
	cpuset_t mask;
	uint64_t rate = DEFAULT_RATE;
	const char *logfile = "default.log";
	int status;
	int option, logfd, kq;

	CPU_ZERO(&mask);
	if (cpuset_getaffinity(CPU_LEVEL_ROOT, CPU_WHICH_PID, -1,
	    sizeof(mask), &mask) == -1)
		err(EX_OSERR, "ERROR: Cannot determine the available CPUs");

	if (pmc_init() < 0)
		err(EX_SOFTWARE, "ERROR: Failed to initialize libpmc");

	while ((option = getopt_long(argc, argv, "c:r:s:t:", longopts, NULL)) != -1) {
		switch (option) {
		case 'c':
			pmcs.push_back(pmc_config(optarg, rate, mask));
			break;
		case 'r':
			rate = strtoull(optarg, NULL, 0);
			break;
		case 's':
			setup_study(optarg, rate, mask);
			break;
		case 't':
			timelimit = atoi(optarg);
			break;
		case '?':
		default:
			usage();
		}
	}
	argc -= optind;
	argv += optind;
	if (argc != 0 && argc != 1) {
		usage();
		exit(EX_USAGE);
	}
	if (argc == 1)
		logfile = argv[0];

	if (pmcs.size() == 0) {
		errx(EX_NOINPUT, "ERROR: Please select one or more counters or performance studies.");
	}

	if ((logfd = open(logfile, O_CREAT|O_EXCL|O_WRONLY,
			S_IRUSR|S_IWUSR|S_IRGRP|S_IROTH)) < 0) {
		errx(EX_OSERR, "ERROR: Cannot open \"%s\" for writing: %s.", logfile,
		    strerror(errno));
	}

	setup_screen();

	for (auto &p : pmcs) {
		p.getcaps();
	}

	write_header(logfd);
	write_sysinfo(logfd);
	write_cpuinfo(logfd);
	write_footer(logfd);

	// Record pmclog
	kq = kqueue();
	if (kq < 0)
		err(EX_OSERR, "ERROR: kqueue creation failed");

	EV_SET(&kev, fileno(stdin), EVFILT_READ, EV_ADD, 0, 0, NULL);
	if (kevent(kq, &kev, 1, NULL, 0, NULL) < 0)
		err(EX_OSERR, "ERROR: kevent failed to register stdin");

	EV_SET(&kev, SIGINT, EVFILT_SIGNAL, EV_ADD, 0, 0, NULL);
	if (kevent(kq, &kev, 1, NULL, 0, NULL) < 0)
		err(EX_OSERR, "ERROR: kevent failed to register SIGINT");
	EV_SET(&kev, SIGIO, EVFILT_SIGNAL, EV_ADD, 0, 0, NULL);
	if (kevent(kq, &kev, 1, NULL, 0, NULL) < 0)
		err(EX_OSERR, "ERROR: kevent failed to register SIGIO");

	EV_SET(&kev, 0, EVFILT_TIMER, EV_ADD, 0, 100, NULL);
	if (kevent(kq, &kev, 1, NULL, 0, NULL) < 0)
		err(EX_OSERR, "ERROR: Cannot register kevent for timer");

	pmc_configure_logfile(logfd);

	for (auto &p : pmcs) {
		p.allocate();
	}

	for (auto &p : pmcs) {
		p.start();
	}

	if (clock_gettime(CLOCK_MONOTONIC, &start) < 0)
		err(EX_OSERR, "ERROR: Could not get the current time");

	printf("Recording performance trace press Ctrl-C to stop\n");

	/*
	 * loop till either the target process (if any) exits, or we
	 * are killed by a SIGINT or we reached the time duration.
	 */
	while (1) {
		status = kevent(kq, NULL, 0, &kev, 1, NULL);
		if (status <= 0) {
			if (errno != EINTR)
				err(EX_OSERR, "ERROR: kevent failed");
			else
				continue;
		}

		if (kev.flags & EV_ERROR)
			errc(EX_OSERR, kev.data, "ERROR: kevent failed");

		switch (kev.filter) {
		case EVFILT_READ:  /* log file data is present */
			if (kev.ident == (unsigned)fileno(stdin)) {
				// Check for exit key
			}
			break;
		case EVFILT_SIGNAL:
			if (kev.ident == SIGIO) {
				fprintf(stderr, "ERROR: IO error");
				status = EX_OSERR;
				goto done;
			} else if (kev.ident == SIGINT) {
				status = EX_OK;
				goto done;
			} else {
				err(EX_OSERR, "Unknown signal recieved");
			}
			break;
		case EVFILT_TIMER:
			if (update_status(logfd) == 1)
				goto done;
			break;
		}
	}

done:
	for (auto &p : pmcs) {
		p.stop();
		p.release();
	}
	close(logfd);

	return (status);
}