#include <cheriintrin.h>
#include <stdint.h>
#include <stddef.h>

#include "riscv_test.h"

// TODO: Write track allocated and free memory avaliable
// Also to calculate the distance between the virtual and 
// physical memory. 

// [1, 2, 3. ....... n] -> physical 
//  |__| (Distance -1 between address)
//  |  | (Distance 1 between address)
// [1, 2, 3. ....... n] -> virtual 

#define HEAP_SIZE 65536  // 64 KB heap

#define PHYS_BASE 0x80000000
#define VIRT_BASE 0xFFFFFFFF80000000

static char heap[HEAP_SIZE];
static size_t bump = 0;

static void * __capability free_list = NULL;

typedef struct free_node {
    void * __capability next;
} free_node_t;



// Add delta value for TLB translation
static inline void * __capability add_delta(void * __capability cap, int offset) {
    void * __capability result;

    asm volatile (
        "candperm %0, %1, %2"
        : "=C" (result)  // Output: %0 (result)
        : "C" (cap),     // Input:  %1 (original cap)
          "r" (offset)   // Input:  %2 (offset register)
        :                // No clobbered registers
    );

    return result;
}


void * __capability malloc(size_t size) {

    // Align to 8 bytes (important for capability safety)
    size = (size + 7) & ~7;

    if (bump + size > HEAP_SIZE)
        return NULL;

    void * __capability base = cheri_ddc_get();

    uintptr_t addr = (uintptr_t)(heap + bump);

    // Create capability to this region
    void * __capability cap = cheri_address_set(base, addr);

    // Enforce bounds (this is the key CHERI feature)
    cap = cheri_bounds_set(cap, size);

    // Hard-coded delta value 
    cap = add_delta(cap, 10);

    bump += size;

    return cap;
}

void bump_reset(void) {
    bump = 0;
}

// Can only free if it's in a stack 
// Need to write a free †hat can work 
// in any sequence 
// In regular malloc the freelist 
// is maintained from mmap. 
void free(void * __capability ptr) {
    if (!ptr) return;

    uintptr_t addr = cheri_address_get(ptr);
    size_t size = cheri_length_get(ptr);

    // Check if this is the most recent allocation
    if ((char *)addr + size == heap + bump) {
        bump -= size;
    }

    // *ptr = NULL;
}

// Quick tests
// int main(void) {
//     char * __capability a = malloc(30);
//     char * __capability b = malloc(16);


//     // if (!a || !b) return -1;

//     a[0] = 'A';
//     a[1] = 'C';
//     b[0] = 'B';

//     // This will fault (out-of-bounds)
//     // a[20] = 'X';

//     if (a[1] != 'C') {
//         while (1);
//     }

//     if (b[0] != 'B') {
//         while (1);
//     }

//     free(&b);
//     b = NULL;

//     if (b[0] == 'B') {
//         while (1);
//     }

//     char * __capability c = malloc(16);

//     // if (c[0] != 'B') {
//     //     while (1);
//     // }


//     return 0;
// }