/* gpu.c - rmtool gpu list/topology/info: enumeration and identity queries. */
#include <getopt.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "rmtool.h"

static int read_sysfs(const char *bdf, const char *file, char *buf, size_t len)
{
	char path[512];
	FILE *fh;

	snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/%s", bdf, file);
	fh = fopen(path, "r");
	if (fh == NULL) {
		buf[0] = '\0';
		return -1;
	}
	if (fgets(buf, (int)len, fh) == NULL)
		buf[0] = '\0';
	fclose(fh);
	buf[strcspn(buf, "\r\n")] = '\0';
	return 0;
}

static int parse_common(int argc, char **argv, struct rm_ctx *ctx)
{
	static const struct option long_opts[] = {
		{ "device-spec", required_argument, NULL, 'd' },
		{ "format", required_argument, NULL, 'f' },
		{ "device", required_argument, NULL, 1000 },
		{ "help", no_argument, NULL, 'h' },
		{ NULL, 0, NULL, 0 },
	};
	int opt;

	optind = 1;
	opterr = 0;
	while ((opt = getopt_long(argc, argv, "d:f:h", long_opts, NULL)) != -1) {
		int rc;

		if (opt == 'h')
			return 1;
		rc = rm_apply_common_opt(ctx, opt, optarg);
		if (rc < 0)
			return -1;
		if (rc == 0) {
			fprintf(stderr, "rmtool: unknown option\n");
			return -1;
		}
	}
	if (optind < argc) {
		fprintf(stderr, "rmtool: unexpected argument '%s'\n",
			argv[optind]);
		return -1;
	}
	return 0;
}

static void list_usage(void)
{
	printf("usage: rmtool gpu list [--format=table|json|csv]\n"
	       "       rmtool gpu topology [--format=...]\n"
	       "       rmtool gpu info [-d N|BDF] [--format=...]\n");
}

static int gpu_list(struct rm_ctx *ctx)
{
	struct rm_gpu_desc gpus[RM_MAX_GPUS];
	const char *headers[] = { "MINOR", "BUS", "MODEL", "UUID", "VBIOS" };
	char cells[RM_MAX_GPUS * 5][RM_UUID_LEN + 8];
	int n;
	int i;

	n = rm_scan_gpus(gpus, RM_MAX_GPUS);
	if (n < 0) {
		fprintf(stderr, "rmtool: cannot read /proc/driver/nvidia/gpus: "
			"%s\n", strerror(-n));
		return 2;
	}
	for (i = 0; i < n; i++) {
		snprintf(cells[i * 5 + 0], sizeof(cells[0]), "%d",
			 gpus[i].minor);
		snprintf(cells[i * 5 + 1], sizeof(cells[0]), "%s",
			 gpus[i].bdf);
		snprintf(cells[i * 5 + 2], sizeof(cells[0]), "%s",
			 gpus[i].model);
		snprintf(cells[i * 5 + 3], sizeof(cells[0]), "%s",
			 gpus[i].uuid);
		snprintf(cells[i * 5 + 4], sizeof(cells[0]), "%s",
			 gpus[i].vbios);
	}
	if (ctx->format == RM_FMT_TABLE)
		printf("%d GPU(s) on this host\n", n);
	rm_print_table("gpus", headers, &cells[0][0], sizeof(cells[0]), n, 5,
		       ctx->format);
	return 0;
}

static int gpu_topology(struct rm_ctx *ctx)
{
	struct rm_gpu_desc gpus[RM_MAX_GPUS];
	const char *headers[] = { "MINOR", "BUS", "LINK", "MAX LINK", "NUMA",
				  "VENDOR:DEVICE" };
	char cells[RM_MAX_GPUS * 6][RM_UUID_LEN];
	int n;
	int i;

	n = rm_scan_gpus(gpus, RM_MAX_GPUS);
	if (n < 0) {
		fprintf(stderr, "rmtool: cannot read /proc/driver/nvidia/gpus: "
			"%s\n", strerror(-n));
		return 2;
	}
	for (i = 0; i < n; i++) {
		char speed[32], width[32], max_speed[32], max_width[32];
		char numa[16], vendor[32], device[32];

		read_sysfs(gpus[i].bdf, "current_link_speed", speed,
			   sizeof(speed));
		read_sysfs(gpus[i].bdf, "current_link_width", width,
			   sizeof(width));
		read_sysfs(gpus[i].bdf, "max_link_speed", max_speed,
			   sizeof(max_speed));
		read_sysfs(gpus[i].bdf, "max_link_width", max_width,
			   sizeof(max_width));
		read_sysfs(gpus[i].bdf, "numa_node", numa, sizeof(numa));
		read_sysfs(gpus[i].bdf, "vendor", vendor, sizeof(vendor));
		read_sysfs(gpus[i].bdf, "device", device, sizeof(device));
		snprintf(cells[i * 6 + 0], sizeof(cells[0]), "%d",
			 gpus[i].minor);
		snprintf(cells[i * 6 + 1], sizeof(cells[0]), "%s",
			 gpus[i].bdf);
		snprintf(cells[i * 6 + 2], sizeof(cells[0]), "%s x%s",
			 speed, width);
		snprintf(cells[i * 6 + 3], sizeof(cells[0]), "%s x%s",
			 max_speed, max_width);
		snprintf(cells[i * 6 + 4], sizeof(cells[0]), "%s", numa);
		snprintf(cells[i * 6 + 5], sizeof(cells[0]), "%s:%s", vendor,
			 device);
	}
	rm_print_table("gpu topology (sysfs, read-only)", headers,
		       &cells[0][0], sizeof(cells[0]), n, 6, ctx->format);
	return 0;
}

