//! Weight Streaming for Large Models
//!
//! This module provides on-demand weight loading for models larger than GPU memory,
//! using an LRU cache to keep frequently used layers in GPU memory while streaming
//! others from host memory or disk.

use crate::device::MetalDevice;
use crate::error::{MetalError, Result};
use crate::memory::{MetalBuffer, MetalBufferUsage};
use std::collections::{HashMap, VecDeque};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Instant;
use parking_lot::{Mutex, RwLock};
use tracing::{debug, trace, warn};

/// Unique identifier for a layer's weights
pub type LayerId = u64;

/// Source for weight data
#[derive(Debug, Clone)]
pub enum WeightSource {
    /// Memory-mapped file
    MmappedFile {
        /// Path to the file
        path: PathBuf,
        /// Offset in file
        offset: usize,
        /// Size in bytes
        size: usize,
    },
    /// Host memory
    HostMemory {
        /// Data in host memory
        data: Arc<Vec<f16>>,
    },
    /// Already in GPU memory (no streaming needed)
    GpuResident,
}

/// Half-precision float for compact storage
#[allow(non_camel_case_types)]
pub type f16 = u16; // Use u16 as f16 placeholder

/// Configuration for weight streaming
#[derive(Debug, Clone)]
pub struct StreamingConfig {
    /// Maximum GPU memory budget for weights (bytes)
    pub gpu_budget: usize,
    /// Number of layers to prefetch ahead
    pub prefetch_lookahead: usize,
    /// Enable async prefetching
    pub async_prefetch: bool,
    /// Eviction policy
    pub eviction_policy: EvictionPolicy,
    /// Pin certain layers (never evict)
    pub pinned_layers: Vec<LayerId>,
}

impl Default for StreamingConfig {
    fn default() -> Self {
        Self {
            gpu_budget: 4 * 1024 * 1024 * 1024, // 4GB
            prefetch_lookahead: 2,
            async_prefetch: true,
            eviction_policy: EvictionPolicy::Lru,
            pinned_layers: Vec::new(),
        }
    }
}

/// Policy for evicting weights from GPU cache
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum EvictionPolicy {
    /// Least Recently Used
    Lru,
    /// Least Frequently Used
    Lfu,
    /// First In First Out
    Fifo,
}

/// Statistics for weight streaming
#[derive(Debug, Default, Clone)]
pub struct StreamingStats {
    /// Total weight loads from source
    pub total_loads: u64,
    /// Cache hits (already in GPU)
    pub cache_hits: u64,
    /// Cache misses (needed to load)
    pub cache_misses: u64,
    /// Total bytes loaded
    pub bytes_loaded: u64,
    /// Total bytes evicted
    pub bytes_evicted: u64,
    /// Average load time (microseconds)
    pub avg_load_time_us: f64,
    /// Current cache size (bytes)
    pub current_cache_size: usize,
    /// Number of cached layers
    pub cached_layers: usize,
}

impl StreamingStats {
    /// Get cache hit rate
    pub fn hit_rate(&self) -> f64 {
        let total = self.cache_hits + self.cache_misses;
        if total == 0 {
            0.0
        } else {
            self.cache_hits as f64 / total as f64
        }
    }
}

/// Entry in the weight cache
struct CacheEntry {
    /// GPU buffer containing weights
    buffer: MetalBuffer<f16>,
    /// Size in bytes
    size: usize,
    /// Last access time
    last_access: Instant,
    /// Access count (for LFU)
    access_count: u64,
    /// Whether this entry is pinned
    pinned: bool,
    /// Source for reloading
    source: WeightSource,
}

/// Weight Streaming Manager
pub struct WeightStreamingManager {
    /// Device
    device: Arc<MetalDevice>,
    /// Configuration
    config: StreamingConfig,
    /// LRU cache of weights in GPU memory
    cache: RwLock<HashMap<LayerId, CacheEntry>>,
    /// LRU order tracking
    lru_order: Mutex<VecDeque<LayerId>>,
    /// Current cache size in bytes
    current_size: Mutex<usize>,
    /// Statistics
    stats: RwLock<StreamingStats>,
    /// Prefetch queue
    prefetch_queue: Mutex<VecDeque<LayerId>>,
    /// Weight sources
    sources: RwLock<HashMap<LayerId, WeightSource>>,
}

