#define _GNU_SOURCE
#include "internal.h"
#include <stdlib.h>
#include <string.h>
#include <sched.h>
#include <dirent.h>
#include <stdio.h>
#include <unistd.h>

#ifdef __x86_64__
#include <cpuid.h>

static void cpuid(uint32_t leaf, uint32_t subleaf,
                  uint32_t *eax, uint32_t *ebx, uint32_t *ecx, uint32_t *edx)
{
    __cpuid_count(leaf, subleaf, *eax, *ebx, *ecx, *edx);
}
#endif

arch_type_t detect_arch(void)
{
#ifdef __x86_64__
    return ARCH_X86_64;
#elif defined(__aarch64__)
    return ARCH_ARM64;
#else
    return ARCH_UNKNOWN;
#endif
}

#ifdef __x86_64__
static bool detect_cat_support(icepick_topology_t *topo)
{
    uint32_t eax, ebx, ecx, edx;

    cpuid(0x10, 0, &eax, &ebx, &ecx, &edx);
    if (!(ebx & (1 << 1)))
        return false;

    cpuid(0x10, 1, &eax, &ebx, &ecx, &edx);

    topo->l3_ways = (eax & 0x1F) + 1;
    topo->max_clos = (edx & 0xFFFF) + 1;
    topo->default_way_mask = ebx;

    if (topo->max_clos > MAX_CLOS)
        topo->max_clos = MAX_CLOS;

    return true;
}

static bool detect_mba_support(icepick_topology_t *topo)
{
    uint32_t eax, ebx, ecx, edx;

    cpuid(0x10, 0, &eax, &ebx, &ecx, &edx);
    if (!(ebx & (1 << 3)))
        return false;

    cpuid(0x10, 3, &eax, &ebx, &ecx, &edx);

    topo->max_mba_thrtl = (eax & 0xFFF) + 1;
    topo->mba_is_linear = (ecx & (1 << 2)) != 0;

    return true;
}

static void detect_amd_cache_topology(icepick_topology_t *topo)
{
    uint32_t eax, ebx, ecx, edx;

    for (int i = 0; i < 16; i++) {
        cpuid(0x8000001D, i, &eax, &ebx, &ecx, &edx);

        uint32_t cache_type = eax & 0x1F;
        if (cache_type == 0)
            break;

        uint32_t cache_level = (eax >> 5) & 0x7;
        if (cache_level != 3)
            continue;

        uint32_t line_size = (ebx & 0xFFF) + 1;
        uint32_t partitions = ((ebx >> 12) & 0x3FF) + 1;
        uint32_t ways = ((ebx >> 22) & 0x3FF) + 1;
        uint32_t sets = ecx + 1;

        topo->l3_size = (size_t)line_size * partitions * ways * sets;
        topo->l3_ways = ways;
        topo->way_size = topo->l3_size / topo->l3_ways;
        break;
    }
}

static void detect_intel_cache_topology(icepick_topology_t *topo)
{
    uint32_t eax, ebx, ecx, edx;

    for (int i = 0; i < 16; i++) {
        cpuid(0x04, i, &eax, &ebx, &ecx, &edx);

        uint32_t cache_type = eax & 0x1F;
        if (cache_type == 0)
            break;

        uint32_t cache_level = (eax >> 5) & 0x7;
        if (cache_level != 3)
            continue;

        uint32_t line_size = (ebx & 0xFFF) + 1;
        uint32_t partitions = ((ebx >> 12) & 0x3FF) + 1;
        uint32_t ways = ((ebx >> 22) & 0x3FF) + 1;
        uint32_t sets = ecx + 1;

        topo->l3_size = (size_t)line_size * partitions * ways * sets;
        topo->l3_ways = ways;
        topo->way_size = topo->l3_size / topo->l3_ways;
        break;
    }
}

static bool is_amd_cpu(void)
{
    uint32_t eax, ebx, ecx, edx;
    cpuid(0, 0, &eax, &ebx, &ecx, &edx);
    return (ebx == 0x68747541 && edx == 0x69746E65 && ecx == 0x444D4163);
}

static bool detect_hybrid_architecture(icepick_topology_t *topo)
{
    (void)topo;
    uint32_t eax, ebx, ecx, edx;

    cpuid(0x00, 0, &eax, &ebx, &ecx, &edx);
    if (eax < 0x1A)
        return false;

    cpuid(0x07, 0, &eax, &ebx, &ecx, &edx);
    if (!(edx & (1 << 15)))
        return false;

    return true;
}

static void enumerate_hybrid_cores(icepick_topology_t *topo)
{
    int num_cpus = sysconf(_SC_NPROCESSORS_ONLN);
    if (num_cpus <= 0)
        return;

    cpu_set_t original_affinity;
    sched_getaffinity(0, sizeof(original_affinity), &original_affinity);

    topo->pcore_count = 0;
    topo->ecore_count = 0;

    for (int cpu = 0; cpu < num_cpus && cpu < 320; cpu++) {
        cpu_set_t cpuset;
        CPU_ZERO(&cpuset);
        CPU_SET(cpu, &cpuset);

        if (sched_setaffinity(0, sizeof(cpuset), &cpuset) != 0)
            continue;

        uint32_t eax, ebx, ecx, edx;
        cpuid(0x1A, 0, &eax, &ebx, &ecx, &edx);

        uint32_t core_type = (eax >> 24) & 0xFF;

        if (core_type == 0x40 && topo->pcore_count < 64) {
            topo->pcore_cpus[topo->pcore_count++] = cpu;
        } else if (core_type == 0x20 && topo->ecore_count < 256) {
            topo->ecore_cpus[topo->ecore_count++] = cpu;
        }
    }

    sched_setaffinity(0, sizeof(original_affinity), &original_affinity);
}
#endif

