//! Arena-based memory allocation for resource-bounded evaluation
//!
//! Provides deterministic memory management with O(1) allocation and bulk deallocation.
//! Based on the REPL resource-bounded evaluation requirements.
//!
//! # Safety
//! 
//! This module uses unsafe code for manual memory management. All unsafe operations
//! are carefully validated and encapsulated within safe abstractions.

use anyhow::{Result, anyhow};
use std::alloc::{alloc, dealloc, Layout};
use std::cell::RefCell;
use std::marker::PhantomData;
use std::ptr::{self, NonNull};
use std::rc::Rc;

// ============================================================================
// Arena Allocator
// ============================================================================

/// Arena allocator for fast, bounded memory allocation
pub struct Arena {
    /// Current chunk being allocated from
    current_chunk: RefCell<Chunk>,
    /// Previous chunks that are full
    chunks: RefCell<Vec<Chunk>>,
    /// Total size limit
    max_size: usize,
    /// Current total allocated across all chunks
    total_allocated: RefCell<usize>,
    /// Statistics
    stats: RefCell<ArenaStats>,
}

/// A single chunk of memory in the arena
struct Chunk {
    /// Pointer to the start of the memory
    ptr: NonNull<u8>,
    /// Size of this chunk
    size: usize,
    /// Current position in the chunk
    pos: usize,
    /// Layout for deallocation
    layout: Layout,
}

/// Statistics for arena usage
#[derive(Debug, Clone, Default)]
pub struct ArenaStats {
    pub allocations: usize,
    pub deallocations: usize,
    pub peak_usage: usize,
    pub fragmentation: f64,
    pub chunks_allocated: usize,
}

impl Arena {
    /// Create a new arena with the given maximum size
    pub fn new(max_size: usize) -> Self {
        // Start with a 64KB chunk
        let initial_size = (64 * 1024).min(max_size);
        let chunk = Chunk::new(initial_size).expect("Failed to allocate initial chunk");
        
        Self {
            current_chunk: RefCell::new(chunk),
            chunks: RefCell::new(Vec::new()),
            max_size,
            total_allocated: RefCell::new(initial_size),
            stats: RefCell::new(ArenaStats {
                chunks_allocated: 1,
                ..Default::default()
            }),
        }
    }
    
    /// Allocate memory for a value of type T
    pub fn alloc<T>(&self, value: T) -> Result<ArenaBox<T>> {
        let layout = Layout::new::<T>();
        let ptr = self.alloc_raw(layout)?;
        
        unsafe {
            ptr::write(ptr.cast::<T>().as_ptr(), value);
        }
        
        Ok(ArenaBox {
            ptr: ptr.cast(),
            arena: self as *const Arena,
            _phantom: PhantomData,
        })
    }
    
    /// Allocate raw memory with the given layout
    pub fn alloc_raw(&self, layout: Layout) -> Result<NonNull<u8>> {
        let size = layout.size();
        let align = layout.align();
        
        // Check if allocation would exceed limit
        if *self.total_allocated.borrow() + size > self.max_size {
            return Err(anyhow!(
                "Arena allocation would exceed limit: {} + {} > {}",
                *self.total_allocated.borrow(),
                size,
                self.max_size
            ));
        }
        
        // Try to allocate from current chunk
        let mut current = self.current_chunk.borrow_mut();
        if let Some(ptr) = current.try_alloc(size, align) {
            self.stats.borrow_mut().allocations += 1;
            self.update_peak_usage();
            return Ok(ptr);
        }
        
        // Current chunk is full, need a new one
        drop(current); // Release borrow
        self.grow_arena(size)?;
        
        // Try again with new chunk
        let mut current = self.current_chunk.borrow_mut();
        current.try_alloc(size, align)
            .ok_or_else(|| anyhow!("Failed to allocate after growing arena"))
    }
    
