//! Unified memory management for heterogeneous compute.

use crate::error::ComputeError;
use crate::processor::ProcessorType;
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Tensor handle for memory management.
#[derive(Clone, Debug)]
pub struct TensorHandle {
    /// Unique ID.
    pub id: TensorId,
    /// Shape.
    pub shape: Vec<usize>,
    /// Data type.
    pub dtype: DataType,
    /// Size in bytes.
    pub size_bytes: u64,
    /// Current locations.
    pub locations: Vec<ProcessorType>,
}

impl TensorHandle {
    /// Creates a new tensor handle.
    pub fn new(shape: Vec<usize>, dtype: DataType) -> Self {
        let size_bytes = shape.iter().product::<usize>() as u64 * dtype.size_bytes() as u64;
        Self {
            id: TensorId::new(),
            shape,
            dtype,
            size_bytes,
            locations: Vec::new(),
        }
    }

    /// Gets the number of elements.
    pub fn numel(&self) -> usize {
        self.shape.iter().product()
    }
}

/// Tensor identifier.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct TensorId(pub u64);

impl TensorId {
    /// Creates a new tensor ID.
    pub fn new() -> Self {
        use rand::Rng;
        TensorId(rand::thread_rng().gen())
    }
}

impl Default for TensorId {
    fn default() -> Self {
        Self::new()
    }
}

/// Data types for tensors.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum DataType {
    Float64,
    Float32,
    Float16,
    BFloat16,
    Int64,
    Int32,
    Int16,
    Int8,
    UInt8,
    Bool,
}

impl DataType {
    /// Returns size in bytes.
    pub fn size_bytes(&self) -> usize {
        match self {
            DataType::Float64 | DataType::Int64 => 8,
            DataType::Float32 | DataType::Int32 => 4,
            DataType::Float16 | DataType::BFloat16 | DataType::Int16 => 2,
            DataType::Int8 | DataType::UInt8 | DataType::Bool => 1,
        }
    }
}

/// Data transfer path between processors.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum TransferPath {
    /// Direct GPU-to-GPU via NVLink.
    NvLink,
    /// Direct GPU-to-GPU via PCIe P2P.
    PciePeerToPeer,
    /// Through CPU memory.
    CpuMediated,
    /// Unified memory (Apple Silicon).
    UnifiedMemory,
    /// Network transfer.
    Network,
    /// Same memory space (no transfer needed).
    SameMemory,
}

impl TransferPath {
    /// Returns approximate bandwidth in GB/s.
    pub fn bandwidth_gbps(&self) -> f64 {
        match self {
            TransferPath::NvLink => 900.0,       // NVLink 4.0
            TransferPath::PciePeerToPeer => 64.0, // PCIe 5.0 x16
            TransferPath::CpuMediated => 50.0,   // DDR5
            TransferPath::UnifiedMemory => 400.0, // Apple unified
            TransferPath::Network => 10.0,       // 100Gbps network
            TransferPath::SameMemory => f64::INFINITY,
        }
    }

    /// Estimates transfer time for given bytes.
    pub fn estimate_transfer_time(&self, bytes: u64) -> std::time::Duration {
        if matches!(self, TransferPath::SameMemory) {
            return std::time::Duration::ZERO;
        }

        let bytes_f64 = bytes as f64;
        let bandwidth = self.bandwidth_gbps() * 1e9; // Convert to bytes/s
        let seconds = bytes_f64 / bandwidth;
        std::time::Duration::from_secs_f64(seconds)
    }
}

/// Unified memory manager.
pub struct MemoryManager {
    /// Allocated tensors.
    tensors: RwLock<HashMap<TensorId, TensorHandle>>,
    /// Memory usage per processor type.
    usage: RwLock<HashMap<ProcessorType, u64>>,
    /// Memory limits per processor type.
    limits: HashMap<ProcessorType, u64>,
}

impl MemoryManager {
    /// Creates a new memory manager.
    pub fn new() -> Self {
        Self {
            tensors: RwLock::new(HashMap::new()),
            usage: RwLock::new(HashMap::new()),
            limits: HashMap::new(),
        }
    }

    /// Sets memory limit for a processor type.
    pub fn set_limit(&mut self, proc_type: ProcessorType, limit_bytes: u64) {
        self.limits.insert(proc_type, limit_bytes);
    }

    /// Allocates a tensor.
    pub fn allocate(&self, shape: Vec<usize>, dtype: DataType) -> Result<TensorHandle, ComputeError> {
        let handle = TensorHandle::new(shape, dtype);
        self.tensors.write().insert(handle.id, handle.clone());
        Ok(handle)
    }

    /// Frees a tensor.
    pub fn free(&self, tensor_id: TensorId) -> Result<(), ComputeError> {
        if let Some(handle) = self.tensors.write().remove(&tensor_id) {
            // Update usage for all locations
            let mut usage = self.usage.write();
            for loc in &handle.locations {
                if let Some(u) = usage.get_mut(loc) {
                    *u = u.saturating_sub(handle.size_bytes);
                }
            }
        }
        Ok(())
    }

    /// Gets a tensor handle.
    pub fn get(&self, tensor_id: TensorId) -> Option<TensorHandle> {
        self.tensors.read().get(&tensor_id).cloned()
    }

    /// Ensures tensor is on specified processor.
    pub fn ensure_on(
        &self,
        tensor_id: TensorId,
        target: ProcessorType,
    ) -> Result<TransferPath, ComputeError> {
        let mut tensors = self.tensors.write();

        if let Some(handle) = tensors.get_mut(&tensor_id) {
            // Check if already on target
            if handle.locations.contains(&target) {
                return Ok(TransferPath::SameMemory);
            }

            // Determine transfer path
            let path = if handle.locations.is_empty() {
                // New tensor, allocate on target
                TransferPath::SameMemory
            } else {
                // Find best transfer path from existing location
                self.find_best_path(&handle.locations[0], &target)
            };

            // Record new location
            handle.locations.push(target.clone());

            // Update usage
            let mut usage = self.usage.write();
            *usage.entry(target).or_default() += handle.size_bytes;

            Ok(path)
        } else {
            Err(ComputeError::Internal("Tensor not found".to_string()))
        }
    }

