//! Main coordinator implementation.

use memkit_gpu::{
    MkGpu, MkGpuBackend, MkBufferUsage,
    MkDeviceBuffer, MkStagingBuffer,
};
use parking_lot::RwLock;
use std::sync::atomic::{AtomicU64, Ordering};

use crate::config::{CoordinatorConfig, BufferingMode};
use crate::frame::{FrameManager, FrameContext};
use crate::scheduler::TransferScheduler;
use crate::staging::{DynamicStagingPool, StagingHandle};
use crate::sync::TransferFence;

/// The main CPU-GPU memory coordinator.
///
/// Bridges memkit CPU allocators with memkit-gpu GPU memory management,
/// providing unified APIs for efficient data transfers.
pub struct MkCoordinator<B: MkGpuBackend> {
    /// GPU allocator handle.
    gpu: MkGpu<B>,
    /// Configuration.
    config: CoordinatorConfig,
    /// Frame manager for double/triple buffering.
    frames: RwLock<FrameManager>,
    /// Transfer scheduler.
    scheduler: TransferScheduler<B>,
    /// Dynamic staging pool.
    staging_pool: DynamicStagingPool<B>,
    /// Statistics.
    stats: CoordinatorStats,
}

/// Statistics tracked by the coordinator.
pub struct CoordinatorStats {
    frames_completed: AtomicU64,
    uploads_completed: AtomicU64,
    bytes_uploaded: AtomicU64,
}

impl Default for CoordinatorStats {
    fn default() -> Self {
        Self {
            frames_completed: AtomicU64::new(0),
            uploads_completed: AtomicU64::new(0),
            bytes_uploaded: AtomicU64::new(0),
        }
    }
}

impl CoordinatorStats {
    /// Get total frames completed.
    pub fn frames_completed(&self) -> u64 {
        self.frames_completed.load(Ordering::Relaxed)
    }
    
    /// Get total uploads completed.
    pub fn uploads_completed(&self) -> u64 {
        self.uploads_completed.load(Ordering::Relaxed)
    }
    
    /// Get total bytes uploaded.
    pub fn bytes_uploaded(&self) -> u64 {
        self.bytes_uploaded.load(Ordering::Relaxed)
    }
}

impl<B: MkGpuBackend> MkCoordinator<B> {
    /// Create a new coordinator with the given backend and configuration.
    pub fn new(backend: B, config: CoordinatorConfig) -> Self {
        let frames = FrameManager::new(config.frame_arena_size, config.buffering);
        let scheduler = TransferScheduler::new(config.max_pending_transfers);
        
        Self {
            gpu: MkGpu::new(backend),
            config,
            frames: RwLock::new(frames),
            scheduler,
            staging_pool: DynamicStagingPool::new(),
            stats: CoordinatorStats::default(),
        }
    }
    
    /// Create a coordinator with default configuration.
    pub fn with_backend(backend: B) -> Self {
        Self::new(backend, CoordinatorConfig::default())
    }
    
    /// Create a coordinator with specific buffering mode.
    pub fn with_buffering(backend: B, mode: BufferingMode) -> Self {
        let config = CoordinatorConfig::default().with_buffering(mode);
        Self::new(backend, config)
    }
    
    /// Get a reference to the underlying GPU allocator.
    pub fn gpu(&self) -> &MkGpu<B> {
        &self.gpu
    }
    
    /// Get a reference to the backend.
    pub fn backend(&self) -> &B {
        self.gpu.backend()
    }
    
    /// Get the configuration.
    pub fn config(&self) -> &CoordinatorConfig {
        &self.config
    }
    
    /// Get coordinator statistics.
    pub fn stats(&self) -> &CoordinatorStats {
        &self.stats
    }
    
    // =========================================================================
    // Frame Management
    // =========================================================================
    
    /// Begin a new frame.
    ///
    /// This rotates to the next frame buffer (for double/triple buffering)
    /// and resets the frame arena.
    pub fn begin_frame(&self) {
        self.frames.write().begin_frame();
    }
    
    /// End the current frame.
    ///
    /// If `auto_flush_on_frame_end` is enabled (default), this will
    /// flush all staged transfers.
    pub fn end_frame(&self) -> Result<(), B::Error> {
        // Process staged uploads from frame context
        let staged = self.frames.read().current().take_staged();
        
        for upload in staged {
            // Create staging buffer and copy data
            let staging = self.staging_pool.acquire(upload.size, self.gpu.backend())?;
            
            if let Some(ptr) = self.gpu.backend().map(staging.handle()) {
                unsafe {
                    std::ptr::copy_nonoverlapping(upload.data_ptr, ptr, upload.size);
                }
                self.gpu.backend().unmap(staging.handle());
            }
            
            // Note: In a real implementation, we'd track the destination buffer
            // For now, just release staging back to pool
            self.staging_pool.release(staging, self.gpu.backend());
        }
        
        // Flush pending transfers if configured
        if self.config.auto_flush_on_frame_end {
            self.flush_transfers()?;
        }
        
        self.frames.write().end_frame();
        self.stats.frames_completed.fetch_add(1, Ordering::Relaxed);
        
        Ok(())
    }
    
