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

static uint32_t compute_way_mask(unsigned num_ways, unsigned start_way)
{
    uint32_t mask = 0;
    for (unsigned i = 0; i < num_ways; i++)
        mask |= (1U << (start_way + i));
    return mask;
}

static unsigned find_start_way(uint32_t default_mask, unsigned num_ways)
{
    for (unsigned i = 0; i <= 32 - num_ways; i++) {
        uint32_t candidate = compute_way_mask(num_ways, i);
        if ((candidate & default_mask) == candidate)
            return i;
    }
    return 0;
}

static void init_pointer_chase(volatile char *ptr, size_t size)
{
    size_t num_lines = size / CACHE_LINE_SIZE;
    volatile size_t *indices = (volatile size_t *)ptr;

    for (size_t i = 0; i < num_lines; i++)
        indices[i * (CACHE_LINE_SIZE / sizeof(size_t))] = ((i + 1) % num_lines) * CACHE_LINE_SIZE;

    for (size_t i = num_lines - 1; i > 0; i--) {
        size_t j = (size_t)rdtsc_start() % (i + 1);
        size_t idx_i = i * (CACHE_LINE_SIZE / sizeof(size_t));
        size_t idx_j = j * (CACHE_LINE_SIZE / sizeof(size_t));
        size_t tmp = indices[idx_i];
        indices[idx_i] = indices[idx_j];
        indices[idx_j] = tmp;
    }
}

static void prime_temporal_sequential(volatile char *ptr, size_t size)
{
    for (size_t offset = 0; offset < size; offset += CACHE_LINE_SIZE)
        (void)ptr[offset];
}

static void prime_temporal_reverse(volatile char *ptr, size_t size)
{
    for (size_t offset = size; offset > 0; offset -= CACHE_LINE_SIZE)
        (void)ptr[offset - CACHE_LINE_SIZE];
}

static void prime_temporal_strided(volatile char *ptr, size_t size, size_t stride)
{
    size_t effective_stride = stride ? stride : (CACHE_LINE_SIZE * 8);
    for (size_t pass = 0; pass < effective_stride; pass += CACHE_LINE_SIZE) {
        for (size_t offset = pass; offset < size; offset += effective_stride)
            (void)ptr[offset];
    }
}

static void prime_temporal_chase(volatile char *ptr, size_t size)
{
    volatile size_t *indices = (volatile size_t *)ptr;
    size_t num_lines = size / CACHE_LINE_SIZE;
    size_t offset = 0;

    for (size_t i = 0; i < num_lines; i++) {
        offset = indices[offset / sizeof(size_t)];
    }

    (void)offset;
}

static void prime_prefetcht2_sequential(volatile char *ptr, size_t size)
{
    for (size_t offset = 0; offset < size; offset += CACHE_LINE_SIZE)
        _mm_prefetch((const char *)&ptr[offset], _MM_HINT_T2);
    _mm_mfence();
    for (size_t offset = 0; offset < size; offset += CACHE_LINE_SIZE)
        (void)ptr[offset];
}

static void prime_prefetcht2_reverse(volatile char *ptr, size_t size)
{
    for (size_t offset = size; offset > 0; offset -= CACHE_LINE_SIZE)
        _mm_prefetch((const char *)&ptr[offset - CACHE_LINE_SIZE], _MM_HINT_T2);
    _mm_mfence();
    for (size_t offset = size; offset > 0; offset -= CACHE_LINE_SIZE)
        (void)ptr[offset - CACHE_LINE_SIZE];
}

static void prime_prefetcht2_strided(volatile char *ptr, size_t size, size_t stride)
{
    size_t effective_stride = stride ? stride : (CACHE_LINE_SIZE * 8);
    for (size_t pass = 0; pass < effective_stride; pass += CACHE_LINE_SIZE) {
        for (size_t offset = pass; offset < size; offset += effective_stride)
            _mm_prefetch((const char *)&ptr[offset], _MM_HINT_T2);
    }
    _mm_mfence();
    for (size_t pass = 0; pass < effective_stride; pass += CACHE_LINE_SIZE) {
        for (size_t offset = pass; offset < size; offset += effective_stride)
            (void)ptr[offset];
    }
}

static void prime_prefetchnta_sequential(volatile char *ptr, size_t size)
{
    for (size_t offset = 0; offset < size; offset += CACHE_LINE_SIZE)
        _mm_prefetch((const char *)&ptr[offset], _MM_HINT_NTA);
    _mm_mfence();
    for (size_t offset = 0; offset < size; offset += CACHE_LINE_SIZE)
        (void)ptr[offset];
}

