aetheric-gpu 0.1.0-alpha

//! digital-gpu — Aetheric Silicon's top-level API.
//!
//! Turns this host's RAM into a virtual GPU endpoint with:
//! - 32× effective capacity via binary GEMM compression (1 GB RAM → 32 GB virtual VRAM)
//! - Slow-hill degradation curve: smooth bandwidth drop from DRAM → swap → cloud
//! - Optional LAN cluster pooling: multiple Aetheric nodes share their RAM
//!
//! Usage:
//!   let gpu = digital_gpu::boot(GpuSpec::new()
//!       .effective_gb(32)          // 32 GB effective VRAM
//!       .physical_gb(1)            // backed by 1 GB of this host's RAM
//!       .enable_cluster()          // discover and pool LAN peers
//!   )?;
//!
//!   // Now use gpu.vram() to allocate, gpu.curve() to query bandwidth at any capacity
//!   let buf = gpu.vram().allocate(2 * 1024 * 1024 * 1024)?; // 2 GB virtual
//!   gpu.gemm().dispatch(...); // run compute

use std::sync::Arc;
use std::collections::HashMap;
use std::time::Duration;
use parking_lot::RwLock;
use thiserror::Error;
use serde::{Deserialize, Serialize};
use rayon::prelude::*;

use pim_core::arena::{Arena, ArenaSpec, ArenaTier, AllocationPolicy};
use pim_core::degradation::DegradationCurve;
use vram_backends::{VramManager, VramRegistry, VramTier, VramBackend, AllocationHandle, VramStats};

pub mod cluster;
pub mod distributed;
pub use cluster::{ClusterDiscovery, spawn_discovery, spawn_discovery_on_port, spawn_discovery_local, unix_ms_now, DEFAULT_DISCOVERY_PORT};

/// Errors for digital-gpu
#[derive(Debug, Error)]
pub enum GpuError {
    #[error("Arena error: {0}")]
    Arena(#[from] pim_core::arena::ArenaError),

    #[error("VRAM backend error: {0}")]
    VramBackend(#[from] vram_backends::VramError),

    #[error("Orchestrator error: {0}")]
    Orchestrator(#[from] orchestrator::OrchestratorError),

    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),

    #[error("Compression backend required but not available")]
    CompressionUnavailable,

    #[error("Cluster node {0} unreachable: {1}")]
    ClusterUnreachable(String, String),

    #[error("Invalid spec: {0}")]
    InvalidSpec(String),
}

/// Compression mode for the virtual GPU
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum CompressionMode {
    /// Binary weights + FP32 activations (32× effective, best for quantized LLM inference)
    BinaryGemm,
    /// INT4 quantization (8× effective, general ML workloads)
    Int4,
    /// INT8 quantization (4× effective, training-friendly)
    Int8,
    /// No compression — raw FP32 (1× effective, maximum precision)
    None,
}

impl CompressionMode {
    pub fn expansion_ratio(&self) -> f32 {
        match self {
            Self::BinaryGemm => 32.0,
            Self::Int4 => 8.0,
            Self::Int8 => 4.0,
            Self::None => 1.0,
        }
    }

    pub fn tier(&self) -> VramTier {
        match self {
            Self::BinaryGemm | Self::Int4 | Self::Int8 => VramTier::CompressedRam,
            Self::None => VramTier::UnifiedMemory,
        }
    }
}

/// Cluster node descriptor (discovered or manually configured)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClusterNode {
    pub id: String,
    pub addr: String,          // e.g., "192.168.1.42:50051"
    pub name: String,          // human-readable
    pub physical_ram_gb: u64,  // host's actual RAM
    pub pledged_gb: u64,       // how much this node will contribute
    pub compression: CompressionMode,
    pub last_seen: u64,        // unix timestamp
}

/// Full specification for a virtual GPU
#[derive(Debug, Clone)]
pub struct GpuSpec {
    /// Target effective VRAM in GB (after compression expansion)
    pub effective_gb: u64,