    /// Get the current frame number.
    pub fn current_frame(&self) -> u64 {
        self.frames.read().current_frame()
    }
    
    /// Access the current frame context for allocations.
    pub fn with_frame<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&FrameContext) -> R,
    {
        let frames = self.frames.read();
        f(frames.current())
    }
    
    // =========================================================================
    // Frame Allocations
    // =========================================================================
    
    /// Allocate a value in the current frame's arena.
    pub fn frame_alloc<T: Default>(&self) -> Option<&mut T> {
        // Safety: The frame arena is valid for the duration of the frame.
        // The returned reference is valid until end_frame() is called.
        let frames = self.frames.read();
        let ptr = frames.current().alloc::<T>()? as *mut T;
        unsafe { Some(&mut *ptr) }
    }
    
    /// Allocate a slice in the current frame's arena.
    pub fn frame_alloc_slice<T: Default + Clone>(&self, len: usize) -> Option<&mut [T]> {
        let frames = self.frames.read();
        let slice = frames.current().alloc_slice::<T>(len)?;
        let ptr = slice.as_mut_ptr();
        let len = slice.len();
        unsafe { Some(std::slice::from_raw_parts_mut(ptr, len)) }
    }
    
    /// Allocate a slice with a fill value in the current frame's arena.
    pub fn frame_alloc_slice_fill<T: Clone>(&self, len: usize, value: T) -> Option<&mut [T]> {
        let frames = self.frames.read();
        let slice = frames.current().alloc_slice_fill(len, value)?;
        let ptr = slice.as_mut_ptr();
        let len = slice.len();
        unsafe { Some(std::slice::from_raw_parts_mut(ptr, len)) }
    }
    
    // =========================================================================
    // Direct Uploads
    // =========================================================================
    
    /// Upload a slice directly to a new GPU buffer.
    ///
    /// Creates a staging buffer, copies data, transfers to device buffer,
    /// and returns the device buffer handle.
    pub fn upload_slice<T: Copy>(
        &self,
        data: &[T],
    ) -> Result<MkDeviceBuffer<B>, B::Error> {
        let staging = self.gpu.staging_buffer_with_data(data)?;
        let size = std::mem::size_of_val(data);
        
        let device = self.gpu.create_device_buffer(
            size,
            MkBufferUsage::STORAGE | MkBufferUsage::TRANSFER_DST,
        )?;
        
        self.gpu.transfer(&staging, &device)?;
        self.gpu.wait_idle()?;
        
        self.stats.uploads_completed.fetch_add(1, Ordering::Relaxed);
        self.stats.bytes_uploaded.fetch_add(size as u64, Ordering::Relaxed);
        
        Ok(device)
    }
    
    /// Upload a slice to an existing device buffer.
    pub fn upload_to_buffer<T: Copy>(
        &self,
        data: &[T],
        device: &MkDeviceBuffer<B>,
    ) -> Result<(), B::Error> {
        let staging = self.gpu.staging_buffer_with_data(data)?;
        let size = std::mem::size_of_val(data);
        
        self.gpu.transfer(&staging, device)?;
        
        self.stats.uploads_completed.fetch_add(1, Ordering::Relaxed);
        self.stats.bytes_uploaded.fetch_add(size as u64, Ordering::Relaxed);
        
        Ok(())
    }
    
    /// Upload raw bytes to a new GPU buffer.
    pub fn upload_bytes(&self, data: &[u8]) -> Result<MkDeviceBuffer<B>, B::Error> {
        self.upload_slice(data)
    }
    
    // =========================================================================
    // Staged Uploads (Batched)
    // =========================================================================
    
    /// Stage data for upload (will be transferred on flush or end_frame).
    pub fn stage_upload<T: Copy>(
        &self,
        data: &[T],
        device: &MkDeviceBuffer<B>,
    ) -> Result<TransferFence, B::Error> {
        let staging = self.gpu.staging_buffer_with_data(data)?;
        let size = std::mem::size_of_val(data);
        
        let fence = self.scheduler.queue(
            staging.handle().clone(),
            device.handle().clone(),
            size,
        );
        
        Ok(fence)
    }
    
    /// Flush all pending staged transfers.
    pub fn flush_transfers(&self) -> Result<Vec<TransferFence>, B::Error> {
        let fences = self.scheduler.flush(self.gpu.backend())?;
        
        let count = fences.len();
        let bytes: u64 = fences.iter().map(|_| 0u64).sum(); // Would need to track bytes
        
        self.stats.uploads_completed.fetch_add(count as u64, Ordering::Relaxed);
        
        Ok(fences)
    }
    
    /// Get the number of pending transfers.
    pub fn pending_transfer_count(&self) -> usize {
        self.scheduler.pending_count()
    }
    
    // =========================================================================
    // Staging Pool Access
    // =========================================================================
    
    /// Acquire a staging buffer from the pool.
    pub fn acquire_staging(&self, size: usize) -> Result<StagingHandle<B>, B::Error> {
        self.staging_pool.acquire(size, self.gpu.backend())
    }
    
    /// Release a staging buffer back to the pool.
    pub fn release_staging(&self, staging: StagingHandle<B>) {
        self.staging_pool.release(staging, self.gpu.backend());
    }
    
    /// Transfer from a staging handle to a device buffer, then release staging.
    pub fn transfer_and_release(
        &self,
        staging: StagingHandle<B>,
        device: &MkDeviceBuffer<B>,
    ) -> Result<(), B::Error> {
        let size = staging.size().min(device.size());
        
        self.gpu.backend().copy_buffer(staging.handle(), device.handle(), size)?;
        self.staging_pool.release(staging, self.gpu.backend());
        
        self.stats.uploads_completed.fetch_add(1, Ordering::Relaxed);
        self.stats.bytes_uploaded.fetch_add(size as u64, Ordering::Relaxed);
        
        Ok(())
    }
    
    // =========================================================================
    // Buffer Creation
    // =========================================================================
    
    /// Create a device buffer.
    pub fn create_device_buffer(
        &self,
        size: usize,
        usage: MkBufferUsage,
    ) -> Result<MkDeviceBuffer<B>, B::Error> {
        self.gpu.create_device_buffer(size, usage)
    }
    
    /// Create a staging buffer.
    pub fn create_staging_buffer(&self, size: usize) -> Result<MkStagingBuffer<B>, B::Error> {
        self.gpu.create_staging_buffer(size)
    }
    
    /// Create a staging buffer with data.
    pub fn create_staging_with_data<T: Copy>(
        &self,
        data: &[T],
    ) -> Result<MkStagingBuffer<B>, B::Error> {
        self.gpu.staging_buffer_with_data(data)
    }
    
    // =========================================================================
    // Synchronization
    // =========================================================================
    
    /// Wait for the GPU to become idle.
    pub fn wait_idle(&self) -> Result<(), B::Error> {
        self.gpu.wait_idle()
    }
    
    /// Wait for a specific transfer to complete.
    pub fn wait_for_fence(&self, fence: &TransferFence) -> Result<(), B::Error> {
        // For now, just wait for idle
        // A real implementation would track fences properly
        if fence.is_pending() {
            self.gpu.wait_idle()?;
        }
        Ok(())
    }
}

