//! Transfer scheduling and batching.

use memkit_gpu::MkGpuBackend;
use parking_lot::Mutex;
use std::collections::VecDeque;

use crate::sync::{TransferFence, FenceIdGenerator};

/// A pending transfer operation.
pub struct PendingTransfer<B: MkGpuBackend> {
    pub(crate) src: B::BufferHandle,
    pub(crate) dst: B::BufferHandle,
    pub(crate) size: usize,
    pub(crate) fence: TransferFence,
}

/// Scheduler for batching and executing GPU transfers.
pub struct TransferScheduler<B: MkGpuBackend> {
    pending: Mutex<VecDeque<PendingTransfer<B>>>,
    fence_gen: FenceIdGenerator,
    max_pending: usize,
    total_transfers: Mutex<u64>,
    total_bytes: Mutex<u64>,
}

impl<B: MkGpuBackend> TransferScheduler<B> {
    /// Create a new transfer scheduler.
    pub fn new(max_pending: usize) -> Self {
        Self {
            pending: Mutex::new(VecDeque::with_capacity(max_pending)),
            fence_gen: FenceIdGenerator::new(),
            max_pending,
            total_transfers: Mutex::new(0),
            total_bytes: Mutex::new(0),
        }
    }
    
    /// Get the number of pending transfers.
    pub fn pending_count(&self) -> usize {
        self.pending.lock().len()
    }
    
    /// Check if the scheduler should auto-flush.
    pub fn should_flush(&self) -> bool {
        self.pending.lock().len() >= self.max_pending
    }
    
    /// Get total transfer count since creation.
    pub fn total_transfers(&self) -> u64 {
        *self.total_transfers.lock()
    }
    
    /// Get total bytes transferred since creation.
    pub fn total_bytes(&self) -> u64 {
        *self.total_bytes.lock()
    }
    
    /// Queue a transfer operation.
    pub fn queue(
        &self,
        src: B::BufferHandle,
        dst: B::BufferHandle,
        size: usize,
    ) -> TransferFence {
        let fence = TransferFence::new(self.fence_gen.next());
        
        self.pending.lock().push_back(PendingTransfer {
            src,
            dst,
            size,
            fence: fence.clone(),
        });
        
        fence
    }
    
    /// Flush all pending transfers.
    pub fn flush(&self, backend: &B) -> Result<Vec<TransferFence>, B::Error> {
        let transfers: Vec<_> = self.pending.lock().drain(..).collect();
        
        if transfers.is_empty() {
            return Ok(Vec::new());
        }
        
        let mut fences = Vec::with_capacity(transfers.len());
        let mut total_bytes = 0u64;
        
        for transfer in transfers {
            match backend.copy_buffer(&transfer.src, &transfer.dst, transfer.size) {
                Ok(()) => {
                    transfer.fence.mark_complete();
                    total_bytes += transfer.size as u64;
                }
                Err(e) => {
                    transfer.fence.mark_failed();
                    return Err(e);
                }
            }
            fences.push(transfer.fence);
        }
        
        // Update stats
        *self.total_transfers.lock() += fences.len() as u64;
        *self.total_bytes.lock() += total_bytes;
        
        // Wait for completion
        backend.wait_idle()?;
        
        Ok(fences)
    }
    
    /// Flush and wait for all transfers to complete.
    pub fn flush_and_wait(&self, backend: &B) -> Result<(), B::Error> {
        self.flush(backend)?;
        backend.wait_idle()
    }
    
    /// Clear all pending transfers without executing them.
    pub fn clear(&self) {
        let transfers = self.pending.lock().drain(..).collect::<Vec<_>>();
        for transfer in transfers {
            transfer.fence.mark_failed();
        }
    }
}

/// Builder for constructing batch transfers.
pub struct TransferBatch<'a, B: MkGpuBackend> {
    scheduler: &'a TransferScheduler<B>,
    transfers: Vec<(B::BufferHandle, B::BufferHandle, usize)>,
}

impl<'a, B: MkGpuBackend> TransferBatch<'a, B> {
    /// Create a new transfer batch.
    pub fn new(scheduler: &'a TransferScheduler<B>) -> Self {
        Self {
            scheduler,
            transfers: Vec::new(),
        }
    }
    
    /// Add a transfer to the batch.
    pub fn copy(mut self, src: B::BufferHandle, dst: B::BufferHandle, size: usize) -> Self {
        self.transfers.push((src, dst, size));
        self
    }
    
    /// Queue all transfers in the batch.
    pub fn queue(self) -> Vec<TransferFence> {
        self.transfers
            .into_iter()
            .map(|(src, dst, size)| self.scheduler.queue(src, dst, size))
            .collect()
    }
    
    /// Queue and immediately flush all transfers.
    pub fn submit(self, backend: &B) -> Result<Vec<TransferFence>, B::Error> {
        let scheduler = self.scheduler;
        let fences: Vec<_> = self.transfers
            .into_iter()
            .map(|(src, dst, size)| scheduler.queue(src, dst, size))
            .collect();
        scheduler.flush(backend)?;
        Ok(fences)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use memkit_gpu::{DummyBackend, DummyBackendConfig, MkBufferUsage, MkMemoryType};
    
    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_scheduler_queue_and_flush() {
        let backend = create_backend();
        let scheduler: TransferScheduler<DummyBackend> = TransferScheduler::new(64);
        
        let src = backend.create_buffer(1024, MkBufferUsage::TRANSFER_SRC, MkMemoryType::HostVisible).unwrap();
        let dst = backend.create_buffer(1024, MkBufferUsage::TRANSFER_DST, MkMemoryType::DeviceLocal).unwrap();
        
        let fence = scheduler.queue(src.clone(), dst.clone(), 1024);
        assert!(fence.is_pending());
        assert_eq!(scheduler.pending_count(), 1);
        
        scheduler.flush(&backend).unwrap();
        assert!(fence.is_complete());
        assert_eq!(scheduler.pending_count(), 0);
        
        backend.destroy_buffer(&src);
        backend.destroy_buffer(&dst);
    }
}