    /// Physical RAM to reserve on this host (GB). Must be <= host RAM.
    /// If None, derives from effective_gb / compression ratio.
    pub physical_gb: Option<u64>,

    /// Compression mode. Determines the expansion ratio.
    pub compression: CompressionMode,

    /// Whether to enable LAN cluster pooling
    pub enable_cluster: bool,

    /// Manually specified cluster peers (for firewalled/private nets)
    pub manual_peers: Vec<ClusterNode>,

    /// Allocation policy for the arena
    pub policy: AllocationPolicy,

    /// Slow-hill curve gentle multiple (default 32.0 — matches BinaryGemm compression
    /// so the entire 32× compression stretch stays at full DRAM bandwidth before
    /// the curve starts to taper).
    pub gentle_multiple: f64,

    /// Slow-hill curve floor multiple (default 4096.0 — gives you effectively
    /// unlimited VVRAM: ~64 TiB on a 16 GiB host before hitting the swap floor).
    pub floor_multiple: f64,
}

impl GpuSpec {
    pub fn new() -> Self {
        Self {
            effective_gb: 32,
            physical_gb: None,
            compression: CompressionMode::BinaryGemm,
            enable_cluster: false,
            manual_peers: Vec::new(),
            policy: AllocationPolicy::Strict,
            gentle_multiple: 32.0,
            floor_multiple: 4096.0,
        }
    }

    pub fn effective_gb(mut self, gb: u64) -> Self {
        self.effective_gb = gb;
        self
    }

    pub fn physical_gb(mut self, gb: u64) -> Self {
        self.physical_gb = Some(gb);
        self
    }

    pub fn compression(mut self, mode: CompressionMode) -> Self {
        self.compression = mode;
        self
    }

    pub fn enable_cluster(mut self) -> Self {
        self.enable_cluster = true;
        self
    }

    pub fn with_peers(mut self, peers: Vec<ClusterNode>) -> Self {
        self.manual_peers = peers;
        self
    }

    pub fn policy(mut self, p: AllocationPolicy) -> Self {
        self.policy = p;
        self
    }

    /// Validate and fill in defaults
    pub fn validate(&mut self, host_ram_gb: u64) -> Result<(), GpuError> {
        if self.effective_gb == 0 {
            return Err(GpuError::InvalidSpec("effective_gb must be > 0".into()));
        }

        // Derive physical GB if not specified
        let physical = self.physical_gb.unwrap_or_else(|| {
            ((self.effective_gb as f32 / self.compression.expansion_ratio()) as u64).max(1)
        });

        // We no longer reject physical_gb > host RAM. VVRAM is virtual —
        // the slow-hill degradation curve handles oversubscription gracefully
        // (full bandwidth through 32× stretch, then gentle taper). The user
        // asks for whatever they need; we report what bandwidth they'll see.
        let _ = host_ram_gb;
        self.physical_gb = Some(physical);
        Ok(())
    }

    /// Physical bytes for arena
    pub fn physical_bytes(&self) -> u64 {
        self.physical_gb.unwrap_or(1) * 1024 * 1024 * 1024
    }

    /// Effective bytes (virtual VRAM after compression)
    pub fn effective_bytes(&self) -> u64 {
        self.effective_gb * 1024 * 1024 * 1024
    }
}

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

/// A live digital GPU handle — the thing the app actually uses
pub struct GpuHandle {
    /// The orchestrator (holds arena + degradation curve + license + security)
    orchestrator: Arc<orchestrator::AethericSilicon>,

    /// VRAM manager (handles compressed allocation + backend selection)
    vram: Arc<VramManager>,

    /// Registered cluster nodes
    cluster: Arc<RwLock<HashMap<String, ClusterNode>>>,

    /// Our spec
    spec: GpuSpec,

    /// Compression-aware arena for binary GEMM workloads
    compressed_arena: Option<Arc<CompressedArena>>,

