#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 0
#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 size_t allocations_count = 0;
static size_t frees_count = 0;

static void hex_dump(const void* mem, size_t len) {
    const uint8_t* data = (const uint8_t*)mem;
    char line[128];
    for (size_t i = 0; i < len; i += 16) {
        int off = snprintf(line, sizeof(line), "%08zx: ", i);
        for (size_t j = 0; j < 16; j++) {
            if (i + j < len) {
                off += snprintf(line + off, sizeof(line) - off, "%02x ", data[i + j]);
            } else {
                off += snprintf(line + off, sizeof(line) - off, "   ");
            }
        }
        off += snprintf(line + off, sizeof(line) - off, " ");
        for (size_t j = 0; j < 16; j++) {
            if (i + j < len) {
                uint8_t c = data[i + j];
                off += snprintf(line + off, sizeof(line) - off, "%c", (c >= 32 && c <= 126) ? c : '.');
            } else {
                off += snprintf(line + off, sizeof(line) - off, " ");
            }
        }
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "%s", line);
    }
}

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) {
            DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "%s: %s: underflow prefix corrupted at %p, byte %zu (expected 0x%02X, got 0x%02X)", location, text, ptr, i, PADDING_FILL, prefix_start[i]);
            error = true;
        }
    }
    if (*pre_canary != CANARY) {
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "%s: %s: underflow canary corrupted at %p (expected 0x%X, got 0x%X)", 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) {
            DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "%s: %s: overflow prefix corrupted at %p, byte %zu (expected 0x%02X, got 0x%02X)", location, text, ptr, i, PADDING_FILL, post_start[i]);
            error = true;
        }
    }
    uint32_t* post_canary = (uint32_t*)(user + size);
    if (*post_canary != CANARY) {
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "%s: %s: overflow canary corrupted at %p (expected 0x%X, got 0x%X)", location, text, ptr, CANARY, *post_canary);
        error = true;
    }
    if (error) {
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "%s: %s: Allocated at: %s", location, text, alloc_loc);
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "%s: %s: Memory block size (user): %u bytes, total allocated: %u bytes", location, text, size, total);
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "%s: %s: Full memory dump at base %p:", 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++;
    allocations_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;
    char* alloc_loc = (char*)base;
    uint32_t size = *(uint32_t*)(base + METADATA_SIZE + BOUNDARY_CHECK_SIZE + 4);
    // 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--;
            frees_count++;
            found = true;
            break;
        }
        prev = curr;
        curr = *(void**)((uint8_t*)curr + NEXT_OFFSET);
    }
    size_t active_count = allocated_count;
    UNLOCK();
    if (!found) {
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "========================================");
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "DOUBLE FREE DETECTED");
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "========================================");
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "Free attempt at:   %s", location);
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "Pointer:           %p", ptr);
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "Base:              %p", (void*)base);
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "Allocated at:      %s", alloc_loc);
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "User size:         %u bytes", size);
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "Total allocated:   %u bytes", size + 12 + 2 * BOUNDARY_CHECK_SIZE + METADATA_SIZE);
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "Active allocs:     %zu", active_count);
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "--- Metadata dump at %p (%u bytes) ---", (void*)base, (unsigned)(METADATA_SIZE + BOUNDARY_CHECK_SIZE + 8));
        hex_dump(base, METADATA_SIZE + BOUNDARY_CHECK_SIZE + 8);
        if (size > 0 && size <= 1024 * 1024) {
            uint32_t dump_size = (size > 256) ? 256 : size;
            DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "--- User data dump at %p (%u of %u bytes) ---", ptr, dump_size, size);
            hex_dump(ptr, dump_size);
        }
#if BACKTRACE_ENABLED
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "--- Backtrace ---");
#  if defined(__linux__) || defined(__APPLE__) || defined(__FreeBSD__)
        void* bt_buf[32];
        int bt_size = backtrace(bt_buf, 32);
        char** bt_symbols = backtrace_symbols(bt_buf, bt_size);
        if (bt_symbols) {
            for (int i = 0; i < bt_size; i++) DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "  #%d %s", i, bt_symbols[i]);
            free(bt_symbols);
        }
#  elif defined(_WIN32)
        void* bt_buf[32];
        USHORT frames = CaptureStackBackTrace(0, 32, bt_buf, NULL);
        for (USHORT i = 0; i < frames; i++) DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "  #%d %p", i, bt_buf[i]);
#  endif
#endif
        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;
    size_t alloc_count = allocations_count;
    size_t free_count = frees_count;
    UNLOCK();
    DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "Memory stats: allocations=%zu, frees=%zu, current=%zu", alloc_count, free_count, count);
    if (count == 0) {
        DEBUG_INFO(DEBUG_CATEGORY_MEMORY, "No unfreed memory blocks.");
        return;
    }
    DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "Unfreed memory blocks: %zu, allocations: %zu frees: %zu", count, allocations_count, frees_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");
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, " - Allocated at: %s, size: %u bytes", alloc_loc, size);
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "   Full memory dump at %p:", user_ptr);
        hex_dump(user_ptr, size);
        curr = *(void**)((uint8_t*)curr + NEXT_OFFSET);
    }
}

size_t u_get_allocated_count(void) {
    return allocated_count;
}