trustformers-core 0.1.1

// Copyright (c) 2025-2026 COOLJAPAN OU (Team KitaSan)
// SPDX-License-Identifier: Apache-2.0

//! Hardware resource allocation and management components
//!
//! This module provides resource allocation strategies, load balancing, memory management,
//! and memory pressure monitoring for hardware devices.

use super::config::{AllocationStrategy, LoadBalancingStrategy, MemoryUsageStats};
use super::traits::{MemoryType, OperationParameter};
use super::HardwareResult;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::SystemTime;

/// Resource allocator for managing hardware resource assignments
#[derive(Debug, Clone)]
pub struct ResourceAllocator {
    /// Current allocation strategy
    pub strategy: AllocationStrategy,
    /// Active resource reservations
    pub reservations: HashMap<String, ResourceReservation>,
    /// Historical allocation records
    pub history: Vec<AllocationRecord>,
    /// Resource limits per device
    pub limits: HashMap<String, ResourceLimits>,
}

/// Resource reservation details
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ResourceReservation {
    /// Target device ID
    pub device_id: String,
    /// Reserved resource amounts by type
    pub resources: HashMap<String, f64>,
    /// Reservation creation timestamp
    pub timestamp: SystemTime,
    /// Optional expiration time
    pub expiration: Option<SystemTime>,
    /// Unique reservation identifier
    pub id: String,
}

/// Allocation record for auditing and analytics
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct AllocationRecord {
    /// Allocated device ID
    pub device_id: String,
    /// Allocation timestamp
    pub timestamp: SystemTime,
    /// Duration of allocation
    pub duration: std::time::Duration,
    /// Resources allocated
    pub resources: HashMap<String, f64>,
    /// Operation parameters
    pub operation_params: Vec<OperationParameter>,
    /// Success indicator
    pub success: bool,
    /// Performance metrics
    pub performance_metrics: HashMap<String, f64>,
}

/// Resource limits configuration per device
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ResourceLimits {
    /// Maximum CPU utilization (0.0 - 1.0)
    pub max_cpu: f64,
    /// Maximum memory utilization (0.0 - 1.0)
    pub max_memory: f64,
    /// Maximum GPU utilization (0.0 - 1.0)
    pub max_gpu: f64,
    /// Maximum power consumption (watts)
    pub max_power: f64,
    /// Maximum bandwidth (bytes/sec)
    pub max_bandwidth: f64,
    /// Custom resource limits
    pub custom_limits: HashMap<String, f64>,
}

/// Load balancer for distributing work across devices
#[derive(Debug, Clone)]
pub struct LoadBalancer {
    /// Active load balancing strategy
    pub strategy: LoadBalancingStrategy,
    /// Device weights for weighted algorithms
    pub weights: HashMap<String, f64>,
    /// Connection counts per device
    pub connections: HashMap<String, u64>,
    /// Load history for trend analysis
    pub load_history: HashMap<String, Vec<(SystemTime, f64)>>,
    /// Adaptive thresholds for dynamic balancing
    pub adaptive_thresholds: HashMap<String, f64>,
}

/// Memory manager for device memory pools and allocation
#[derive(Debug, Clone)]
pub struct MemoryManager {
    /// Memory pools per device
    pub pools: HashMap<String, MemoryPool>,
    /// Memory usage tracking
    pub usage_tracking: HashMap<String, MemoryUsageStats>,
    /// Garbage collection schedules
    pub gc_schedule: HashMap<String, SystemTime>,
    /// Memory pressure monitor
    pub pressure_monitor: MemoryPressureMonitor,
}

/// Memory pool representation
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MemoryPool {
    /// Pool identifier
    pub id: String,
    /// Associated device ID
    pub device_id: String,
    /// Total pool size in bytes
    pub total_size: usize,
    /// Currently used size in bytes
    pub used_size: usize,
    /// Available size in bytes
    pub available_size: usize,
    /// Allocated memory blocks
    pub allocated_blocks: Vec<MemoryBlock>,
    /// Free memory blocks
    pub free_blocks: Vec<MemoryBlock>,
    /// Memory fragmentation ratio (0.0 - 1.0)
    pub fragmentation_ratio: f64,
}