    /// Discovery thread handle (None if cluster disabled)
    discovery_handle: Option<std::sync::Arc<ClusterDiscovery>>,
    _discovery_thread: Option<std::thread::JoinHandle<()>>,
}

/// A wrapper around Arena that knows about compression expansion
pub struct CompressedArena {
    arena: Arc<dyn ArenaOps>,
    expansion_ratio: f32,
}

/// Trait to allow both Arena and CompressedArena to be used polymorphically
trait ArenaOps: Send + Sync {
    fn allocate(&self, bytes: u64) -> Result<pim_core::arena::Allocation<'_>, pim_core::arena::ArenaError>;
    fn capacity(&self) -> u64;
    fn live_bytes(&self) -> u64;
    fn arena_tier(&self) -> ArenaTier;
    fn bandwidth_gib_s(&self) -> f64;
}

impl ArenaOps for Arc<pim_core::Arena> {
    fn allocate(&self, bytes: u64) -> Result<pim_core::arena::Allocation<'_>, pim_core::arena::ArenaError> {
        let s: &pim_core::Arena = &**self;
        s.allocate(bytes)
    }
    fn capacity(&self) -> u64 { (**self).capacity() }
    fn live_bytes(&self) -> u64 { (**self).live_bytes() }
    fn arena_tier(&self) -> ArenaTier { (**self).arena_tier() }
    fn bandwidth_gib_s(&self) -> f64 { (**self).arena_tier().bandwidth_gib_s() }
}

impl CompressedArena {
    fn new(arena: Arc<dyn ArenaOps>, expansion_ratio: f32) -> Self {
        Self { arena, expansion_ratio }
    }

    /// Allocate `effective_bytes` of virtual VRAM (expanded by compression)
    /// Returns an Allocation whose `.data` is the physical backing
    pub fn allocate_effective(&self, effective_bytes: u64) -> Result<CompressedAllocation<'_>, GpuError> {
        let physical_bytes = (effective_bytes as f32 / self.expansion_ratio) as u64;
        let alloc = self.arena.allocate(physical_bytes)?;
        Ok(CompressedAllocation {
            physical: alloc,
            expansion_ratio: self.expansion_ratio,
        })
    }
}

/// An allocation from a compressed arena
pub struct CompressedAllocation<'a> {
    physical: pim_core::arena::Allocation<'a>,
    expansion_ratio: f32,
}

impl<'a> CompressedAllocation<'a> {
    pub fn physical_slice(&self) -> &[u8] { self.physical.data }
    pub fn physical_slice_mut(&mut self) -> &mut [u8] { self.physical.data }
    pub fn effective_bytes(&self) -> u64 { (self.physical.len() as f32 * self.expansion_ratio) as u64 }
    pub fn physical_bytes(&self) -> u64 { self.physical.len() }
    pub fn offset(&self) -> u64 { self.physical.offset() }
    pub fn bandwidth_gib_s(&self) -> f64 { self.physical.bandwidth().gib_s }
    pub fn deepest_tier(&self) -> ArenaTier { self.physical.bandwidth().deepest }
}

