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

static uintptr_t next_virtual = 0x10000000;

uintptr_t compute_physical_base(size_t size) {
    // 1. Determine the required alignment for this size in CHERI
    size_t mask = cheri_representable_alignment_mask(size);
    
    // 2. Round up the current next_virtual to the required alignment
    uintptr_t base = (next_virtual + ~mask) & mask;
    
    // 3. Ensure the length itself is representable
    size_t representable_len = cheri_representable_length(size);
    
    // 4. Update the global pointer for the next call
    next_virtual = base + representable_len;
    
    return base;
}


// 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
// Malloc wrapper
// void * __capability malloc(size_t size) {

//     uintptr_t raw = (uintptr_t)tiny_malloc(size);

//     // void *__capability cap = (void *__capability)raw;
//     void *__capability cap = cheri_ddc_get();

//     // // Set address from raw pointer
//     cap = cheri_address_set(cap, raw);

//     int delta = 12;

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

//     cap = add_delta(cap, delta);

//     size_t aligned = cheri_representable_length(size);

//     cap = cheri_bounds_set(cap, aligned); 

//     // // 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 * __capability malloc(size_t size) {
    void * phys_ptr = tiny_malloc(size);
    if (!phys_ptr) return NULL;

    uintptr_t raw_phys = (uintptr_t)phys_ptr;
    intptr_t p_base = compute_physical_base(size);
    intptr_t delta = p_base - (intptr_t)raw_phys;

    void * __capability cap = cheri_ddc_get();
    cap = cheri_address_set(cap, raw_phys);

    cap = add_delta(cap, (int)delta);

    size_t aligned_size = cheri_representable_length(size);
    cap = cheri_bounds_set(cap, aligned_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;
// }