/// Memory block allocation unit
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MemoryBlock {
    /// Block identifier
    pub id: String,
    /// Memory offset
    pub offset: usize,
    /// Block size in bytes
    pub size: usize,
    /// Memory type
    pub memory_type: MemoryType,
    /// Allocation timestamp
    pub allocated_at: SystemTime,
    /// Optional tags for categorization
    pub tags: Vec<String>,
}

/// Memory pressure monitor for tracking memory pressure levels
#[derive(Debug, Clone)]
pub struct MemoryPressureMonitor {
    /// Current pressure levels per device
    pub pressure_levels: HashMap<String, MemoryPressureLevel>,
    /// Historical pressure data
    pub pressure_history: HashMap<String, Vec<(SystemTime, f64)>>,
    /// Pressure thresholds configuration
    pub thresholds: HashMap<String, MemoryPressureThresholds>,
}

/// Memory pressure level indicators
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum MemoryPressureLevel {
    /// Low pressure - optimal conditions
    Low,
    /// Medium pressure - some concern
    Medium,
    /// High pressure - action recommended
    High,
    /// Critical pressure - immediate action required
    Critical,
}

/// Memory pressure threshold configuration
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MemoryPressureThresholds {
    /// Low pressure threshold (0.0 - 1.0)
    pub low: f64,
    /// Medium pressure threshold (0.0 - 1.0)
    pub medium: f64,
    /// High pressure threshold (0.0 - 1.0)
    pub high: f64,
    /// Critical pressure threshold (0.0 - 1.0)
    pub critical: f64,
}

impl ResourceAllocator {
    /// Create a new resource allocator
    pub fn new(strategy: AllocationStrategy) -> Self {
        Self {
            strategy,
            reservations: HashMap::new(),
            history: Vec::new(),
            limits: HashMap::new(),
        }
    }

    /// Allocate resources on the best available device
    pub fn allocate(&mut self, requirements: &HashMap<String, f64>) -> HardwareResult<String> {
        // Placeholder implementation - in practice, this would implement
        // sophisticated allocation logic based on the strategy
        let device_id = match self.strategy {
            AllocationStrategy::FirstAvailable => "device_0".to_string(),
            AllocationStrategy::BestFit => self.find_best_fit_device(requirements)?,
            AllocationStrategy::RoundRobin => self.next_round_robin_device(),
            AllocationStrategy::LoadAware => self.find_least_loaded_device()?,
            AllocationStrategy::PerformanceOptimized => self.find_highest_performance_device()?,
            AllocationStrategy::PowerEfficient => self.find_most_power_efficient_device()?,
        };

        // Record the allocation
        let record = AllocationRecord {
            device_id: device_id.clone(),
            timestamp: SystemTime::now(),
            duration: std::time::Duration::from_secs(0), // Will be updated on completion
            resources: requirements.clone(),
            operation_params: vec![],
            success: true,
            performance_metrics: HashMap::new(),
        };
        self.history.push(record);

        Ok(device_id)
    }

    /// Find device with best fit for requirements
    fn find_best_fit_device(&self, _requirements: &HashMap<String, f64>) -> HardwareResult<String> {
        // Placeholder - would implement actual best-fit algorithm
        Ok("best_fit_device".to_string())
    }

    /// Get next device in round-robin order
    fn next_round_robin_device(&self) -> String {
        // Placeholder - would maintain round-robin state
        "round_robin_device".to_string()
    }

    /// Find device with lowest current load
    fn find_least_loaded_device(&self) -> HardwareResult<String> {
        // Placeholder - would check actual device loads
        Ok("least_loaded_device".to_string())
    }

    /// Find device with highest performance rating
    fn find_highest_performance_device(&self) -> HardwareResult<String> {
        // Placeholder - would check performance metrics
        Ok("high_perf_device".to_string())
    }

    /// Find most power-efficient device
    fn find_most_power_efficient_device(&self) -> HardwareResult<String> {
        // Placeholder - would check power efficiency ratings
        Ok("power_efficient_device".to_string())
    }

    /// Set resource limits for a device
    pub fn set_limits(&mut self, device_id: &str, limits: ResourceLimits) {
        self.limits.insert(device_id.to_string(), limits);
    }

    /// Get allocation history
    pub fn get_history(&self) -> &[AllocationRecord] {
        &self.history
    }
}