impl GpuHandle {
    /// Boot a digital GPU from the given spec
    pub fn boot(spec: GpuSpec) -> Result<Self, GpuError> {
        let host_ram = detect_host_ram_gb();
        let mut spec = spec;
        spec.validate(host_ram)?;

        let orchestrator = Arc::new(orchestrator::AethericSilicon::new(spec.physical_bytes())?);

        let vram = Arc::new(VramManager::new_blocking()?);

        // Create compressed arena if using compression
        let compressed_arena = if spec.compression != CompressionMode::None {
            let arena = Arc::new(orchestrator.arena.clone()) as Arc<dyn ArenaOps>;
            Some(Arc::new(CompressedArena::new(arena, spec.compression.expansion_ratio())))
        } else {
            None
        };

        let mut discovery_handle: Option<Arc<ClusterDiscovery>> = None;
        let mut discovery_thread: Option<std::thread::JoinHandle<()>> = None;

        if spec.enable_cluster {
            let physical = spec.physical_gb.unwrap_or(host_ram);
            let discovery = Arc::new(ClusterDiscovery::new(
                physical,
                physical,
                spec.compression,
            ));
            let d_for_thread = discovery.clone();
            let h = cluster::spawn_discovery(d_for_thread);
            discovery_thread = Some(h);
            discovery_handle = Some(discovery);
        }

        let gpu = Self {
            orchestrator,
            vram,
            cluster: Arc::new(RwLock::new(HashMap::new())),
            spec: spec.clone(),
            compressed_arena,
            discovery_handle,
            _discovery_thread: discovery_thread,
        };

        // Copy any peers already in our data path into the runtime table
        if let Some(d) = gpu.discovery_handle.as_ref() {
            for peer in d.peer_list() {
                gpu.cluster.write().insert(peer.id.clone(), peer);
            }
        }

        // Add manual peers
        for peer in &spec.manual_peers {
            gpu.cluster.write().insert(peer.id.clone(), peer.clone());
        }

        log::info!(
            target: "aetheric",
            "Digital GPU booted: effective={} GB, physical={} GB, compression={:?}, expansion={:.1}x",
            spec.effective_gb,
            spec.physical_gb.unwrap_or(1),
            spec.compression,
            spec.compression.expansion_ratio()
        );

        Ok(gpu)
    }

    /// Get the orchestrator (for status, curve, tier info)
    pub fn orchestrator(&self) -> &orchestrator::AethericSilicon {
        &self.orchestrator
    }

    /// Get the VRAM manager (for backend-aware allocation)
    pub fn vram(&self) -> &VramManager {
        &self.vram
    }

    /// Get the compressed arena (for binary GEMM / quantized workloads)
    pub fn compressed_arena(&self) -> Option<&CompressedArena> {
        self.compressed_arena.as_deref()
    }

    /// Get the degradation curve (bandwidth at any virtual capacity)
    pub fn curve(&self) -> &DegradationCurve {
        self.orchestrator.curve()
    }

    /// Current cluster nodes
    pub fn cluster_nodes(&self) -> Vec<ClusterNode> {
        self.cluster.read().values().cloned().collect()
    }

    /// Discover Aetheric Silicon peers on LAN (mDNS + broadcast)
    fn discover_cluster(&self) -> Result<(), GpuError> {
        // TODO: implement mDNS / UDP broadcast discovery
        // For now, this is a stub that would:
        // 1. Listen on UDP 50050 for "Aetheric Hello" packets
        // 2. Broadcast our presence
        // 3. Maintain peer list with heartbeats
        Ok(())
    }

    /// Attempt to add a cluster node
    pub fn add_node(&self, node: ClusterNode) {
        self.cluster.write().insert(node.id.clone(), node);
    }

    /// Effective capacity (virtual VRAM)
    pub fn effective_capacity_gb(&self) -> u64 {
        self.spec.effective_gb
    }

    /// Physical capacity (host RAM committed)
    pub fn physical_capacity_gb(&self) -> u64 {
        self.spec.physical_gb.unwrap_or(1)
    }

    /// Effective bandwidth at a given virtual capacity
    pub fn bandwidth_at_effective(&self, virtual_bytes: u64) -> f64 {
        // The degradation curve works on physical bytes.
        // We need to map virtual → physical using our expansion ratio.
        let expansion = self.spec.compression.expansion_ratio();
        let physical_bytes = (virtual_bytes as f32 / expansion) as u64;
        self.curve().bandwidth_at(physical_bytes)
    }