impl<B: MkGpuBackend> Drop for MkCoordinator<B> {
    fn drop(&mut self) {
        // Clear staging pool
        self.staging_pool.clear(self.gpu.backend());
        
        // Clear pending transfers
        self.scheduler.clear();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use memkit_gpu::{DummyBackend, DummyBackendConfig};
    
    fn create_backend() -> DummyBackend {
        DummyBackend::with_config(DummyBackendConfig {
            max_buffer_size: 64 * 1024 * 1024,
            simulate_device_local: true,
            transfer_delay_us: 0,
        })
    }
    
    #[test]
    fn test_coordinator_creation() {
        let backend = create_backend();
        let coordinator = MkCoordinator::with_backend(backend);
        
        assert_eq!(coordinator.current_frame(), 0);
    }
    
    #[test]
    fn test_coordinator_upload_slice() {
        let backend = create_backend();
        let coordinator = MkCoordinator::with_backend(backend);
        
        let data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0];
        let buffer = coordinator.upload_slice(&data).unwrap();
        
        assert_eq!(buffer.size(), data.len() * std::mem::size_of::<f32>());
        assert_eq!(coordinator.stats().uploads_completed(), 1);
    }
    
    #[test]
    fn test_coordinator_frame_lifecycle() {
        let backend = create_backend();
        let coordinator = MkCoordinator::with_backend(backend);
        
        coordinator.begin_frame();
        assert_eq!(coordinator.current_frame(), 1);
        
        coordinator.end_frame().unwrap();
        assert_eq!(coordinator.stats().frames_completed(), 1);
    }
    
    #[test]
    fn test_coordinator_staging_pool() {
        let backend = create_backend();
        let coordinator = MkCoordinator::with_backend(backend);
        
        let staging = coordinator.acquire_staging(1024).unwrap();
        assert!(staging.size() >= 1024);
        
        coordinator.release_staging(staging);
    }
}