impl WeightStreamingManager {
    /// Create a new weight streaming manager
    pub fn new(device: Arc<MetalDevice>, config: StreamingConfig) -> Self {
        Self {
            device,
            config,
            cache: RwLock::new(HashMap::new()),
            lru_order: Mutex::new(VecDeque::new()),
            current_size: Mutex::new(0),
            stats: RwLock::new(StreamingStats::default()),
            prefetch_queue: Mutex::new(VecDeque::new()),
            sources: RwLock::new(HashMap::new()),
        }
    }

    /// Create with default configuration
    pub fn default_config(device: Arc<MetalDevice>) -> Self {
        Self::new(device, StreamingConfig::default())
    }

    /// Register a weight source
    pub fn register_weights(&self, layer_id: LayerId, source: WeightSource) {
        let mut sources = self.sources.write();
        sources.insert(layer_id, source);
    }

    /// Register weights from host memory
    pub fn register_from_host(&self, layer_id: LayerId, data: Vec<f16>) {
        self.register_weights(
            layer_id,
            WeightSource::HostMemory {
                data: Arc::new(data),
            },
        );
    }

    /// Register weights from file
    pub fn register_from_file(&self, layer_id: LayerId, path: PathBuf, offset: usize, size: usize) {
        self.register_weights(
            layer_id,
            WeightSource::MmappedFile { path, offset, size },
        );
    }

    /// Get weights for a layer, loading from source if needed
    pub fn get_weights(&self, layer_id: LayerId) -> Result<WeightRef> {
        // Check if already cached
        {
            let cache = self.cache.read();
            if cache.contains_key(&layer_id) {
                // Update stats and access time
                self.record_hit(layer_id);
                return Ok(WeightRef {
                    manager: self,
                    layer_id,
                });
            }
        }

        // Cache miss - need to load
        self.load_weights(layer_id)?;
        Ok(WeightRef {
            manager: self,
            layer_id,
        })
    }

    /// Check if weights are cached
    pub fn is_cached(&self, layer_id: LayerId) -> bool {
        self.cache.read().contains_key(&layer_id)
    }

    /// Load weights from source into GPU cache
    fn load_weights(&self, layer_id: LayerId) -> Result<()> {
        let source = {
            let sources = self.sources.read();
            sources.get(&layer_id).cloned()
        };

        let source = source.ok_or_else(|| {
            MetalError::internal(format!("No source registered for layer {}", layer_id))
        })?;

        let start = Instant::now();

        let (buffer, size) = match &source {
            WeightSource::HostMemory { data } => {
                let size = data.len() * std::mem::size_of::<f16>();

                // Evict if needed
                self.evict_for_space(size)?;

                // Create GPU buffer
                let buffer = MetalBuffer::<f16>::new(
                    &self.device,
                    data.len(),
                    MetalBufferUsage::Shared,
                )?;

                // Copy data
                buffer.write(data)?;

                (buffer, size)
            }
            WeightSource::MmappedFile { path, offset, size } => {
                // Evict if needed
                self.evict_for_space(*size)?;

                // In a real implementation, we would mmap the file
                // For now, we read it conventionally
                let data = std::fs::read(path).map_err(|e| MetalError::internal(e.to_string()))?;

                let num_elements = (*size) / std::mem::size_of::<f16>();
                let buffer = MetalBuffer::<f16>::new(
                    &self.device,
                    num_elements,
                    MetalBufferUsage::Shared,
                )?;

                // Copy relevant portion
                let start_elem = (*offset) / std::mem::size_of::<f16>();
                let end_elem = start_elem + num_elements;

                if end_elem * 2 <= data.len() {
                    let slice: Vec<f16> = data[*offset..*offset + *size]
                        .chunks(2)
                        .map(|c| u16::from_le_bytes([c[0], c.get(1).copied().unwrap_or(0)]))
                        .collect();
                    buffer.write(&slice)?;
                }

                (buffer, *size)
            }
            WeightSource::GpuResident => {
                return Err(MetalError::internal("Cannot load GPU-resident weights"));
            }
        };

        let duration = start.elapsed();

        // Add to cache
        {
            let mut cache = self.cache.write();
            cache.insert(
                layer_id,
                CacheEntry {
                    buffer,
                    size,
                    last_access: Instant::now(),
                    access_count: 1,
                    pinned: self.config.pinned_layers.contains(&layer_id),
                    source,
                },
            );
        }

        // Update LRU order
        {
            let mut lru = self.lru_order.lock();
            lru.push_back(layer_id);
        }

        // Update size tracking
        {
            let mut current = self.current_size.lock();
            *current += size;
        }

        // Update stats
        {
            let mut stats = self.stats.write();
            stats.total_loads += 1;
            stats.cache_misses += 1;
            stats.bytes_loaded += size as u64;
            stats.cached_layers += 1;
            stats.current_cache_size += size;

            // Update average load time
            let total = stats.total_loads;
            stats.avg_load_time_us = (stats.avg_load_time_us * (total - 1) as f64
                + duration.as_micros() as f64)
                / total as f64;
        }

        debug!(
            "Loaded weights for layer {} ({} bytes) in {:?}",
            layer_id, size, duration
        );

        Ok(())
    }