    /// Report the full slow-hill curve for the frontend to render
    pub fn curve_sample(&self, max_virtual_bytes: u64, samples: usize) -> Vec<CurvePoint> {
        let curve = self.curve();
        let expansion = self.spec.compression.expansion_ratio();
        let max_physical = (max_virtual_bytes as f32 / expansion) as u64;
        let max = max_physical.max(1);
        let step = max / samples.max(1) as u64;
        let mut out = Vec::with_capacity(samples);
        let mut i = 0u64;
        while i <= max {
            let physical_bw = curve.bandwidth_at(i);
            let virtual_bytes = (i as f32 * expansion) as u64;
            out.push(CurvePoint {
                bytes: virtual_bytes,
                bandwidth_gib_s: physical_bw,
                tier: curve.tier_name_at(i).to_string(),
            });
            i = i.saturating_add(step.max(1));
        }
        out
    }
}

/// One sample of the bandwidth-vs-capacity curve
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CurvePoint {
    pub bytes: u64,            // virtual bytes (after compression)
    pub bandwidth_gib_s: f64,  // actual GiB/s at that virtual capacity
    pub tier: String,          // tier label
}

/// Detect host RAM in GB
fn detect_host_ram_gb() -> u64 {
    #[cfg(target_os = "macos")]
    {
        if let Ok(out) = std::process::Command::new("sysctl")
            .args(["-n", "hw.memsize"])
            .output()
        {
            if let Ok(s) = std::str::from_utf8(&out.stdout) {
                if let Ok(n) = s.trim().parse::<u64>() {
                    return n / (1024 * 1024 * 1024);
                }
            }
        }
    }
    #[cfg(target_os = "linux")]
    {
        if let Ok(s) = std::fs::read_to_string("/proc/meminfo") {
            for line in s.lines() {
                if let Some(rest) = line.strip_prefix("MemTotal:") {
                    if let Some(kib) = rest.trim().split_whitespace().next() {
                        if let Ok(n) = kib.parse::<u64>() {
                            return (n * 1024) / (1024 * 1024 * 1024);
                        }
                    }
                }
            }
        }
    }
    16 // conservative fallback
}

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

    #[test]
    fn spec_defaults() {
        let spec = GpuSpec::new();
        assert_eq!(spec.effective_gb, 32);
        assert_eq!(spec.compression, CompressionMode::BinaryGemm);
        assert!(!spec.enable_cluster);
    }

    #[test]
    fn spec_builder_pattern() {
        let spec = GpuSpec::new()
            .effective_gb(64)
            .compression(CompressionMode::Int4)
            .enable_cluster();
        assert_eq!(spec.effective_gb, 64);
        assert_eq!(spec.compression, CompressionMode::Int4);
        assert!(spec.enable_cluster);
    }

    #[test]
    fn spec_validates_physical_ram() {
        // Oversubscription is now allowed — the slow-hill curve handles it
        let mut spec = GpuSpec::new()
            .effective_gb(128)
            .physical_gb(4 * 1024);  // 4 TB >> host RAM — allowed via oversubscription
        assert!(spec.validate(32).is_ok());
    }

    #[test]
    fn spec_accepts_packed_under_host_ram() {
        // 1 GB physical backing 64 GB effective via 32× ratio
        let mut spec = GpuSpec::new()
            .effective_gb(64)
            .physical_gb(2);  // 2 GB physical fits in 32 GB host
        assert!(spec.validate(32).is_ok());
        assert_eq!(spec.physical_bytes(), 2 * 1024 * 1024 * 1024);
        assert_eq!(spec.effective_bytes(), 64 * 1024 * 1024 * 1024);
    }

    #[test]
    fn spec_rejects_zero() {
        let mut spec = GpuSpec::new().effective_gb(0);
        let err = spec.validate(32).unwrap_err();
        assert!(matches!(err, GpuError::InvalidSpec(_)));
    }

    #[test]
    fn compression_ratios() {
        assert_eq!(CompressionMode::BinaryGemm.expansion_ratio(), 32.0);
        assert_eq!(CompressionMode::Int4.expansion_ratio(), 8.0);
        assert_eq!(CompressionMode::Int8.expansion_ratio(), 4.0);
        assert_eq!(CompressionMode::None.expansion_ratio(), 1.0);
    }
}