//! Frame-based coordination for per-frame allocations.

use memkit::MkFastArena;
use parking_lot::Mutex;
use std::sync::atomic::{AtomicU64, AtomicBool, Ordering};

use crate::config::BufferingMode;

/// Context for a single frame's allocations.
pub struct FrameContext {
    /// Frame number.
    frame_number: u64,
    /// Whether frame is currently active.
    active: AtomicBool,
    /// Arena for this frame's CPU allocations.
    arena: MkFastArena,
    /// Staged upload requests for this frame.
    staged_uploads: Mutex<Vec<StagedUpload>>,
}

/// A staged upload request.
pub(crate) struct StagedUpload {
    pub data_ptr: *const u8,
    pub size: usize,
    pub alignment: usize,
}

// Safety: StagedUpload only contains raw pointers to arena memory
// which is valid for the duration of the frame.
unsafe impl Send for StagedUpload {}
unsafe impl Sync for StagedUpload {}

impl FrameContext {
    /// Create a new frame context.
    pub fn new(arena_size: usize) -> Self {
        Self {
            frame_number: 0,
            active: AtomicBool::new(false),
            arena: MkFastArena::new(arena_size),
            staged_uploads: Mutex::new(Vec::new()),
        }
    }
    
    /// Get the frame number.
    pub fn frame_number(&self) -> u64 {
        self.frame_number
    }
    
    /// Check if frame is active.
    pub fn is_active(&self) -> bool {
        self.active.load(Ordering::Acquire)
    }
    
    /// Begin a new frame.
    pub(crate) fn begin(&mut self, frame_number: u64) {
        self.frame_number = frame_number;
        self.active.store(true, Ordering::Release);
        self.arena.reset();
        self.staged_uploads.lock().clear();
    }
    
    /// End the current frame.
    pub(crate) fn end(&mut self) {
        self.active.store(false, Ordering::Release);
    }
    
    /// Allocate a value in the frame arena.
    /// 
    /// # Safety
    /// 
    /// The returned reference is valid until the frame ends.
    pub fn alloc<T: Default>(&self) -> Option<&mut T> {
        unsafe { self.arena.alloc(T::default()) }
    }
    
    /// Allocate a slice in the frame arena with default values.
    /// 
    /// # Safety
    /// 
    /// The returned reference is valid until the frame ends.
    pub fn alloc_slice<T: Default + Clone>(&self, len: usize) -> Option<&mut [T]> {
        unsafe { self.arena.alloc_slice_default::<T>(len) }
    }
    
    /// Allocate a slice with a fill value.
    /// 
    /// # Safety
    /// 
    /// The returned reference is valid until the frame ends.
    pub fn alloc_slice_fill<T: Clone>(&self, len: usize, value: T) -> Option<&mut [T]> {
        unsafe { self.arena.alloc_slice_fill(len, value) }
    }
    
    /// Stage data for GPU upload.
    pub fn stage<T: Copy>(&self, data: &T) {
        let ptr = data as *const T as *const u8;
        let size = std::mem::size_of::<T>();
        let alignment = std::mem::align_of::<T>();
        
        self.staged_uploads.lock().push(StagedUpload {
            data_ptr: ptr,
            size,
            alignment,
        });
    }
    
    /// Stage a slice for GPU upload.
    pub fn stage_slice<T: Copy>(&self, data: &[T]) {
        let ptr = data.as_ptr() as *const u8;
        let size = std::mem::size_of_val(data);
        let alignment = std::mem::align_of::<T>();
        
        self.staged_uploads.lock().push(StagedUpload {
            data_ptr: ptr,
            size,
            alignment,
        });
    }
    
    /// Get the number of staged uploads.
    pub fn staged_count(&self) -> usize {
        self.staged_uploads.lock().len()
    }
    
    /// Take all staged uploads.
    pub(crate) fn take_staged(&self) -> Vec<StagedUpload> {
        std::mem::take(&mut *self.staged_uploads.lock())
    }
    
    /// Get arena stats (bytes allocated).
    pub fn arena_used(&self) -> usize {
        self.arena.allocated()
    }
    
    /// Get arena capacity.
    pub fn arena_capacity(&self) -> usize {
        self.arena.capacity()
    }
}

/// Manager for multiple frame contexts (double/triple buffering).
pub struct FrameManager {
    contexts: Vec<FrameContext>,
    current_index: AtomicU64,
    frame_counter: AtomicU64,
    buffering_mode: BufferingMode,
}

impl FrameManager {
    /// Create a new frame manager.
    pub fn new(arena_size: usize, buffering_mode: BufferingMode) -> Self {
        let count = buffering_mode.buffer_count();
        let contexts = (0..count)
            .map(|_| FrameContext::new(arena_size))
            .collect();
        
        Self {
            contexts,
            current_index: AtomicU64::new(0),
            frame_counter: AtomicU64::new(0),
            buffering_mode,
        }
    }
    
    /// Get the buffering mode.
    pub fn buffering_mode(&self) -> BufferingMode {
        self.buffering_mode
    }
    
    /// Get the current frame number.
    pub fn current_frame(&self) -> u64 {
        self.frame_counter.load(Ordering::Acquire)
    }
    
    /// Get the current frame context.
    pub fn current(&self) -> &FrameContext {
        let idx = self.current_index.load(Ordering::Acquire) as usize;
        &self.contexts[idx]
    }
    
    /// Get the current frame context mutably.
    pub fn current_mut(&mut self) -> &mut FrameContext {
        let idx = self.current_index.load(Ordering::Acquire) as usize;
        &mut self.contexts[idx]
    }
    
    /// Begin a new frame, rotating to the next context.
    pub fn begin_frame(&mut self) -> &mut FrameContext {
        let frame_num = self.frame_counter.fetch_add(1, Ordering::AcqRel);
        let idx = (frame_num as usize) % self.contexts.len();
        
        self.current_index.store(idx as u64, Ordering::Release);
        self.contexts[idx].begin(frame_num);
        
        &mut self.contexts[idx]
    }
    
    /// End the current frame.
    pub fn end_frame(&mut self) {
        let idx = self.current_index.load(Ordering::Acquire) as usize;
        self.contexts[idx].end();
    }
    
    /// Get a specific frame context by index.
    pub fn context(&self, index: usize) -> Option<&FrameContext> {
        self.contexts.get(index)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_frame_context_alloc() {
        let ctx = FrameContext::new(1024 * 1024);
        
        let value = ctx.alloc::<u32>().unwrap();
        *value = 42;
        assert_eq!(*value, 42);
        
        let slice = ctx.alloc_slice_fill(10, 1.0f32).unwrap();
        assert_eq!(slice.len(), 10);
        assert_eq!(slice[0], 1.0);
    }
    
    #[test]
    fn test_frame_manager_rotation() {
        let mut manager = FrameManager::new(1024 * 1024, BufferingMode::Double);
        
        // Frame 0
        let ctx = manager.begin_frame();
        assert_eq!(ctx.frame_number(), 0);
        manager.end_frame();
        
        // Frame 1
        let ctx = manager.begin_frame();
        assert_eq!(ctx.frame_number(), 1);
        manager.end_frame();
        
        // Frame 2 - should rotate back
        let ctx = manager.begin_frame();
        assert_eq!(ctx.frame_number(), 2);
    }
}