impl LoadBalancer {
    /// Create a new load balancer
    pub fn new(strategy: LoadBalancingStrategy) -> Self {
        Self {
            strategy,
            weights: HashMap::new(),
            connections: HashMap::new(),
            load_history: HashMap::new(),
            adaptive_thresholds: HashMap::new(),
        }
    }

    /// Select next device based on load balancing strategy
    pub fn select_device(&mut self, available_devices: &[String]) -> HardwareResult<String> {
        if available_devices.is_empty() {
            return Err(super::TrustformersError::hardware_error(
                "No devices available",
                "allocate",
            ));
        }

        let selected = match self.strategy {
            LoadBalancingStrategy::RoundRobin => self.round_robin_select(available_devices),
            LoadBalancingStrategy::LeastConnections => {
                self.least_connections_select(available_devices)
            },
            LoadBalancingStrategy::LeastUtilization => {
                self.least_utilization_select(available_devices)
            },
            LoadBalancingStrategy::WeightedRoundRobin => {
                self.weighted_round_robin_select(available_devices)
            },
            LoadBalancingStrategy::PerformanceBased => {
                self.performance_based_select(available_devices)
            },
            LoadBalancingStrategy::Adaptive => self.adaptive_select(available_devices),
        };

        // Update connection count
        *self.connections.entry(selected.clone()).or_insert(0) += 1;

        Ok(selected)
    }

    fn round_robin_select(&self, devices: &[String]) -> String {
        // Placeholder - would maintain state for round-robin
        devices[0].clone()
    }

    fn least_connections_select(&self, devices: &[String]) -> String {
        devices
            .iter()
            .min_by_key(|device| self.connections.get(*device).unwrap_or(&0))
            .expect("devices slice must be non-empty (checked by caller)")
            .clone()
    }

    fn least_utilization_select(&self, devices: &[String]) -> String {
        // Placeholder - would check actual utilization
        devices[0].clone()
    }

    fn weighted_round_robin_select(&self, devices: &[String]) -> String {
        // Placeholder - would implement weighted selection
        devices[0].clone()
    }

    fn performance_based_select(&self, devices: &[String]) -> String {
        // Placeholder - would select based on performance metrics
        devices[0].clone()
    }

    fn adaptive_select(&self, devices: &[String]) -> String {
        // Placeholder - would implement adaptive selection
        devices[0].clone()
    }

    /// Update device weight
    pub fn set_weight(&mut self, device_id: &str, weight: f64) {
        self.weights.insert(device_id.to_string(), weight);
    }
}

impl MemoryManager {
    /// Create a new memory manager
    pub fn new() -> Self {
        Self {
            pools: HashMap::new(),
            usage_tracking: HashMap::new(),
            gc_schedule: HashMap::new(),
            pressure_monitor: MemoryPressureMonitor::new(),
        }
    }

    /// Allocate memory block
    pub fn allocate_memory(
        &mut self,
        device_id: &str,
        size: usize,
        memory_type: MemoryType,
    ) -> HardwareResult<MemoryBlock> {
        let pool = self
            .pools
            .entry(device_id.to_string())
            .or_insert_with(|| MemoryPool::new(device_id));

        pool.allocate(size, memory_type)
    }

    /// Deallocate memory block
    pub fn deallocate_memory(&mut self, device_id: &str, block_id: &str) -> HardwareResult<()> {
        if let Some(pool) = self.pools.get_mut(device_id) {
            pool.deallocate(block_id)
        } else {
            Err(super::TrustformersError::hardware_error(
                "Device not found",
                "deallocate",
            ))
        }
    }

    /// Trigger garbage collection for a device
    pub fn trigger_gc(&mut self, device_id: &str) -> HardwareResult<()> {
        if let Some(pool) = self.pools.get_mut(device_id) {
            pool.garbage_collect()?;
            self.gc_schedule.insert(device_id.to_string(), SystemTime::now());
        }
        Ok(())
    }

    /// Get memory usage statistics
    pub fn get_usage_stats(&self, device_id: &str) -> Option<&MemoryUsageStats> {
        self.usage_tracking.get(device_id)
    }
}

