#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 

void* tiny_malloc(size_t);
void    tiny_free(void*);

#define ALIGN 16

// #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;
}


// Malloc wrapper
void * __capability malloc(size_t size) {

    void *raw = tiny_malloc(size);

    void *__capability cap = (void *__capability)raw;
    
    cap = cheri_bounds_set(cap, size); 

    int delta = 12;

    delta = (delta + ALIGN - 1) & ~(ALIGN - 1);

    cap = add_delta(cap, delta);

    // // 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;
// }

// Free wrapper
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;
    // }

    // Extract raw address from capability
    void *raw = (void *)cheri_address_get(ptr);

    tiny_free(raw);
}

// 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;
// }