#include "debug_config.h"
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <stdio.h>
#include <stdbool.h>
#include "mem.h"

#ifndef BOUNDARY_CHECK_SIZE
#define BOUNDARY_CHECK_SIZE 256
#endif

#define CANARY 0xDEADBEEF
#define PADDING_FILL 0xAA
#define METADATA_SIZE 128
#define POINTER_SIZE sizeof(void*)
#define LOCATION_MAX (METADATA_SIZE - POINTER_SIZE)
#define NEXT_OFFSET LOCATION_MAX

#if defined(__linux__) || defined(__APPLE__) || defined(__FreeBSD__)
#include <execinfo.h>
#include <pthread.h>
#define BACKTRACE_ENABLED 1
#elif defined(_WIN32)
#include <windows.h>
#include <dbghelp.h>
#define BACKTRACE_ENABLED 1
#else
#define BACKTRACE_ENABLED 0
#endif

#if defined(__linux__) || defined(__APPLE__) || defined(__FreeBSD__)
static pthread_mutex_t alloc_mutex = PTHREAD_MUTEX_INITIALIZER;
#define LOCK() pthread_mutex_lock(&alloc_mutex)
#define UNLOCK() pthread_mutex_unlock(&alloc_mutex)
#elif defined(_WIN32)
static CRITICAL_SECTION alloc_mutex;
static bool mutex_initialized = false;
static void init_mutex() {
    if (!mutex_initialized) {
        InitializeCriticalSection(&alloc_mutex);
        mutex_initialized = true;
    }
}
#define LOCK() init_mutex(); EnterCriticalSection(&alloc_mutex)
#define UNLOCK() LeaveCriticalSection(&alloc_mutex)
#else
#error "Unsupported platform for thread safety"
#endif

static void* allocated_head = NULL;
static size_t allocated_count = 0;

static void hex_dump(const void* mem, size_t len) {
    const uint8_t* data = (const uint8_t*)mem;
    for (size_t i = 0; i < len; i += 16) {
        printf("%08zx: ", i);
        for (size_t j = 0; j < 16; j++) {
            if (i + j < len) {
                printf("%02x ", data[i + j]);
            } else {
                printf("   ");
            }
        }
        printf("  ");
        for (size_t j = 0; j < 16; j++) {
            if (i + j < len) {
                uint8_t c = data[i + j];
                printf("%c", (c >= 32 && c <= 126) ? c : '.');
            } else {
                printf(" ");
            }
        }
        printf("\n");
    }
}

void u_check(void* ptr, const char* text, const char* location) {
    if (!ptr) return;
    uint8_t* base = (uint8_t*)ptr - 8 - BOUNDARY_CHECK_SIZE - METADATA_SIZE;
    char* alloc_loc = (char*)base;
    uint8_t* prefix_start = base + METADATA_SIZE;
    uint32_t* pre_canary = (uint32_t*)(prefix_start + BOUNDARY_CHECK_SIZE);
    uint32_t size = *(uint32_t*)(prefix_start + BOUNDARY_CHECK_SIZE + 4);
    uint32_t total = size + 12 + 2 * BOUNDARY_CHECK_SIZE + METADATA_SIZE;
    uint8_t* user = (uint8_t*)ptr;
    bool error = false;
    for (size_t i = 0; i < BOUNDARY_CHECK_SIZE; i++) {
        if (prefix_start[i] != PADDING_FILL) {
            printf("%s: %s: underflow prefix corrupted at %p, byte %zu (expected 0x%02X, got 0x%02X)\n", location, text, ptr, i, PADDING_FILL, prefix_start[i]);
            DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "%s: %s: underflow prefix corrupted at %p, byte %zu (expected 0x%02X, got 0x%02X)\n", location, text, ptr, i, PADDING_FILL, prefix_start[i]);
            error = true;
        }
    }
    if (*pre_canary != CANARY) {
        printf("%s: %s: underflow canary corrupted at %p (expected 0x%X, got 0x%X)\n", location, text, ptr, CANARY, *pre_canary);
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "%s: %s: underflow canary corrupted at %p (expected 0x%X, got 0x%X)\n", location, text, ptr, CANARY, *pre_canary);
        error = true;
    }
    uint8_t* post_start = user + size + 4;
    for (size_t i = 0; i < BOUNDARY_CHECK_SIZE; i++) {
        if (post_start[i] != PADDING_FILL) {
            printf("%s: %s: overflow prefix corrupted at %p, byte %zu (expected 0x%02X, got 0x%02X)\n", location, text, ptr, i, PADDING_FILL, post_start[i]);
            DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "%s: %s: overflow prefix corrupted at %p, byte %zu (expected 0x%02X, got 0x%02X)\n", location, text, ptr, i, PADDING_FILL, post_start[i]);
            error = true;
        }
    }
    uint32_t* post_canary = (uint32_t*)(user + size);
    if (*post_canary != CANARY) {
        printf("%s: %s: overflow canary corrupted at %p (expected 0x%X, got 0x%X)\n", location, text, ptr, CANARY, *post_canary);
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "%s: %s: overflow canary corrupted at %p (expected 0x%X, got 0x%X)\n", location, text, ptr, CANARY, *post_canary);
        error = true;
    }
    if (error) {
        printf("%s: %s: Allocated at: %s\n", location, text, alloc_loc);
        printf("%s: %s: Memory block size (user): %u bytes, total allocated: %u bytes\n", location, text, size, total);
        printf("%s: %s: Full memory dump at base %p:\n", location, text, base);
        hex_dump(base, total);
        exit(EXIT_FAILURE);
    }
}

