#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdint.h>
#include <string.h>
#include <dirent.h>
#include <cpuid.h>
#include <math.h>

#define AMD_STRING "AuthenticAMD"
#define ZEN_FAMILY 0x17
#define ZEN3_FAMILY 0x19

#define MESUREMENT_TIME 0.1

static unsigned int cores = 0;
static double energy_unit = 0;
static struct cpudev *cpu_dev_ids;

static int *msr_files = NULL;

static unsigned long package_eng_b = 0;
static unsigned long package_eng_a = 0;
static unsigned long *core_eng_b = NULL;
static unsigned long *core_eng_a = NULL;

float package_power;
float *core_power;
float *core_fid;

struct cpudev {
	short coreid;
	short cpuid;
};


static int check_zen() {
    unsigned int eax = 0, ebx = 0, ecx = 0, edx = 0, ext_family;
    char vendor[13];

    __get_cpuid(0, &eax, &ebx, &ecx, &edx);

    memcpy(vendor, &ebx, 4);
    memcpy(vendor+4, &edx, 4);
    memcpy(vendor+8, &ecx, 4);
    vendor[12] = 0;

    if (strcmp(vendor, AMD_STRING) != 0){
        return 0;
    }

    __get_cpuid(1, &eax, &ebx, &ecx, &edx);

    ext_family = ((eax >> 8) & 0xF) + ((eax >> 20) & 0xFF);
    if (ext_family != ZEN_FAMILY && ext_family != ZEN3_FAMILY){
        return 0;
    }

    return 1;
}

static unsigned int get_core_count() {
    unsigned int eax = 0, ebx = 0, ecx = 0, edx = 0;
    unsigned int logical_cpus, threads_per_code;

    // AMD PPR: page 57 - CPUID_Fn00000001_EBX
    __get_cpuid(1, &eax, &ebx, &ecx, &edx);
    logical_cpus = (ebx >> 16) & 0xFF;

    // AMD PPR: page 82 - CPUID_Fn8000001E_EBX
    __get_cpuid(0x8000001E, &eax, &ebx, &ecx, &edx);
    threads_per_code = ((ebx >> 8) & 0xF) + 1;

    if (threads_per_code == 0)
        return logical_cpus;

    return logical_cpus / threads_per_code;
}

static struct cpudev* get_cpu_dev_ids() {
    unsigned int num_cores = get_core_count();
    struct cpudev* cores = malloc(num_cores * sizeof(struct cpudev));

    for (int i = 0; i < num_cores; ++i) {
        cores[i].coreid = i;
        cores[i].cpuid = -1;
    }

    DIR* d;
    struct dirent* dir;
    d = opendir("/sys/devices/system/cpu");

    if (d) {
        while ((dir = readdir(d)) != NULL) {
            if (strncmp(dir->d_name, "cpu", 3) == 0 && atoi(dir->d_name + 3) > 0) {
                char path[256];
				int core_id;
				FILE* f;
				
				snprintf(path, sizeof(path), "/sys/devices/system/cpu/%s/topology/core_id", dir->d_name);
				f = fopen(path, "r");
				if (f) {
					fscanf(f, "%d", &core_id);
					fclose(f);
				}
				
                snprintf(path, sizeof(path), "/sys/devices/system/cpu/%s/topology/thread_siblings_list", dir->d_name);
                f = fopen(path, "r");

                if (f) {
                    int cpuid, sibling;

                    if (fscanf(f, "%d,%d", &cpuid, &sibling) == 2) {
                        // The lower numbered CPU is considered the representative of the core
                        if (cores[core_id].cpuid == -1 || cores[core_id].cpuid > cpuid) {
                            cores[core_id].cpuid = cpuid;
                        }
                    }
                    fclose(f);
                }
            }
        }
        closedir(d);
    }

    return cores;
}

static int open_msr(short devid) {
    char msr_path[20];
    sprintf(msr_path, "/dev/cpu/%d/msr", devid);
    return open(msr_path, O_RDONLY);
}