static int gpu_info(struct rm_ctx *ctx)
{
	struct rm_gpu_desc gpus[RM_MAX_GPUS];
	NV2080_CTRL_MC_ARCH_INFO_PARAMS arch;
	NV2080_CTRL_BUS_GET_PCI_INFO_PARAMS pci;
	NV2080_CTRL_GPU_GET_ENGINES_V2_PARAMS engines;
	struct rm_kv kv[20];
	char v[20][96];
	int n = 0;
	int i, n_gpus, rc;
	NvU32 status;

	n_gpus = rm_scan_gpus(gpus, RM_MAX_GPUS);
	rc = rm_ctx_open(ctx);
	if (rc != 0) {
		fprintf(stderr, "rmtool: cannot open RM client "
			"(need the driver bound and read access): rc=%d\n", rc);
		return 2;
	}

	status = rm_query_name(ctx, v[0], sizeof(v[0]));
	kv[n].key = "Name";
	kv[n++].value = status == 0 ? v[0] : "query failed";

	memset(v[1], 0, sizeof(v[1]));
	snprintf(v[1], sizeof(v[1]), "%d", ctx->rm.gpu_minor);
	kv[n].key = "GPU minor";
	kv[n++].value = v[1];

	if (n_gpus > 0) {
		const struct rm_gpu_desc *g = NULL;

		for (i = 0; i < n_gpus; i++) {
			if (gpus[i].minor == ctx->rm.gpu_minor)
				g = &gpus[i];
		}
		if (g == NULL)
			g = &gpus[0];
		kv[n].key = "Model (proc)";
		kv[n++].value = g->model;
		kv[n].key = "Bus location";
		kv[n++].value = g->bdf;
		kv[n].key = "GPU UUID";
		kv[n++].value = g->uuid;
		kv[n].key = "Video BIOS";
		kv[n++].value = g->vbios;
	}

	status = rm_query_mc_arch(ctx, &arch);
	if (status == 0) {
		snprintf(v[2], sizeof(v[2]), "%s (0x%x/0x%x rev %u)",
			 rm_arch_str(arch.architecture), arch.architecture,
			 arch.implementation, arch.revision);
		kv[n].key = "Architecture";
		kv[n++].value = v[2];
	} else {
		snprintf(v[2], sizeof(v[2]), "MC_GET_ARCH_INFO: %s",
			 rm_status_str(status));
		kv[n].key = "Architecture";
		kv[n++].value = v[2];
	}

	status = rm_query_pci(ctx, &pci);
	if (status == 0) {
		snprintf(v[3], sizeof(v[3]), "%04x:%04x (subsystem %04x)",
			 pci.pciDeviceId & 0xffff,
			 (pci.pciDeviceId >> 16) & 0xffff,
			 pci.pciSubSystemId);
		kv[n].key = "PCI device";
		kv[n++].value = v[3];
		snprintf(v[4], sizeof(v[4]), "%u", pci.pciRevisionId);
		kv[n].key = "PCI revision";
		kv[n++].value = v[4];
	}

	status = rm_query_engines(ctx, &engines);
	snprintf(v[5], sizeof(v[5]), "%s",
		 status == 0 ? "ok" : rm_status_str(status));
	kv[n].key = "Engine inventory";
	kv[n++].value = v[5];
	if (status == 0) {
		snprintf(v[6], sizeof(v[6]), "%u", engines.engineCount);
		kv[n].key = "Engine count";
		kv[n++].value = v[6];
	}

	rm_print_kv("rmtool gpu info (read-only RM queries)", kv, n,
		    ctx->format);
	return 0;
}

int cmd_gpu(int argc, char **argv, struct rm_ctx *ctx)
{
	const char *sub;

	if (argc < 2) {
		list_usage();
		return 2;
	}
	sub = argv[1];
	if (strcmp(sub, "list") == 0) {
		int rc = parse_common(argc - 1, argv + 1, ctx);

		if (rc < 0)
			return 2;
		if (rc == 1) {
			list_usage();
			return 0;
		}
		return gpu_list(ctx);
	}
	if (strcmp(sub, "topology") == 0) {
		int rc = parse_common(argc - 1, argv + 1, ctx);

		if (rc < 0)
			return 2;
		if (rc == 1) {
			list_usage();
			return 0;
		}
		return gpu_topology(ctx);
	}
	if (strcmp(sub, "info") == 0) {
		int rc = parse_common(argc - 1, argv + 1, ctx);

		if (rc < 0)
			return 2;
		if (rc == 1) {
			list_usage();
			return 0;
		}
		return gpu_info(ctx);
	}
	list_usage();
	return 2;
}