#include "memory_manager.h"
#include <cstring>

#ifdef ARDUINO
#include <Arduino.h>
#define DEBUG_PRINT(x) Serial.print(x)
#define DEBUG_PRINTLN(x) Serial.println(x)
#else
#include <cstdio>
#define DEBUG_PRINT(x) printf("%s", (x))
#define DEBUG_PRINTLN(x) printf("%s\n", (x))
#endif

// Global memory pool instance
MemoryPool globalMemoryPool;

MemoryPool::MemoryPool() : next_free_offset(0), allocation_count(0), deallocation_count(0) {
    // Initialize all blocks as unused
    for (size_t i = 0; i < MAX_BLOCKS; ++i) {
        blocks[i].ptr = nullptr;
        blocks[i].size = 0;
        blocks[i].in_use = false;
        blocks[i].magic = 0;
    }
    
    // Clear the memory pool
    memset(pool_memory, 0, POOL_SIZE);
}

MemoryPool::~MemoryPool() {
    // Check for memory leaks
    uint32_t active_blocks = 0;
    for (size_t i = 0; i < MAX_BLOCKS; ++i) {
        if (blocks[i].in_use) {
            active_blocks++;
        }
    }
    
    if (active_blocks > 0) {
#ifdef DEBUG_MEMORY
        DEBUG_PRINT("WARNING: Memory pool destroyed with ");
        DEBUG_PRINT(active_blocks);
        DEBUG_PRINTLN(" active blocks!");
#endif
    }
}

void* MemoryPool::allocate(size_t size) {
    if (size == 0 || size > POOL_SIZE) {
        return nullptr;
    }
    
    // Align size to 4-byte boundary for better performance
    size = (size + 3) & ~3;
    
    // Check if we have enough space
    if (next_free_offset + size > POOL_SIZE) {
        // Try to find a free block that was previously deallocated
        for (size_t i = 0; i < MAX_BLOCKS; ++i) {
            if (!blocks[i].in_use && blocks[i].ptr != nullptr && blocks[i].size >= size) {
                blocks[i].in_use = true;
                blocks[i].magic = MAGIC_NUMBER;
                allocation_count++;
                return blocks[i].ptr;
            }
        }
        return nullptr; // Out of memory
    }
    
    // Find a free block descriptor
    size_t block_index = MAX_BLOCKS;
    for (size_t i = 0; i < MAX_BLOCKS; ++i) {
        if (!blocks[i].in_use && blocks[i].ptr == nullptr) {
            block_index = i;
            break;
        }
    }
    
    if (block_index == MAX_BLOCKS) {
        return nullptr; // No free block descriptors
    }
    
    // Allocate from the pool
    void* ptr = &pool_memory[next_free_offset];
    
    // Set up the block descriptor
    blocks[block_index].ptr = ptr;
    blocks[block_index].size = size;
    blocks[block_index].in_use = true;
    blocks[block_index].magic = MAGIC_NUMBER;
    
    next_free_offset += size;
    allocation_count++;
    
    return ptr;
}

bool MemoryPool::deallocate(void* ptr) {
    if (ptr == nullptr) {
        return false;
    }
    
    // Find the block
    for (size_t i = 0; i < MAX_BLOCKS; ++i) {
        if (blocks[i].ptr == ptr && blocks[i].in_use) {
            // Check magic number for corruption
            if (blocks[i].magic != MAGIC_NUMBER) {
#ifdef DEBUG_MEMORY
                DEBUG_PRINTLN("ERROR: Memory corruption detected during deallocation!");
#endif
                return false;
            }
            
            blocks[i].in_use = false;
            blocks[i].magic = 0;
            deallocation_count++;
            
            // Clear the memory for security
            memset(ptr, 0, blocks[i].size);
            
            return true;
        }
    }
    
    return false; // Pointer not found
}

MemoryPool::MemoryStats MemoryPool::getStats() const {
    MemoryStats stats;
    stats.total_size = POOL_SIZE;
    stats.allocations = allocation_count;
    stats.deallocations = deallocation_count;
    stats.active_blocks = 0;
    stats.used_size = 0;
    
    for (size_t i = 0; i < MAX_BLOCKS; ++i) {
        if (blocks[i].in_use) {
            stats.active_blocks++;
            stats.used_size += blocks[i].size;
        }
    }
    
    stats.free_size = POOL_SIZE - stats.used_size;
    
    return stats;
}

bool MemoryPool::checkIntegrity() const {
    for (size_t i = 0; i < MAX_BLOCKS; ++i) {
        if (blocks[i].in_use && blocks[i].magic != MAGIC_NUMBER) {
            return false;
        }
    }
    return true;
}