static int read_msr(int file, unsigned int index, unsigned long long *data) {
    if (file < 0)
        return 0;

    return pread(file, data, sizeof *data, index) == sizeof *data;
}

static double get_energy_unit() {
    unsigned long long data;
    // AMD OSRR: page 139 - MSRC001_0299
    if (!read_msr(msr_files[0], 0xC0010299, &data))
        return 0.0;

    return pow(1.0/2.0, (double)((data >> 8) & 0x1F));
}

static unsigned long get_package_energy() {
    unsigned long long data;
    // AMD OSRR: page 139 - MSRC001_029B
    if (!read_msr(msr_files[0], 0xC001029B, &data))
        return 0;

    return data;
}

static unsigned long get_core_energy(int core) {
    unsigned long long data;
    // AMD OSRR: page 139 - MSRC001_029A
    if (!read_msr(msr_files[core], 0xC001029A, &data))
        return 0;

    return data;
}

static double get_core_fid(int core) {
    double ratio;
    unsigned long long data;

    // By reverse-engineering Ryzen Master, we know that
    //  this undocumented MSR is responsible for returning
    //  the FID and FDID for the core used for calculating the
    //  effective frequency.
    //
    // The FID is returned in bits [8:0]
    // The FDID is returned in bits [14:8]
    if (!read_msr(msr_files[core], 0xC0010293, &data))
        return 0;

    ratio = (double)(data & 0xff) / (double)((data >> 8) & 0x3F);

    // The effective ratio is based on increments of 200 MHz.
    return ratio * 200.0 / 1000.0;
}

static int msr_init() {
    unsigned int i;

    if (!check_zen())
        return 0;

    cores = get_core_count();
    if (cores == 0)
        return 0;

    cpu_dev_ids = get_cpu_dev_ids();
    msr_files = malloc(cores * sizeof (int));
    for (i = 0; i < cores; i++) {
        msr_files[i] = open_msr(cpu_dev_ids[i].cpuid);
		if (msr_files[i] < 0)
			return 0;
    }

    energy_unit = get_energy_unit();
    if (energy_unit == 0)
        return 0;

    core_eng_b = malloc(cores * sizeof (unsigned long));
    core_eng_a = malloc(cores * sizeof (unsigned long));
    core_power = malloc(cores * sizeof (float));
    core_fid = malloc(cores * sizeof (float));
    
    /*msr_update();*/

    return 1;
}

static void msr_update() {
    int i;

    package_eng_b = get_package_energy();
    for (i = 0; i < cores; i++) {
        core_eng_b[i] = get_core_energy(i);
    }

    usleep(MESUREMENT_TIME*1000000);

    package_eng_a = get_package_energy();
    for (i = 0; i < cores; i++) {
        core_eng_a[i] = get_core_energy(i);
    }

    if (package_eng_a >= package_eng_b) {
        package_power = (package_eng_a - package_eng_b) * energy_unit / MESUREMENT_TIME;
    }

    for (i = 0; i < cores; i++) {
        if (core_eng_a[i] >= core_eng_b[i]) {
            core_power[i] = (core_eng_a[i] - core_eng_b[i]) * energy_unit / MESUREMENT_TIME;
        }

        core_fid[i] = get_core_fid(i);
    }
}


int main(int argc, char *argv[]) {
	int init_ok, i;
	init_ok = msr_init();
	
	if (init_ok) {
		msr_update();
	} else {
		printf("Error reading RAPL (running average power limit) sensors.\nAre you root? Did you `modprobe msr`? Are you on AMD Zen?\n");
		return 1;
	}

	printf("zen3-rapl\n");
	printf("%-12s   %6.2f W\n", "Package:", package_power);
	char core_name[256];
	/* Print core powers */
	for (i = 0; i < cores; i++) {
		sprintf(core_name, "Core_%d:", i);
		printf("%-12s   %6.2f W\n", core_name, i, core_power[i]);
	}
	
	/* Print core effective frequencies */
	for (i = 0; i < cores; i++) {
		sprintf(core_name, "Core_%d_eff:", i);
		printf("%-12s %6.3f GHz\n", core_name, i, core_fid[i]);
	}
    return 0;
}