    /// Grow the arena by allocating a new chunk
    fn grow_arena(&self, min_size: usize) -> Result<()> {
        // Move current chunk to the full chunks list
        let old_chunk = self.current_chunk.replace(Chunk::empty());
        self.chunks.borrow_mut().push(old_chunk);
        
        // Calculate new chunk size (double previous, but respect limits)
        let new_size = (*self.total_allocated.borrow())
            .saturating_mul(2)
            .max(min_size)
            .min(self.max_size - *self.total_allocated.borrow());
        
        if new_size == 0 {
            return Err(anyhow!("Arena is full"));
        }
        
        // Allocate new chunk
        let new_chunk = Chunk::new(new_size)?;
        *self.current_chunk.borrow_mut() = new_chunk;
        *self.total_allocated.borrow_mut() += new_size;
        self.stats.borrow_mut().chunks_allocated += 1;
        
        Ok(())
    }
    
    /// Reset the arena, deallocating all memory
    pub fn reset(&self) {
        // Clear current chunk
        self.current_chunk.borrow_mut().reset();
        
        // Clear all other chunks
        self.chunks.borrow_mut().clear();
        
        // Reset stats
        self.stats.borrow_mut().deallocations = self.stats.borrow().allocations;
        
        // Reset to initial chunk size
        let initial_size = (64 * 1024).min(self.max_size);
        if self.current_chunk.borrow().size != initial_size {
            *self.current_chunk.borrow_mut() = 
                Chunk::new(initial_size).expect("Failed to allocate chunk");
            *self.total_allocated.borrow_mut() = initial_size;
        }
    }
    
    /// Get current memory usage
    pub fn used(&self) -> usize {
        let mut total = self.current_chunk.borrow().pos;
        for chunk in self.chunks.borrow().iter() {
            total += chunk.pos;
        }
        total
    }
    
    /// Get statistics
    pub fn stats(&self) -> ArenaStats {
        self.stats.borrow().clone()
    }
    
    fn update_peak_usage(&self) {
        let current = self.used();
        let mut stats = self.stats.borrow_mut();
        if current > stats.peak_usage {
            stats.peak_usage = current;
        }
        
        // Calculate fragmentation
        let total_allocated = *self.total_allocated.borrow();
        if total_allocated > 0 {
            stats.fragmentation = 1.0 - (current as f64 / total_allocated as f64);
        }
    }
}

impl Drop for Arena {
    fn drop(&mut self) {
        // Chunks will deallocate their memory when dropped
    }
}

impl Chunk {
    fn new(size: usize) -> Result<Self> {
        let layout = Layout::from_size_align(size, 8)
            .map_err(|e| anyhow!("Invalid layout: {}", e))?;
        
        let ptr = unsafe {
            let ptr = alloc(layout);
            if ptr.is_null() {
                return Err(anyhow!("Failed to allocate chunk of size {}", size));
            }
            NonNull::new_unchecked(ptr)
        };
        
        Ok(Self {
            ptr,
            size,
            pos: 0,
            layout,
        })
    }
    
    fn empty() -> Self {
        Self {
            ptr: NonNull::dangling(),
            size: 0,
            pos: 0,
            layout: Layout::from_size_align(1, 1).unwrap(),
        }
    }
    
    fn try_alloc(&mut self, size: usize, align: usize) -> Option<NonNull<u8>> {
        // Align the current position
        let aligned_pos = (self.pos + align - 1) & !(align - 1);
        
        // Check if we have enough space
        if aligned_pos + size > self.size {
            return None;
        }
        
        // Calculate pointer
        let ptr = unsafe {
            NonNull::new_unchecked(self.ptr.as_ptr().add(aligned_pos))
        };
        
        self.pos = aligned_pos + size;
        Some(ptr)
    }
    
    fn reset(&mut self) {
        self.pos = 0;
    }
}

impl Drop for Chunk {
    fn drop(&mut self) {
        if self.size > 0 {
            unsafe {
                dealloc(self.ptr.as_ptr(), self.layout);
            }
        }
    }
}

// ============================================================================
// Arena Box - Smart Pointer for Arena Allocation
// ============================================================================

/// A smart pointer to a value allocated in an arena
pub struct ArenaBox<T> {
    ptr: NonNull<T>,
    arena: *const Arena,
    _phantom: PhantomData<T>,
}