void MemoryPool::reset() {
    // Clear all blocks
    for (size_t i = 0; i < MAX_BLOCKS; ++i) {
        blocks[i].ptr = nullptr;
        blocks[i].size = 0;
        blocks[i].in_use = false;
        blocks[i].magic = 0;
    }
    
    next_free_offset = 0;
    allocation_count = 0;
    deallocation_count = 0;
    
    // Clear the memory pool
    memset(pool_memory, 0, POOL_SIZE);
}

void MemoryPool::printStats() const {
    MemoryStats stats = getStats();
    
#ifdef ARDUINO
    Serial.println("=== Memory Pool Statistics ===");
    Serial.print("Total Size: ");
    Serial.print(stats.total_size);
    Serial.println(" bytes");
    Serial.print("Used Size: ");
    Serial.print(stats.used_size);
    Serial.println(" bytes");
    Serial.print("Free Size: ");
    Serial.print(stats.free_size);
    Serial.println(" bytes");
    Serial.print("Active Blocks: ");
    Serial.println(stats.active_blocks);
    Serial.print("Total Allocations: ");
    Serial.println(stats.allocations);
    Serial.print("Total Deallocations: ");
    Serial.println(stats.deallocations);
    Serial.print("Memory Integrity: ");
    Serial.println(checkIntegrity() ? "OK" : "CORRUPTED");
    Serial.println("==============================");
#else
    printf("=== Memory Pool Statistics ===\n");
    printf("Total Size: %zu bytes\n", stats.total_size);
    printf("Used Size: %zu bytes\n", stats.used_size);
    printf("Free Size: %zu bytes\n", stats.free_size);
    printf("Active Blocks: %u\n", stats.active_blocks);
    printf("Total Allocations: %u\n", stats.allocations);
    printf("Total Deallocations: %u\n", stats.deallocations);
    printf("Memory Integrity: %s\n", checkIntegrity() ? "OK" : "CORRUPTED");
    printf("==============================\n");
#endif
}

// SafeBuffer implementation
SafeBuffer::SafeBuffer(size_t capacity, MemoryPool* pool) 
    : data_(nullptr), size_(0), capacity_(capacity), pool_(pool), owns_memory_(false) {
    
    if (capacity > 0) {
        if (pool_) {
            data_ = static_cast<uint8_t*>(pool_->allocate(capacity));
        } else {
#ifdef ARDUINO
            data_ = static_cast<uint8_t*>(malloc(capacity));
#else
            data_ = new uint8_t[capacity];
#endif
        }
        
        if (data_) {
            owns_memory_ = true;
            memset(data_, 0, capacity);
        }
    }
}

SafeBuffer::~SafeBuffer() {
    if (data_ && owns_memory_) {
        if (pool_) {
            pool_->deallocate(data_);
        } else {
#ifdef ARDUINO
            free(data_);
#else
            delete[] data_;
#endif
        }
    }
}

SafeBuffer::SafeBuffer(SafeBuffer&& other) noexcept 
    : data_(other.data_), size_(other.size_), capacity_(other.capacity_), 
      pool_(other.pool_), owns_memory_(other.owns_memory_) {
    other.data_ = nullptr;
    other.size_ = 0;
    other.capacity_ = 0;
    other.owns_memory_ = false;
}

SafeBuffer& SafeBuffer::operator=(SafeBuffer&& other) noexcept {
    if (this != &other) {
        // Clean up current resources
        if (data_ && owns_memory_) {
            if (pool_) {
                pool_->deallocate(data_);
            } else {
#ifdef ARDUINO
                free(data_);
#else
                delete[] data_;
#endif
            }
        }
        
        // Move from other
        data_ = other.data_;
        size_ = other.size_;
        capacity_ = other.capacity_;
        pool_ = other.pool_;
        owns_memory_ = other.owns_memory_;
        
        // Reset other
        other.data_ = nullptr;
        other.size_ = 0;
        other.capacity_ = 0;
        other.owns_memory_ = false;
    }
    return *this;
}

bool SafeBuffer::write(const void* data, size_t size, size_t offset) {
    if (!data_ || !data || size == 0) {
        return false;
    }
    
    if (offset + size > capacity_) {
        return false; // Would exceed buffer capacity
    }
    
    memcpy(data_ + offset, data, size);
    
    // Update size if we wrote beyond current size
    if (offset + size > size_) {
        size_ = offset + size;
    }
    
    return true;
}

bool SafeBuffer::read(void* data, size_t size, size_t offset) const {
    if (!data_ || !data || size == 0) {
        return false;
    }
    
    if (offset + size > size_) {
        return false; // Would read beyond valid data
    }
    
    memcpy(data, data_ + offset, size);
    return true;
}

bool SafeBuffer::append(const void* data, size_t size) {
    return write(data, size, size_);
}

void SafeBuffer::clear() {
    if (data_) {
        memset(data_, 0, capacity_);
        size_ = 0;
    }
}