    /// Evict weights to make space
    fn evict_for_space(&self, needed: usize) -> Result<()> {
        let current = *self.current_size.lock();

        if current + needed <= self.config.gpu_budget {
            return Ok(());
        }

        let to_free = current + needed - self.config.gpu_budget;
        let mut freed = 0;

        while freed < to_free {
            let evict_id = self.select_for_eviction();

            match evict_id {
                Some(id) => {
                    let evicted_size = self.evict(id)?;
                    freed += evicted_size;
                }
                None => {
                    warn!("Cannot free enough space for weights: need {}, freed {}", to_free, freed);
                    break;
                }
            }
        }

        Ok(())
    }

    /// Select a layer for eviction based on policy
    fn select_for_eviction(&self) -> Option<LayerId> {
        match self.config.eviction_policy {
            EvictionPolicy::Lru => {
                let lru = self.lru_order.lock();
                let cache = self.cache.read();

                // Find first non-pinned entry
                lru.iter()
                    .find(|&&id| {
                        cache.get(&id).map(|e| !e.pinned).unwrap_or(false)
                    })
                    .copied()
            }
            EvictionPolicy::Lfu => {
                let cache = self.cache.read();
                cache
                    .iter()
                    .filter(|(_, e)| !e.pinned)
                    .min_by_key(|(_, e)| e.access_count)
                    .map(|(&id, _)| id)
            }
            EvictionPolicy::Fifo => {
                let lru = self.lru_order.lock();
                let cache = self.cache.read();

                lru.iter()
                    .find(|&&id| {
                        cache.get(&id).map(|e| !e.pinned).unwrap_or(false)
                    })
                    .copied()
            }
        }
    }

    /// Evict a specific layer
    fn evict(&self, layer_id: LayerId) -> Result<usize> {
        let size = {
            let mut cache = self.cache.write();
            if let Some(entry) = cache.remove(&layer_id) {
                entry.size
            } else {
                return Ok(0);
            }
        };

        // Remove from LRU order
        {
            let mut lru = self.lru_order.lock();
            lru.retain(|&id| id != layer_id);
        }

        // Update size tracking
        {
            let mut current = self.current_size.lock();
            *current = current.saturating_sub(size);
        }

        // Update stats
        {
            let mut stats = self.stats.write();
            stats.bytes_evicted += size as u64;
            stats.cached_layers = stats.cached_layers.saturating_sub(1);
            stats.current_cache_size = stats.current_cache_size.saturating_sub(size);
        }

        trace!("Evicted layer {} ({} bytes)", layer_id, size);

        Ok(size)
    }

    /// Record a cache hit
    fn record_hit(&self, layer_id: LayerId) {
        // Update access time and count
        {
            let mut cache = self.cache.write();
            if let Some(entry) = cache.get_mut(&layer_id) {
                entry.last_access = Instant::now();
                entry.access_count += 1;
            }
        }

        // Move to back of LRU queue
        {
            let mut lru = self.lru_order.lock();
            lru.retain(|&id| id != layer_id);
            lru.push_back(layer_id);
        }

        // Update stats
        {
            let mut stats = self.stats.write();
            stats.cache_hits += 1;
        }
    }

