#ifndef ICEPICK_INTERNAL_H
#define ICEPICK_INTERNAL_H

#include "../include/icepick.h"
#include <stdint.h>
#include <stdatomic.h>
#include <pthread.h>

#define MSR_IA32_PQR_ASSOC       0xC8F
#define MSR_IA32_L3_QOS_MASK_0   0xC90
#define MSR_IA32_L2_QOS_CFG      0xC82
#define MSR_IA32_MBA_THRTL_BASE  0xD50

#define MPAM_PARTID_BITS         16
#define MPAM_MAX_PARTID          ((1 << MPAM_PARTID_BITS) - 1)

#define CACHE_LINE_SIZE  64
#define MAX_CLOS         16
#define MAX_CCXS         16
#define MAX_CORE_TYPES   2

typedef enum {
    ARCH_X86_64,
    ARCH_ARM64,
    ARCH_UNKNOWN
} arch_type_t;

#define MONITOR_DEFAULT_PMU_POLL_NS   100000
#define MONITOR_DEFAULT_PROBE_NS      10000
#define MONITOR_DEFAULT_MISS_THRESH   16
#define MONITOR_PROBE_CONFIRM_COUNT   8
#define MONITOR_L3_LATENCY_THRESH_NS  50

struct icepick_monitor;

struct icepick_topology {
    arch_type_t arch;
    bool cat_supported;
    bool mba_supported;
    bool mba_is_linear;
    bool is_hybrid;
    bool mpam_supported;
    unsigned l3_ways;
    size_t l3_size;
    size_t way_size;
    unsigned max_clos;
    unsigned max_mba_thrtl;
    unsigned max_partid;
    unsigned ccx_count;
    uint32_t default_way_mask;
    int ccx_numa_node[MAX_CCXS];
    unsigned pcore_count;
    unsigned ecore_count;
    int pcore_cpus[64];
    int ecore_cpus[256];
};

struct icepick_region {
    void *ptr;
    size_t size;
    unsigned clos_id;
    int numa_node;
    uint32_t way_mask;
    unsigned mba_throttle;
    icepick_core_type_t core_type;
    icepick_topology_t *topo;
    struct icepick_monitor *monitor;
    icepick_prime_strategy_t prime_strategy;
    icepick_access_pattern_t access_pattern;
    size_t stride_bytes;
    unsigned prime_iterations;
};

struct icepick_monitor {
    icepick_region_t *region;
    pthread_t thread;
    int perf_fd;
    int msr_fd;
    int cpu;
    atomic_bool running;
    atomic_uint_fast64_t eviction_count;
    atomic_uint_fast64_t reprime_count;
    uint64_t pmu_poll_interval_ns;
    uint64_t probe_interval_ns;
    uint32_t miss_threshold;
    bool degraded;
};

typedef struct {
    bool allocated;
    uint32_t way_mask;
} clos_entry_t;

extern clos_entry_t clos_table[MAX_CLOS];

int msr_open(int cpu);
void msr_close(int fd);
int msr_read(int fd, uint32_t msr, uint64_t *value);
int msr_write(int fd, uint32_t msr, uint64_t value);

int clos_init(icepick_topology_t *topo);
int clos_allocate(unsigned clos_id, uint32_t way_mask);
int clos_release(unsigned clos_id);
int clos_configure_mask(int msr_fd, unsigned clos_id, uint32_t way_mask);
int clos_associate_thread(int msr_fd, unsigned clos_id);
uint32_t clos_get_default_mask(void);

int region_alloc(size_t size, int numa_node, bool huge_pages, void **out_ptr);
void region_free(void *ptr, size_t size);

uint64_t rdtsc_start(void);
uint64_t rdtsc_end(void);
uint64_t cycles_to_ns(uint64_t cycles);

int bench_compare(icepick_topology_t *topo, size_t size, size_t iterations);

void prime_region(volatile char *ptr, size_t size, unsigned iterations,
                  icepick_prime_strategy_t strategy, icepick_access_pattern_t pattern,
                  size_t stride);
int monitor_perf_open(int cpu, bool is_amd);

int mba_configure(int msr_fd, unsigned clos_id, unsigned throttle);
int find_cpu_for_core_type(const icepick_topology_t *topo, icepick_core_type_t core_type,
                            int numa_node);

arch_type_t detect_arch(void);
int mpam_configure_partition(unsigned partid, uint32_t cpbm, unsigned mbw_max);
int mpam_associate_thread(unsigned partid);

#endif