#include <stdio.h>
#include <string.h>
#include <dirent.h>
#include <stdint.h>
#include <math.h>
#include <stdbool.h>

#include "blackwell-sensors.h"

/* Number of die sensors, 4 for Blackwell */
#define BLACKWELL_NUM_DIE_SENSORS 4
/* Memory spacing between DRAM FBPA blocks */
#define BLACKWELL_FBPA_STRIDE 0x4000
/* Maximum count of FBPA blocks, technically up to 18 but no consumer card has more than 8 */
#define BLACKWELL_MAX_NUM_FBPA 9

/* Yeah I know this is duplicated */
#define HEADER_COLOR_START "\x1b[36m"
#define HEADER_COLOR_END "\x1b[39m"

static float read_temp_sensor(void *addr);
static int read_dram_temp_sensor(void *addr);

struct blackwell_card_info {
	char *card_name;
	unsigned short vendor_id;
	unsigned short device_id;
	int die_temp_offset[BLACKWELL_NUM_DIE_SENSORS];
	int sys_temp_offset;
	int dram_temp_offset_start;
	int clamshell;			/* If the card has dual-side ram, if not then half the DRAM temps are repeated */
};

static struct blackwell_card_info blackwell_card_list[] =
{
	{"NVIDIA RTX PRO 6000 Blackwell Workstation Edition",	NVIDIA_VEN,		NVIDIA_RTXPRO6000_WORKSTATION_DEV,	{0xAD0A90, 0xAD0A94, 0xAD0A98, 0xAD0A9C},	0xAD06EC,	0x9024C0,	true}
};


/* Search all i2c device files for ones are on a PCI device of a supported GPU,
   and respond with the correct iCX3 version information */
int find_blackwell_gpus(struct card_info *infos, int num_gpus, int max_gpus)
{
	char pci_devices_path[PATH_MAX];
	char device_path[PATH_MAX];
	char dev_file[PATH_MAX];
	char pci_path[PATH_MAX];
	char *pci_addr;

	DIR *dir;
	struct dirent *ent;
	
	unsigned short pci_vendor, pci_device = 0;
	
	/* Start looking for PCI */
	strcpy(pci_devices_path, "/sys/bus/pci/devices/");
	dir = opendir(pci_devices_path);
	
	/* make sure we can open the PCI device */
	if(dir == NULL)
		return num_gpus;
	
	/* loop over all PCI devices */
	while((ent = readdir(dir)) != NULL)
	{				
		strcpy(device_path, pci_devices_path);
		strcat(device_path, ent->d_name);
						
		pci_vendor = read_pci_id(device_path, "/vendor");
		pci_device = read_pci_id(device_path, "/device");
		
		/* See if it's a supported Blackwell card*/
		for (int i = 0; i < (sizeof(blackwell_card_list) / sizeof(struct blackwell_card_info)); i++) {
			
			if (pci_vendor == blackwell_card_list[i].vendor_id &&
			    pci_device == blackwell_card_list[i].device_id) {
				
				infos[num_gpus].card_name     = blackwell_card_list[i].card_name;
				infos[num_gpus].pci_id        = strdup(ent->d_name);
				infos[num_gpus].pci_dev_path  = strdup(device_path);
				infos[num_gpus].pci_device_id = pci_device;
				infos[num_gpus].i2c_dev_path  = NULL;
				infos[num_gpus].i2c_fd        = -1;
				
				mmap_bar0(&infos[num_gpus]);
				num_gpus++;
				break;
			}
		}
		
		if (num_gpus == max_gpus)
			break;
	}
	
	return num_gpus;
}

void print_blackwell_info(int compact, int overwrite, struct card_info *gpu)
{
	if (gpu->bar0_mmap == NULL)
		return;
	
	for (int i = 0; i < (sizeof(blackwell_card_list) / sizeof(struct blackwell_card_info)); i++) {
		if (!compact)
			printf("%s @ %s\n", gpu->card_name, gpu->pci_id);
		float temp;
		
		if (blackwell_card_list[i].device_id == gpu->pci_device_id) {
			/* Print die sensors */
			if (compact)
				printf("%sDIE%s", HEADER_COLOR_START, HEADER_COLOR_END);
			for (int die_sens_num = 0; die_sens_num < BLACKWELL_NUM_DIE_SENSORS; die_sens_num++) {
				temp = read_temp_sensor(gpu->bar0_mmap + blackwell_card_list[i].die_temp_offset[die_sens_num]);
				if (compact)
					printf(" %3.0f", temp);
				else
					printf("Die %d: %+.1f°C\n", die_sens_num, temp);
			}
			
			/* Sys temp */
			temp = read_temp_sensor(gpu->bar0_mmap + blackwell_card_list[i].sys_temp_offset);
			if (compact)
				printf(" %sSYS%s %3.0f", HEADER_COLOR_START, temp, HEADER_COLOR_END);
			else
				printf("System: %+.1f°C\n", temp);
			
			/* DRAM temps */
			int dram_num = 0;
			int num_channels = blackwell_card_list[i].clamshell ? 4 : 2;
			int dram_max = 0;
			int dram_min = 120;
			int dram_total = 0.0;
			for (int fbpa_num = 0; fbpa_num < BLACKWELL_MAX_NUM_FBPA; fbpa_num++) {
				for (int chan_num = 0; chan_num < num_channels; chan_num++) {
					temp = read_dram_temp_sensor(gpu->bar0_mmap + 
						blackwell_card_list[i].dram_temp_offset_start + 
						BLACKWELL_FBPA_STRIDE * fbpa_num +
						0x4 * chan_num + 
						2);
						
					if (temp == 0.0)
						break;
					
					dram_max = (temp > dram_max) ? temp : dram_max;
					dram_min = (temp < dram_min) ? temp : dram_min;
					dram_total += temp;
					
					if (!compact)
						printf("DRAM %i: %+.0f°C\n", dram_num, temp);
					dram_num++;
				}
			}
			if (compact) {
				printf(" %sDRAM HI%s %3.0f %sAVG%s %5.1f %sLO%s %3.0f",
					HEADER_COLOR_START,
					HEADER_COLOR_END,
					(float)dram_max,
					HEADER_COLOR_START,
					HEADER_COLOR_END,
					(float)dram_total/dram_num,
					HEADER_COLOR_START,
					HEADER_COLOR_END,
					(float)dram_min);
			} else {
				printf("DRAM average: %+.1f°C\n", (float)dram_total/dram_num);
				printf("DRAM min: %+.0f°C\n", (float)dram_min);
				printf("DRAM max: %+.0f°C\n", (float)dram_max);
			}
			return;
		}
	}
}

float read_temp_sensor(void *addr)
{
	return (*(uint16_t *)addr) / (float)0x100;
}

int read_dram_temp_sensor(void *addr)
{
	uint16_t raw = *(uint16_t *)addr;
	if (raw == 0xBADF)
		return 0.0;
	
	int code = (raw >> 8) & 0xFF;
	if (code > 80) code = 80;
	return (code > 19) ? (code - 20) * 2 : -(40 - code * 2);
}