static void prime_nt_store_sequential(volatile char *ptr, size_t size)
{
    __m128i zero = _mm_setzero_si128();
    for (size_t offset = 0; offset < size; offset += CACHE_LINE_SIZE) {
        _mm_stream_si128((__m128i *)&ptr[offset], zero);
        _mm_stream_si128((__m128i *)&ptr[offset + 16], zero);
        _mm_stream_si128((__m128i *)&ptr[offset + 32], zero);
        _mm_stream_si128((__m128i *)&ptr[offset + 48], zero);
    }
    _mm_sfence();
    for (size_t offset = 0; offset < size; offset += CACHE_LINE_SIZE)
        (void)ptr[offset];
}

void prime_region(volatile char *ptr, size_t size, unsigned iterations,
                  icepick_prime_strategy_t strategy, icepick_access_pattern_t pattern,
                  size_t stride)
{
    if (pattern == ICEPICK_PATTERN_POINTER_CHASE)
        init_pointer_chase(ptr, size);

    for (unsigned iter = 0; iter < iterations; iter++) {
        switch (strategy) {
        case ICEPICK_PRIME_TEMPORAL:
            switch (pattern) {
            case ICEPICK_PATTERN_SEQUENTIAL:
                prime_temporal_sequential(ptr, size);
                break;
            case ICEPICK_PATTERN_REVERSE:
                prime_temporal_reverse(ptr, size);
                break;
            case ICEPICK_PATTERN_STRIDED:
                prime_temporal_strided(ptr, size, stride);
                break;
            case ICEPICK_PATTERN_POINTER_CHASE:
                prime_temporal_chase(ptr, size);
                break;
            }
            break;

        case ICEPICK_PRIME_PREFETCHT2:
            switch (pattern) {
            case ICEPICK_PATTERN_SEQUENTIAL:
                prime_prefetcht2_sequential(ptr, size);
                break;
            case ICEPICK_PATTERN_REVERSE:
                prime_prefetcht2_reverse(ptr, size);
                break;
            case ICEPICK_PATTERN_STRIDED:
                prime_prefetcht2_strided(ptr, size, stride);
                break;
            case ICEPICK_PATTERN_POINTER_CHASE:
                prime_prefetcht2_sequential(ptr, size);
                prime_temporal_chase(ptr, size);
                break;
            }
            break;

        case ICEPICK_PRIME_PREFETCHNTA:
            prime_prefetchnta_sequential(ptr, size);
            break;

        case ICEPICK_PRIME_NT_STORE:
            prime_nt_store_sequential(ptr, size);
            break;
        }
    }
}

