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

#include "zen3-rapl.h"

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

#define MEASUREMENT_TIME 0.25

static float energy_unit = 0;

static int msr_file = -1;

static float package_power = 0.0;

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 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 float get_energy_unit()
{
    unsigned long long data;
    // AMD OSRR: page 139 - MSRC001_0299
    if (!read_msr(msr_file, 0xC0010299, &data))
        return 0.0;

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

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

    return data;
}

int init_rapl()
{
	/* Check for supported Zen CPU */
	if (!check_zen())
		return 0;
	
	/* Open MSR file */
	char *msr_path = "/dev/cpu/0/msr";
    msr_file = open(msr_path, O_RDONLY);
	if (msr_file < 0)
		return 0;
	
	/* Energy unit */
	energy_unit = get_energy_unit();
    if (energy_unit == 0)
        return 0;
	
	return 1;
}

float get_rapl_package_power()
{
	unsigned long package_eng_b = 0;
	unsigned long package_eng_a = 0;

    package_eng_b = get_package_energy();

    usleep(MEASUREMENT_TIME*1000000);

    package_eng_a = get_package_energy();

	/* Only update if we computed a good power reading, otherwise return the most recent measurement */
    if (package_eng_a >= package_eng_b)
        package_power = (package_eng_a - package_eng_b) * energy_unit / MEASUREMENT_TIME;
	
	return package_power;
}