    /// Finds best transfer path between processors.
    fn find_best_path(&self, from: &ProcessorType, to: &ProcessorType) -> TransferPath {
        // Check for unified memory (Apple Silicon)
        if self.shares_memory(from, to) {
            return TransferPath::UnifiedMemory;
        }

        // Check for NVLink between NVIDIA GPUs
        if matches!(from, ProcessorType::Gpu(crate::processor::GpuVariant::NvidiaCuda { .. }))
            && matches!(to, ProcessorType::Gpu(crate::processor::GpuVariant::NvidiaCuda { .. }))
        {
            return TransferPath::NvLink;
        }

        // Check for PCIe P2P between GPUs
        if from.is_gpu() && to.is_gpu() {
            return TransferPath::PciePeerToPeer;
        }

        // Default to CPU-mediated transfer
        TransferPath::CpuMediated
    }

    /// Checks if two processor types share memory.
    fn shares_memory(&self, a: &ProcessorType, b: &ProcessorType) -> bool {
        use crate::processor::{CpuVariant, GpuVariant, NpuVariant};

        match (a, b) {
            // Apple Silicon unified memory
            (ProcessorType::Cpu(CpuVariant::Arm64 { .. }), ProcessorType::Gpu(GpuVariant::AppleMetal))
            | (ProcessorType::Gpu(GpuVariant::AppleMetal), ProcessorType::Cpu(CpuVariant::Arm64 { .. }))
            | (ProcessorType::Npu(NpuVariant::AppleNeuralEngine { .. }), ProcessorType::Cpu(CpuVariant::Arm64 { .. }))
            | (ProcessorType::Npu(NpuVariant::AppleNeuralEngine { .. }), ProcessorType::Gpu(GpuVariant::AppleMetal)) => true,
            // Same type
            _ if a == b => true,
            _ => false,
        }
    }

    /// Gets current memory usage for a processor type.
    pub fn usage(&self, proc_type: ProcessorType) -> u64 {
        self.usage.read().get(&proc_type).copied().unwrap_or(0)
    }

    /// Gets available memory for a processor type.
    pub fn available(&self, proc_type: ProcessorType) -> u64 {
        let limit = self.limits.get(&proc_type).copied().unwrap_or(u64::MAX);
        let used = self.usage(proc_type);
        limit.saturating_sub(used)
    }

    /// Gets total allocated tensors.
    pub fn tensor_count(&self) -> usize {
        self.tensors.read().len()
    }
}

impl Default for MemoryManager {
    fn default() -> Self {
        Self::new()
    }
}

/// Unified memory abstraction for zero-copy sharing.
pub struct UnifiedMemory {
    /// Base pointer (in unified address space).
    pub base: u64,
    /// Size in bytes.
    pub size: u64,
    /// Accessible from these processor types.
    pub accessible_from: Vec<ProcessorType>,
}

impl UnifiedMemory {
    /// Creates new unified memory region.
    pub fn new(size: u64) -> Self {
        Self {
            base: 0, // Would be actual pointer in real implementation
            size,
            accessible_from: Vec::new(),
        }
    }

    /// Checks if accessible from processor type.
    pub fn is_accessible_from(&self, proc_type: &ProcessorType) -> bool {
        self.accessible_from.contains(proc_type)
    }
}

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

    #[test]
    fn test_tensor_handle() {
        let handle = TensorHandle::new(vec![1024, 1024], DataType::Float32);
        assert_eq!(handle.numel(), 1024 * 1024);
        assert_eq!(handle.size_bytes, 1024 * 1024 * 4);
    }

    #[test]
    fn test_data_type_sizes() {
        assert_eq!(DataType::Float64.size_bytes(), 8);
        assert_eq!(DataType::Float32.size_bytes(), 4);
        assert_eq!(DataType::Float16.size_bytes(), 2);
        assert_eq!(DataType::Int8.size_bytes(), 1);
    }

    #[test]
    fn test_transfer_path_bandwidth() {
        assert!(TransferPath::NvLink.bandwidth_gbps() > TransferPath::PciePeerToPeer.bandwidth_gbps());
        assert!(TransferPath::SameMemory.bandwidth_gbps().is_infinite());
    }

    #[test]
    fn test_memory_manager() {
        let manager = MemoryManager::new();

        let handle = manager.allocate(vec![1024, 1024], DataType::Float32).unwrap();
        assert_eq!(manager.tensor_count(), 1);

        manager.free(handle.id).unwrap();
        assert_eq!(manager.tensor_count(), 0);
    }

    #[test]
    fn test_ensure_on() {
        let manager = MemoryManager::new();

        let handle = manager.allocate(vec![1024], DataType::Float32).unwrap();

        // First ensure should allocate
        let path = manager.ensure_on(
            handle.id,
            ProcessorType::Gpu(crate::processor::GpuVariant::NvidiaCuda {
                compute_capability: (8, 0),
            }),
        ).unwrap();

        assert_eq!(path, TransferPath::SameMemory);

        // Second ensure to same location should be same memory
        let path = manager.ensure_on(
            handle.id,
            ProcessorType::Gpu(crate::processor::GpuVariant::NvidiaCuda {
                compute_capability: (8, 0),
            }),
        ).unwrap();

        assert_eq!(path, TransferPath::SameMemory);
    }
}