void* u_malloc_impl(uint32_t size, const char* location) {
    uint32_t total = size + 12 + 2 * BOUNDARY_CHECK_SIZE + METADATA_SIZE;
    uint8_t* base = (uint8_t*)malloc(total);
    if (!base && size > 0) return NULL;
    memset(base, PADDING_FILL, total);
    // Set metadata
    size_t loc_len = strlen(location);
    size_t max_copy = (loc_len >= LOCATION_MAX) ? LOCATION_MAX - 1 : loc_len;
    memcpy(base, location, max_copy);
    base[max_copy] = 0;
    // Set allocation data
    uint8_t* prefix_start = base + METADATA_SIZE;
    *(uint32_t*)(prefix_start + BOUNDARY_CHECK_SIZE) = CANARY;
    *(uint32_t*)(prefix_start + BOUNDARY_CHECK_SIZE + 4) = size;
    *(uint32_t*)(prefix_start + BOUNDARY_CHECK_SIZE + 8 + size) = CANARY;
    // Add to list (thread-safe)
    LOCK();
    void** next_ptr = (void**)(base + NEXT_OFFSET);
    *next_ptr = allocated_head;
    allocated_head = base;
    allocated_count++;
    UNLOCK();
    return prefix_start + BOUNDARY_CHECK_SIZE + 8;
}

void* u_calloc_impl(uint32_t nmemb, uint32_t size, const char* location) {
    uint32_t total_size = nmemb * size;
    void* ptr = u_malloc_impl(total_size, location);
    if (ptr) memset(ptr, 0, total_size);
    return ptr;
}

void* u_realloc_impl(void* ptr, uint32_t size, const char* location) {
    if (!ptr) return u_malloc_impl(size, location);
    if (size == 0) {
        u_free_impl(ptr, location);
        return NULL;
    }
    u_check(ptr, "u_realloc", location);
    uint8_t* old_base = (uint8_t*)ptr - 8 - BOUNDARY_CHECK_SIZE - METADATA_SIZE;
    uint32_t old_size = *(uint32_t*)(old_base + METADATA_SIZE + BOUNDARY_CHECK_SIZE + 4);
    void* new_ptr = u_malloc_impl(size, location);
    if (!new_ptr) return NULL;
    uint32_t copy_size = (old_size < size) ? old_size : size;
    memcpy(new_ptr, ptr, copy_size);
    u_free_impl(ptr, location);
    return new_ptr;
}

void u_free_impl(void* ptr, const char* location) {
    if (!ptr) return;
    u_check(ptr, "u_free", location);
    uint8_t* base = (uint8_t*)ptr - 8 - BOUNDARY_CHECK_SIZE - METADATA_SIZE;
    // Remove from linked list (thread-safe)
    LOCK();
    void* prev = NULL;
    void* curr = allocated_head;
    bool found = false;
    while (curr) {
        if (curr == base) {
            if (prev) {
                *(void**)((uint8_t*)prev + NEXT_OFFSET) = *(void**)((uint8_t*)curr + NEXT_OFFSET);
            } else {
                allocated_head = *(void**)((uint8_t*)curr + NEXT_OFFSET);
            }
            allocated_count--;
            found = true;
            break;
        }
        prev = curr;
        curr = *(void**)((uint8_t*)curr + NEXT_OFFSET);
    }
    UNLOCK();
    if (!found) {
        printf("%s: u_free: Attempt to free unallocated or already freed pointer %p\n", location, ptr);
        exit(EXIT_FAILURE);
    }
    free(base);
}

char* u_strdup_impl(const char* s, const char* location) {
    if (!s) return NULL;
    size_t len = strlen(s) + 1;
    char* copy = (char*)u_malloc_impl((uint32_t)len, location);
    if (copy) memcpy(copy, s, len);
    return copy;
}

void u_report_unfreed_blocks(void) {
    LOCK();
    size_t count = allocated_count;
    void* head = allocated_head;
    UNLOCK();
    if (count == 0) {
        printf("No unfreed memory blocks.\n");
        return;
    }
    printf("Unfreed memory blocks: %zu\n", count);
    void* curr = head;
    while (curr) {
        char* alloc_loc = (char*)curr;
        uint8_t* prefix_start = (uint8_t*)curr + METADATA_SIZE;
        uint32_t size = *(uint32_t*)(prefix_start + BOUNDARY_CHECK_SIZE + 4);
        void* user_ptr = prefix_start + BOUNDARY_CHECK_SIZE + 8;
        u_check(user_ptr, "unfreed_check", "u_report_unfreed_blocks");
        printf(" - Allocated at: %s, size: %u bytes\n", alloc_loc, size);
        curr = *(void**)((uint8_t*)curr + NEXT_OFFSET);
    }
}