impl<T> ArenaBox<T> {
    /// Get a reference to the value
    pub fn get(&self) -> &T {
        unsafe { self.ptr.as_ref() }
    }
    
    /// Get a mutable reference to the value
    pub fn get_mut(&mut self) -> &mut T {
        unsafe { self.ptr.as_mut() }
    }
}

impl<T> std::ops::Deref for ArenaBox<T> {
    type Target = T;
    
    fn deref(&self) -> &Self::Target {
        self.get()
    }
}

impl<T> std::ops::DerefMut for ArenaBox<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.get_mut()
    }
}

// Note: ArenaBox doesn't implement Drop because the arena handles deallocation

// ============================================================================
// Transactional Arena - Supports Checkpoints and Rollback
// ============================================================================

/// Arena with transactional support for checkpointing and rollback
pub struct TransactionalArena {
    arena: Arena,
    checkpoints: Vec<ArenaCheckpoint>,
}

/// A checkpoint in the arena's allocation state
#[derive(Debug, Clone)]
struct ArenaCheckpoint {
    /// Position in current chunk
    chunk_pos: usize,
    /// Number of chunks at checkpoint
    num_chunks: usize,
    /// Total allocated at checkpoint
    total_allocated: usize,
    /// Stats at checkpoint
    stats: ArenaStats,
}

impl TransactionalArena {
    /// Create a new transactional arena
    pub fn new(max_size: usize) -> Self {
        Self {
            arena: Arena::new(max_size),
            checkpoints: Vec::new(),
        }
    }
    
    /// Create a checkpoint for later rollback
    pub fn checkpoint(&mut self) -> usize {
        let checkpoint = ArenaCheckpoint {
            chunk_pos: self.arena.current_chunk.borrow().pos,
            num_chunks: self.arena.chunks.borrow().len(),
            total_allocated: *self.arena.total_allocated.borrow(),
            stats: self.arena.stats(),
        };
        
        self.checkpoints.push(checkpoint);
        self.checkpoints.len() - 1
    }
    
    /// Rollback to a specific checkpoint
    pub fn rollback(&mut self, checkpoint_id: usize) -> Result<()> {
        if checkpoint_id >= self.checkpoints.len() {
            return Err(anyhow!("Invalid checkpoint ID"));
        }
        
        let checkpoint = &self.checkpoints[checkpoint_id];
        
        // Restore chunk position
        self.arena.current_chunk.borrow_mut().pos = checkpoint.chunk_pos;
        
        // Remove chunks allocated after checkpoint
        let mut chunks = self.arena.chunks.borrow_mut();
        chunks.truncate(checkpoint.num_chunks);
        
        // Restore total allocated
        *self.arena.total_allocated.borrow_mut() = checkpoint.total_allocated;
        
        // Restore stats
        *self.arena.stats.borrow_mut() = checkpoint.stats.clone();
        
        // Remove checkpoints after this one
        self.checkpoints.truncate(checkpoint_id + 1);
        
        Ok(())
    }
    
    /// Commit the current state, removing the last checkpoint
    pub fn commit(&mut self) -> Result<()> {
        if self.checkpoints.is_empty() {
            return Err(anyhow!("No checkpoint to commit"));
        }
        
        self.checkpoints.pop();
        Ok(())
    }
    
    /// Get the underlying arena
    pub fn arena(&self) -> &Arena {
        &self.arena
    }
    
    /// Reset the arena and all checkpoints
    pub fn reset(&mut self) {
        self.arena.reset();
        self.checkpoints.clear();
    }
}

// ============================================================================
// Pool Allocator - For Frequent Same-Size Allocations
// ============================================================================

/// Pool allocator for objects of the same size
pub struct Pool<T> {
    free_list: RefCell<Vec<*mut T>>,
    arena: Rc<Arena>,
    _phantom: PhantomData<T>,
}

impl<T> Pool<T> {
    /// Create a new pool using the given arena
    pub fn new(arena: Rc<Arena>) -> Self {
        Self {
            free_list: RefCell::new(Vec::new()),
            arena,
            _phantom: PhantomData,
        }
    }
    