    /// Request async prefetch for upcoming layers
    pub fn prefetch(&self, layer_ids: &[LayerId]) {
        let mut queue = self.prefetch_queue.lock();
        for &id in layer_ids {
            if !self.is_cached(id) && !queue.contains(&id) {
                queue.push_back(id);
            }
        }
    }

    /// Prefetch layers ahead of current
    pub fn prefetch_ahead(&self, current_layer: LayerId) {
        let lookahead = self.config.prefetch_lookahead;
        let layers: Vec<LayerId> = (current_layer + 1..=current_layer + lookahead as u64).collect();
        self.prefetch(&layers);
    }

    /// Process prefetch queue
    pub fn process_prefetch(&self) -> Result<usize> {
        let mut processed = 0;

        loop {
            let layer_id = {
                let mut queue = self.prefetch_queue.lock();
                queue.pop_front()
            };

            match layer_id {
                Some(id) => {
                    if !self.is_cached(id) {
                        if let Err(e) = self.load_weights(id) {
                            warn!("Prefetch failed for layer {}: {}", id, e);
                        } else {
                            processed += 1;
                        }
                    }
                }
                None => break,
            }
        }

        Ok(processed)
    }

    /// Pin a layer (prevent eviction)
    pub fn pin_layer(&self, layer_id: LayerId) {
        let mut cache = self.cache.write();
        if let Some(entry) = cache.get_mut(&layer_id) {
            entry.pinned = true;
        }
    }

    /// Unpin a layer
    pub fn unpin_layer(&self, layer_id: LayerId) {
        let mut cache = self.cache.write();
        if let Some(entry) = cache.get_mut(&layer_id) {
            entry.pinned = false;
        }
    }

    /// Clear cache
    pub fn clear(&self) {
        let mut cache = self.cache.write();
        cache.clear();

        let mut lru = self.lru_order.lock();
        lru.clear();

        let mut current = self.current_size.lock();
        *current = 0;

        let mut stats = self.stats.write();
        stats.cached_layers = 0;
        stats.current_cache_size = 0;
    }

    /// Get statistics
    pub fn stats(&self) -> StreamingStats {
        self.stats.read().clone()
    }

    /// Get configuration
    pub fn config(&self) -> &StreamingConfig {
        &self.config
    }

    /// Get current cache size
    pub fn cache_size(&self) -> usize {
        *self.current_size.lock()
    }

    /// Get number of cached layers
    pub fn cached_count(&self) -> usize {
        self.cache.read().len()
    }
}

impl std::fmt::Debug for WeightStreamingManager {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let stats = self.stats();
        f.debug_struct("WeightStreamingManager")
            .field("cached_layers", &stats.cached_layers)
            .field("cache_size", &stats.current_cache_size)
            .field("hit_rate", &format!("{:.1}%", stats.hit_rate() * 100.0))
            .finish()
    }
}

/// Reference to cached weights
pub struct WeightRef<'a> {
    manager: &'a WeightStreamingManager,
    layer_id: LayerId,
}