int icepick_lock(icepick_topology_t *topo, const icepick_config_t *cfg,
                 icepick_region_t **region)
{
    if (!topo || !cfg || !region)
        return ICEPICK_E_INVALID;

    if (cfg->clos_id == 0 || cfg->clos_id >= topo->max_clos)
        return ICEPICK_E_INVALID;

    unsigned ways_needed = (cfg->size + topo->way_size - 1) / topo->way_size;
    if (ways_needed == 0)
        ways_needed = 1;

    if (ways_needed >= topo->l3_ways)
        return ICEPICK_E_TOO_LARGE;

    int ret = clos_init(topo);
    if (ret < 0)
        return ret;

    unsigned start_way = find_start_way(topo->default_way_mask, ways_needed);
    uint32_t way_mask = compute_way_mask(ways_needed, start_way);

    ret = clos_allocate(cfg->clos_id, way_mask);
    if (ret < 0)
        return ret;

    icepick_region_t *r = calloc(1, sizeof(*r));
    if (!r) {
        clos_release(cfg->clos_id);
        return ICEPICK_E_ALLOC;
    }

    r->clos_id = cfg->clos_id;
    r->numa_node = cfg->numa_node;
    r->way_mask = way_mask;
    r->mba_throttle = cfg->mba_throttle;
    r->topo = topo;
    r->size = ways_needed * topo->way_size;
    r->monitor = NULL;
    r->prime_strategy = cfg->prime_strategy;
    r->access_pattern = cfg->access_pattern;
    r->stride_bytes = cfg->stride_bytes;
    r->prime_iterations = cfg->prime_iterations ? cfg->prime_iterations : 3;

    ret = region_alloc(r->size, cfg->numa_node, cfg->huge_pages, &r->ptr);
    if (ret < 0) {
        clos_release(cfg->clos_id);
        free(r);
        return ret;
    }

    int cpu = sched_getcpu();
    if (cpu < 0)
        cpu = 0;

    cpu_set_t old_affinity, new_affinity;
    CPU_ZERO(&new_affinity);
    CPU_SET(cpu, &new_affinity);
    sched_getaffinity(0, sizeof(old_affinity), &old_affinity);
    sched_setaffinity(0, sizeof(new_affinity), &new_affinity);

    int msr_fd = msr_open(cpu);
    if (msr_fd < 0) {
        sched_setaffinity(0, sizeof(old_affinity), &old_affinity);
        region_free(r->ptr, r->size);
        clos_release(cfg->clos_id);
        free(r);
        return msr_fd;
    }

    ret = clos_configure_mask(msr_fd, 0, clos_table[0].way_mask);
    if (ret < 0)
        goto cleanup;

    ret = clos_configure_mask(msr_fd, cfg->clos_id, way_mask);
    if (ret < 0)
        goto cleanup;

    if (topo->mba_supported && cfg->mba_throttle > 0) {
        ret = mba_configure(msr_fd, cfg->clos_id, cfg->mba_throttle);
        if (ret < 0)
            goto cleanup;
    }

    ret = clos_associate_thread(msr_fd, cfg->clos_id);
    if (ret < 0)
        goto cleanup;

    prime_region(r->ptr, r->size, r->prime_iterations,
                 r->prime_strategy, r->access_pattern, r->stride_bytes);

    clos_associate_thread(msr_fd, 0);

    msr_close(msr_fd);
    sched_setaffinity(0, sizeof(old_affinity), &old_affinity);

    if (cfg->verify) {
        icepick_latency_stats_t stats;
        ret = icepick_verify(r, &stats);
        if (ret < 0) {
            icepick_unlock(r);
            return ret;
        }
    }

    if (cfg->auto_monitor) {
        icepick_monitor_t *mon;
        ret = icepick_monitor_start_ex(r, cfg->pmu_poll_interval_ns,
                                       cfg->probe_interval_ns,
                                       cfg->miss_threshold, &mon);
        if (ret < 0) {
            icepick_unlock(r);
            return ret;
        }
        r->monitor = mon;
    }

    *region = r;
    return 0;

cleanup:
    msr_close(msr_fd);
    sched_setaffinity(0, sizeof(old_affinity), &old_affinity);
    region_free(r->ptr, r->size);
    clos_release(cfg->clos_id);
    free(r);
    return ret;
}

int icepick_unlock(icepick_region_t *region)
{
    if (!region)
        return ICEPICK_E_INVALID;

    if (region->monitor)
        icepick_monitor_stop(region->monitor);

    int cpu = sched_getcpu();
    if (cpu < 0)
        cpu = 0;

    int msr_fd = msr_open(cpu);
    if (msr_fd >= 0) {
        clos_release(region->clos_id);
        clos_configure_mask(msr_fd, region->clos_id, 0);
        if (region->topo->mba_supported && region->mba_throttle > 0)
            mba_configure(msr_fd, region->clos_id, 0);
        clos_configure_mask(msr_fd, 0, clos_get_default_mask());
        msr_close(msr_fd);
    }

    region_free(region->ptr, region->size);
    free(region);

    return 0;
}

void *icepick_region_ptr(const icepick_region_t *region)
{
    return region ? region->ptr : NULL;
}

size_t icepick_region_size(const icepick_region_t *region)
{
    return region ? region->size : 0;
}

unsigned icepick_region_clos(const icepick_region_t *region)
{
    return region ? region->clos_id : 0;
}

const char *icepick_strerror(int err)
{
    switch (err) {
    case 0:                     return "success";
    case ICEPICK_E_NO_CAT:      return "cat not supported on this cpu";
    case ICEPICK_E_PERMISSION:  return "permission denied (need root or cap_sys_rawio)";
    case ICEPICK_E_NO_CLOS:     return "no clos available or already allocated";
    case ICEPICK_E_TOO_LARGE:   return "requested size exceeds available cache ways";
    case ICEPICK_E_NUMA:        return "invalid numa node or binding failed";
    case ICEPICK_E_HUGEPAGE:    return "huge page allocation failed";
    case ICEPICK_E_VERIFY:      return "verification failed - data may not be cache-resident";
    case ICEPICK_E_INVALID:     return "invalid argument";
    case ICEPICK_E_ALLOC:       return "memory allocation failed";
    case ICEPICK_E_MSR:         return "msr operation failed";
    case ICEPICK_E_THREAD:      return "monitor thread creation failed";
    default:                    return "unknown error";
    }
}