    /// Allocate a value from the pool
    pub fn alloc(&self, value: T) -> Result<PoolBox<T>> {
        // Try to reuse from free list
        if let Some(ptr) = self.free_list.borrow_mut().pop() {
            unsafe {
                ptr::write(ptr, value);
            }
            return Ok(PoolBox {
                ptr: unsafe { NonNull::new_unchecked(ptr) },
                pool: self as *const Pool<T>,
            });
        }
        
        // Allocate new from arena
        let layout = Layout::new::<T>();
        let ptr = self.arena.alloc_raw(layout)?;
        
        unsafe {
            ptr::write(ptr.cast::<T>().as_ptr(), value);
        }
        
        Ok(PoolBox {
            ptr: ptr.cast(),
            pool: self as *const Pool<T>,
        })
    }
    
    /// Return a pointer to the free list
    fn free(&self, ptr: *mut T) {
        self.free_list.borrow_mut().push(ptr);
    }
}

/// Smart pointer for pool-allocated values
pub struct PoolBox<T> {
    ptr: NonNull<T>,
    pool: *const Pool<T>,
}

impl<T> Drop for PoolBox<T> {
    fn drop(&mut self) {
        unsafe {
            // Drop the value
            ptr::drop_in_place(self.ptr.as_ptr());
            // Return to pool
            (*self.pool).free(self.ptr.as_ptr());
        }
    }
}

impl<T> std::ops::Deref for PoolBox<T> {
    type Target = T;
    
    fn deref(&self) -> &Self::Target {
        unsafe { self.ptr.as_ref() }
    }
}

impl<T> std::ops::DerefMut for PoolBox<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { self.ptr.as_mut() }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_arena_allocation() {
        let arena = Arena::new(1024 * 1024); // 1MB limit
        
        // Allocate some values
        let v1 = arena.alloc(42i32).unwrap();
        let v2 = arena.alloc("hello".to_string()).unwrap();
        let v3 = arena.alloc(vec![1, 2, 3]).unwrap();
        
        assert_eq!(*v1, 42);
        assert_eq!(*v2, "hello");
        assert_eq!(*v3, vec![1, 2, 3]);
        
        let stats = arena.stats();
        assert_eq!(stats.allocations, 3);
    }
    
    #[test]
    fn test_arena_reset() {
        let arena = Arena::new(1024 * 1024);
        
        // Allocate some memory
        for i in 0..100 {
            arena.alloc(i).unwrap();
        }
        
        let used_before = arena.used();
        assert!(used_before > 0);
        
        // Reset arena
        arena.reset();
        
        let used_after = arena.used();
        assert_eq!(used_after, 0);
    }
    
    #[test]
    fn test_transactional_arena() {
        let mut arena = TransactionalArena::new(1024 * 1024);
        
        // Initial allocation
        arena.arena().alloc(1).unwrap();
        
        // Create checkpoint
        let checkpoint = arena.checkpoint();
        
        // More allocations
        arena.arena().alloc(2).unwrap();
        arena.arena().alloc(3).unwrap();
        
        let used_before = arena.arena().used();
        
        // Rollback
        arena.rollback(checkpoint).unwrap();
        
        let used_after = arena.arena().used();
        assert!(used_after < used_before);
    }
    
    #[test]
    fn test_memory_limit() {
        let arena = Arena::new(100); // Very small limit
        
        // This should succeed
        arena.alloc([0u8; 50]).unwrap();
        
        // This should fail - would exceed limit
        let result = arena.alloc([0u8; 60]);
        assert!(result.is_err());
    }
    
    #[test]
    fn test_pool_allocator() {
        let arena = Rc::new(Arena::new(1024 * 1024));
        let pool: Pool<i32> = Pool::new(arena);
        
        // Allocate and free
        {
            let v1 = pool.alloc(42).unwrap();
            assert_eq!(*v1, 42);
        } // v1 dropped, returned to pool
        
        // Next allocation should reuse
        let v2 = pool.alloc(100).unwrap();
        assert_eq!(*v2, 100);
    }
}