impl<'a> WeightRef<'a> {
    /// Get the buffer
    pub fn buffer(&self) -> Option<impl std::ops::Deref<Target = MetalBuffer<f16>> + '_> {
        let cache = self.manager.cache.read();
        if cache.contains_key(&self.layer_id) {
            Some(parking_lot::RwLockReadGuard::map(cache, |c| {
                &c.get(&self.layer_id).unwrap().buffer
            }))
        } else {
            None
        }
    }

    /// Get the layer ID
    pub fn layer_id(&self) -> LayerId {
        self.layer_id
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_streaming_config_default() {
        let config = StreamingConfig::default();
        assert_eq!(config.gpu_budget, 4 * 1024 * 1024 * 1024);
        assert_eq!(config.prefetch_lookahead, 2);
    }

    #[test]
    fn test_streaming_stats_hit_rate() {
        let mut stats = StreamingStats::default();
        stats.cache_hits = 80;
        stats.cache_misses = 20;
        assert!((stats.hit_rate() - 0.8).abs() < 0.001);
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_streaming_manager_creation() {
        if !MetalDevice::is_available() {
            return;
        }

        let device = Arc::new(MetalDevice::system_default().unwrap());
        let manager = WeightStreamingManager::default_config(device);

        assert_eq!(manager.cached_count(), 0);
        assert_eq!(manager.cache_size(), 0);
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_register_and_load() {
        if !MetalDevice::is_available() {
            return;
        }

        let device = Arc::new(MetalDevice::system_default().unwrap());
        let manager = WeightStreamingManager::default_config(device);

        // Register weights from host
        let data: Vec<f16> = (0..1024).map(|i| i as f16).collect();
        manager.register_from_host(1, data);

        // Load weights
        let _weight_ref = manager.get_weights(1).unwrap();
        assert!(manager.is_cached(1));
        assert_eq!(manager.cached_count(), 1);
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_cache_hit() {
        if !MetalDevice::is_available() {
            return;
        }

        let device = Arc::new(MetalDevice::system_default().unwrap());
        let manager = WeightStreamingManager::default_config(device);

        let data: Vec<f16> = (0..1024).map(|i| i as f16).collect();
        manager.register_from_host(1, data);

        // First access - miss
        let _ref1 = manager.get_weights(1).unwrap();

        // Second access - hit
        let _ref2 = manager.get_weights(1).unwrap();

        let stats = manager.stats();
        assert_eq!(stats.cache_misses, 1);
        assert_eq!(stats.cache_hits, 1);
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_eviction() {
        if !MetalDevice::is_available() {
            return;
        }

        let device = Arc::new(MetalDevice::system_default().unwrap());
        let config = StreamingConfig {
            gpu_budget: 4096, // Very small budget
            ..Default::default()
        };
        let manager = WeightStreamingManager::new(device, config);

        // Register multiple layers that exceed budget
        for i in 0..5 {
            let data: Vec<f16> = (0..512).map(|j| (i * 512 + j) as f16).collect();
            manager.register_from_host(i as LayerId, data);
            let _ = manager.get_weights(i as LayerId);
        }

        // Should have evicted some
        let stats = manager.stats();
        assert!(stats.bytes_evicted > 0);
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_pinning() {
        if !MetalDevice::is_available() {
            return;
        }

        let device = Arc::new(MetalDevice::system_default().unwrap());
        let config = StreamingConfig {
            gpu_budget: 2048,
            ..Default::default()
        };
        let manager = WeightStreamingManager::new(device, config);

        // Load and pin first layer
        let data: Vec<f16> = (0..256).map(|i| i as f16).collect();
        manager.register_from_host(1, data);
        let _ = manager.get_weights(1).unwrap();
        manager.pin_layer(1);

        // Load more layers to trigger eviction
        for i in 2..5 {
            let data: Vec<f16> = (0..256).map(|j| (i * 256 + j) as f16).collect();
            manager.register_from_host(i as LayerId, data);
            let _ = manager.get_weights(i as LayerId);
        }

        // Pinned layer should still be cached
        assert!(manager.is_cached(1));
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_prefetch() {
        if !MetalDevice::is_available() {
            return;
        }

        let device = Arc::new(MetalDevice::system_default().unwrap());
        let manager = WeightStreamingManager::default_config(device);

        // Register layers
        for i in 0..5 {
            let data: Vec<f16> = (0..256).map(|j| (i * 256 + j) as f16).collect();
            manager.register_from_host(i as LayerId, data);
        }

        // Request prefetch
        manager.prefetch(&[1, 2, 3]);

        // Process prefetch
        let processed = manager.process_prefetch().unwrap();
        assert_eq!(processed, 3);

        // All should be cached
        assert!(manager.is_cached(1));
        assert!(manager.is_cached(2));
        assert!(manager.is_cached(3));
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_clear() {
        if !MetalDevice::is_available() {
            return;
        }

        let device = Arc::new(MetalDevice::system_default().unwrap());
        let manager = WeightStreamingManager::default_config(device);

        // Load some weights
        for i in 0..3 {
            let data: Vec<f16> = (0..256).map(|j| (i * 256 + j) as f16).collect();
            manager.register_from_host(i as LayerId, data);
            let _ = manager.get_weights(i as LayerId);
        }

        assert!(manager.cached_count() > 0);

        manager.clear();
        assert_eq!(manager.cached_count(), 0);
        assert_eq!(manager.cache_size(), 0);
    }
}