impl MemoryPool {
    /// Create a new memory pool
    pub fn new(device_id: &str) -> Self {
        Self {
            id: format!("pool_{}", device_id),
            device_id: device_id.to_string(),
            total_size: 1024 * 1024 * 1024, // 1GB default
            used_size: 0,
            available_size: 1024 * 1024 * 1024,
            allocated_blocks: Vec::new(),
            free_blocks: Vec::new(),
            fragmentation_ratio: 0.0,
        }
    }

    /// Allocate a memory block
    pub fn allocate(
        &mut self,
        size: usize,
        memory_type: MemoryType,
    ) -> HardwareResult<MemoryBlock> {
        if self.available_size < size {
            return Err(super::TrustformersError::hardware_error(
                "Insufficient memory",
                "allocate",
            ));
        }

        let block = MemoryBlock {
            id: format!(
                "block_{}_{}",
                self.allocated_blocks.len(),
                chrono::Utc::now().timestamp_nanos_opt().unwrap_or(0)
            ),
            offset: self.used_size,
            size,
            memory_type,
            allocated_at: SystemTime::now(),
            tags: vec![],
        };

        self.allocated_blocks.push(block.clone());
        self.used_size += size;
        self.available_size -= size;

        Ok(block)
    }

    /// Deallocate a memory block
    pub fn deallocate(&mut self, block_id: &str) -> HardwareResult<()> {
        if let Some(pos) = self.allocated_blocks.iter().position(|b| b.id == block_id) {
            let block = self.allocated_blocks.remove(pos);
            self.used_size -= block.size;
            self.available_size += block.size;
            self.free_blocks.push(block);
            Ok(())
        } else {
            Err(super::TrustformersError::hardware_error(
                "Block not found",
                "deallocate",
            ))
        }
    }

    /// Perform garbage collection
    pub fn garbage_collect(&mut self) -> HardwareResult<()> {
        // Coalesce free blocks and update fragmentation ratio
        self.free_blocks.sort_by_key(|b| b.offset);
        // Implementation would coalesce adjacent free blocks
        self.fragmentation_ratio = self.calculate_fragmentation();
        Ok(())
    }

    fn calculate_fragmentation(&self) -> f64 {
        if self.free_blocks.is_empty() {
            return 0.0;
        }
        // Simplified fragmentation calculation
        self.free_blocks.len() as f64 / (self.total_size / 1024) as f64
    }
}

impl MemoryPressureMonitor {
    /// Create a new memory pressure monitor
    pub fn new() -> Self {
        Self {
            pressure_levels: HashMap::new(),
            pressure_history: HashMap::new(),
            thresholds: HashMap::new(),
        }
    }

    /// Update pressure level for a device
    pub fn update_pressure(&mut self, device_id: &str, utilization: f64) {
        let default_thresholds = MemoryPressureThresholds::default();
        let thresholds = self.thresholds.get(device_id).unwrap_or(&default_thresholds);

        let level = if utilization < thresholds.low {
            MemoryPressureLevel::Low
        } else if utilization < thresholds.medium {
            MemoryPressureLevel::Medium
        } else if utilization < thresholds.high {
            MemoryPressureLevel::High
        } else {
            MemoryPressureLevel::Critical
        };

        self.pressure_levels.insert(device_id.to_string(), level);

        // Record in history
        let entry = self.pressure_history.entry(device_id.to_string()).or_default();
        entry.push((SystemTime::now(), utilization));

        // Keep only last 1000 entries
        if entry.len() > 1000 {
            entry.drain(..500);
        }
    }

    /// Get current pressure level
    pub fn get_pressure_level(&self, device_id: &str) -> Option<MemoryPressureLevel> {
        self.pressure_levels.get(device_id).copied()
    }

    /// Set pressure thresholds for a device
    pub fn set_thresholds(&mut self, device_id: &str, thresholds: MemoryPressureThresholds) {
        self.thresholds.insert(device_id.to_string(), thresholds);
    }
}

impl Default for ResourceLimits {
    fn default() -> Self {
        Self {
            max_cpu: 0.8,
            max_memory: 0.9,
            max_gpu: 0.95,
            max_power: 300.0,
            max_bandwidth: 10_000_000_000.0, // 10 GB/s
            custom_limits: HashMap::new(),
        }
    }
}

impl Default for MemoryPressureThresholds {
    fn default() -> Self {
        Self {
            low: 0.5,
            medium: 0.7,
            high: 0.85,
            critical: 0.95,
        }
    }
}

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

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