#ifdef __aarch64__
static bool detect_mpam_support(icepick_topology_t *topo)
{
    FILE *f = fopen("/sys/devices/system/cpu/cpu0/cache/index3/id", "r");
    if (!f)
        return false;
    fclose(f);

    topo->max_partid = 16;
    return true;
}

static void detect_arm64_cache_topology(icepick_topology_t *topo)
{
    FILE *f = fopen("/sys/devices/system/cpu/cpu0/cache/index3/size", "r");
    if (f) {
        unsigned size_kb = 0;
        if (fscanf(f, "%uK", &size_kb) == 1)
            topo->l3_size = (size_t)size_kb * 1024;
        fclose(f);
    }

    f = fopen("/sys/devices/system/cpu/cpu0/cache/index3/ways_of_associativity", "r");
    if (f) {
        fscanf(f, "%u", &topo->l3_ways);
        fclose(f);
    }

    if (topo->l3_size > 0 && topo->l3_ways > 0)
        topo->way_size = topo->l3_size / topo->l3_ways;

    if (topo->l3_ways > 0)
        topo->default_way_mask = (1U << topo->l3_ways) - 1;
}
#endif

static int count_ccxs(void)
{
    DIR *dir = opendir("/sys/devices/system/node");
    if (!dir)
        return 1;

    int count = 0;
    struct dirent *entry;
    while ((entry = readdir(dir))) {
        if (strncmp(entry->d_name, "node", 4) == 0)
            count++;
    }
    closedir(dir);

    return count > 0 ? count : 1;
}

static void detect_numa_topology(icepick_topology_t *topo)
{
    topo->ccx_count = count_ccxs();
    if (topo->ccx_count > MAX_CCXS)
        topo->ccx_count = MAX_CCXS;

    for (unsigned i = 0; i < topo->ccx_count; i++)
        topo->ccx_numa_node[i] = i;
}

int find_cpu_for_core_type(const icepick_topology_t *topo, icepick_core_type_t core_type,
                            int numa_node)
{
    (void)numa_node;

    if (!topo->is_hybrid || core_type == ICEPICK_CORE_ANY) {
        int cpu = sched_getcpu();
        return cpu >= 0 ? cpu : 0;
    }

    if (core_type == ICEPICK_CORE_PCORE && topo->pcore_count > 0)
        return topo->pcore_cpus[0];

    if (core_type == ICEPICK_CORE_ECORE && topo->ecore_count > 0)
        return topo->ecore_cpus[0];

    return 0;
}

int icepick_discover_topology(icepick_topology_t **topo)
{
    icepick_topology_t *t = calloc(1, sizeof(*t));
    if (!t)
        return ICEPICK_E_ALLOC;

    t->arch = detect_arch();

#ifdef __x86_64__
    if (t->arch == ARCH_X86_64) {
        t->cat_supported = detect_cat_support(t);
        if (!t->cat_supported) {
            free(t);
            return ICEPICK_E_NO_CAT;
        }

        t->mba_supported = detect_mba_support(t);

        bool is_amd = is_amd_cpu();
        if (is_amd)
            detect_amd_cache_topology(t);
        else
            detect_intel_cache_topology(t);

        if (!is_amd) {
            t->is_hybrid = detect_hybrid_architecture(t);
            if (t->is_hybrid)
                enumerate_hybrid_cores(t);
        }
    }
#endif

#ifdef __aarch64__
    if (t->arch == ARCH_ARM64) {
        t->mpam_supported = detect_mpam_support(t);
        if (!t->mpam_supported) {
            free(t);
            return ICEPICK_E_NO_CAT;
        }

        t->cat_supported = t->mpam_supported;
        detect_arm64_cache_topology(t);
        t->max_clos = t->max_partid;
    }
#endif

    if (t->l3_size == 0 || t->l3_ways == 0) {
        t->way_size = 2 * 1024 * 1024;
        t->l3_size = t->way_size * t->l3_ways;
    }

    detect_numa_topology(t);

    *topo = t;
    return 0;
}

void icepick_free_topology(icepick_topology_t *topo)
{
    free(topo);
}

unsigned icepick_topology_l3_ways(const icepick_topology_t *topo)
{
    return topo->l3_ways;
}

size_t icepick_topology_l3_size(const icepick_topology_t *topo)
{
    return topo->l3_size;
}

size_t icepick_topology_way_size(const icepick_topology_t *topo)
{
    return topo->way_size;
}

unsigned icepick_topology_max_clos(const icepick_topology_t *topo)
{
    return topo->max_clos;
}

unsigned icepick_topology_ccx_count(const icepick_topology_t *topo)
{
    return topo->ccx_count;
}

bool icepick_topology_cat_supported(const icepick_topology_t *topo)
{
    return topo->cat_supported;
}

bool icepick_topology_mba_supported(const icepick_topology_t *topo)
{
    return topo->mba_supported;
}

bool icepick_topology_mba_is_linear(const icepick_topology_t *topo)
{
    return topo->mba_is_linear;
}

unsigned icepick_topology_max_mba_thrtl(const icepick_topology_t *topo)
{
    return topo->max_mba_thrtl;
}

bool icepick_topology_is_hybrid(const icepick_topology_t *topo)
{
    return topo->is_hybrid;
}

unsigned icepick_topology_pcore_count(const icepick_topology_t *topo)
{
    return topo->pcore_count;
}

unsigned icepick_topology_ecore_count(const icepick_topology_t *topo)
{
    return topo->ecore_count;
}