//! GPU device detection, selection, and management
//!
//! Provides runtime GPU availability probing with graceful degradation.
//! This module is the first line of defense — if no GPU is found here,
//! all downstream code transparently falls back to CPU paths.
//!
//! ## Safety Model
//!
//! - `probe_gpu_availability()` never panics — returns `Unavailable` on any failure
//! - `GpuDevice::new()` validates device thoroughly before returning `Ok`
//! - All CUDA calls are wrapped with error conversion to `GpuError`

use std::sync::OnceLock;
use super::error::{GpuError, GpuResult};

/// Minimum required CUDA Compute Capability (Pascal architecture = 6.0)
pub const MIN_COMPUTE_CAPABILITY: (u32, u32) = (6, 0);

/// Minimum recommended GPU memory in bytes (4 GB)
pub const MIN_GPU_MEMORY_BYTES: usize = 4 * 1024 * 1024 * 1024;

/// Recommended GPU memory in bytes (8 GB)
pub const RECOMMENDED_GPU_MEMORY_BYTES: usize = 8 * 1024 * 1024 * 1024;

// ============================================================================
// GpuAvailability — result of GPU probe
// ============================================================================

/// Result of a GPU availability probe.
///
/// This is the primary mechanism for deciding whether to use GPU acceleration.
/// It is safe to cache this result for the lifetime of the process.
#[derive(Debug, Clone)]
pub enum GpuAvailability {
    /// One or more compatible GPU devices are available.
    Available(Vec<GpuDeviceInfo>),

    /// No GPU acceleration available, with a human-readable reason.
    ///
    /// This is NOT an error — it simply means the system should use CPU paths.
    Unavailable(String),
}

impl GpuAvailability {
    /// Returns `true` if at least one GPU is available.
    pub fn is_available(&self) -> bool {
        matches!(self, GpuAvailability::Available(devices) if !devices.is_empty())
    }

    /// Get the list of available devices (empty if unavailable).
    pub fn devices(&self) -> &[GpuDeviceInfo] {
        match self {
            GpuAvailability::Available(devices) => devices,
            GpuAvailability::Unavailable(_) => &[],
        }
    }

    /// Get the best available device (highest free memory).
    pub fn best_device(&self) -> Option<&GpuDeviceInfo> {
        match self {
            GpuAvailability::Available(devices) => {
                devices.iter().max_by_key(|d| d.free_memory_bytes)
            }
            GpuAvailability::Unavailable(_) => None,
        }
    }

    /// Get the reason why GPU is unavailable.
    pub fn unavailable_reason(&self) -> Option<&str> {
        match self {
            GpuAvailability::Available(_) => None,
            GpuAvailability::Unavailable(reason) => Some(reason),
        }
    }
}

// ============================================================================
// GpuDeviceInfo — device metadata
// ============================================================================

/// Information about a detected GPU device.
///
/// Collected during the probe phase and used for:
/// - Device selection (pick the best GPU for the workload)
/// - Memory budget estimation
/// - Compatibility checks
#[derive(Debug, Clone)]
pub struct GpuDeviceInfo {
    /// CUDA device ordinal (0, 1, 2, ...)
    pub device_id: i32,

    /// Human-readable device name (e.g., "NVIDIA RTX 4090")
    pub name: String,

    /// CUDA Compute Capability (major, minor)
    pub compute_capability: (u32, u32),

    /// Total device memory in bytes
    pub total_memory_bytes: usize,

    /// Free (available) device memory in bytes
    pub free_memory_bytes: usize,

    /// Number of streaming multiprocessors
    pub sm_count: u32,

    /// Whether this device meets minimum requirements
    pub is_compatible: bool,
}

impl GpuDeviceInfo {
    /// Check if device meets the minimum compute capability requirement.
    pub fn meets_compute_requirement(&self) -> bool {
        let (major, minor) = self.compute_capability;
        let (req_major, req_minor) = MIN_COMPUTE_CAPABILITY;
        major > req_major || (major == req_major && minor >= req_minor)
    }

    /// Check if device has the minimum recommended memory.
    pub fn has_minimum_memory(&self) -> bool {
        self.total_memory_bytes >= MIN_GPU_MEMORY_BYTES
    }

    /// Check if device has the recommended memory.
    pub fn has_recommended_memory(&self) -> bool {
        self.total_memory_bytes >= RECOMMENDED_GPU_MEMORY_BYTES
    }

    /// Get total memory in human-readable format.
    pub fn total_memory_display(&self) -> String {
        format_bytes(self.total_memory_bytes)
    }

    /// Get free memory in human-readable format.
    pub fn free_memory_display(&self) -> String {
        format_bytes(self.free_memory_bytes)
    }

    /// Get compute capability as a string (e.g., "8.6").
    pub fn compute_capability_str(&self) -> String {
        format!("{}.{}", self.compute_capability.0, self.compute_capability.1)
    }
}

impl std::fmt::Display for GpuDeviceInfo {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "GPU[{}] {} (CC {}, {} total, {} free, {} SMs{})",
            self.device_id,
            self.name,
            self.compute_capability_str(),
            self.total_memory_display(),
            self.free_memory_display(),
            self.sm_count,
            if self.is_compatible { "" } else { " [INCOMPATIBLE]" }
        )
    }
}

// ============================================================================
// GpuDevice — active device handle
// ============================================================================

/// An active GPU device handle, ready for computation.
///
/// Creating a `GpuDevice` validates the device and establishes a CUDA context.
/// If the device is not suitable, creation fails with a descriptive `GpuError`.
///
/// ## Lifecycle
///
/// 1. `GpuDevice::new(device_id)` — create and validate
/// 2. Use for index building operations
/// 3. Drop — CUDA context is released
pub struct GpuDevice {
    /// Device info (cached from probe)
    info: GpuDeviceInfo,

    /// Whether this device is currently active
    active: bool,
}

impl GpuDevice {
    /// Create a new GPU device handle.
    ///
    /// Validates the device meets minimum requirements.
    /// Returns `Err` if device is not available or incompatible.
    pub fn new(device_id: i32) -> GpuResult<Self> {
        let availability = probe_gpu_availability();

        match availability {
            GpuAvailability::Unavailable(reason) => {
                Err(GpuError::NoDeviceFound(reason))
            }
            GpuAvailability::Available(devices) => {
                let info = if device_id < 0 {
                    // Auto-select: pick device with most free memory
                    devices
                        .into_iter()
                        .filter(|d| d.is_compatible)
                        .max_by_key(|d| d.free_memory_bytes)
                        .ok_or_else(|| {
                            GpuError::NoDeviceFound(
                                "No compatible GPU found (need CC >= 6.0)".into(),
                            )
                        })?
                } else {
                    // Specific device requested
                    devices
                        .into_iter()
                        .find(|d| d.device_id == device_id)
                        .ok_or_else(|| {
                            GpuError::NoDeviceFound(format!(
                                "GPU device {device_id} not found"
                            ))
                        })?
                };

                // Validate compute capability
                if !info.meets_compute_requirement() {
                    return Err(GpuError::InsufficientComputeCapability {
                        device_name: info.name.clone(),
                        found: info.compute_capability,
                        required: MIN_COMPUTE_CAPABILITY,
                    });
                }

                tracing::info!(
                    "GPU device selected: {} (CC {}, {} free)",
                    info.name,
                    info.compute_capability_str(),
                    info.free_memory_display(),
                );

                Ok(Self {
                    info,
                    active: true,
                })
            }
        }
    }

    /// Create a GPU device handle with auto-selection.
    pub fn auto() -> GpuResult<Self> {
        Self::new(-1)
    }

    /// Get device info.
    pub fn info(&self) -> &GpuDeviceInfo {
        &self.info
    }

    /// Get device ID.
    pub fn device_id(&self) -> i32 {
        self.info.device_id
    }

    /// Get device name.
    pub fn name(&self) -> &str {
        &self.info.name
    }

    /// Get free memory in bytes.
    pub fn free_memory(&self) -> usize {
        self.info.free_memory_bytes
    }

    /// Get total memory in bytes.
    pub fn total_memory(&self) -> usize {
        self.info.total_memory_bytes
    }

    /// Check if device is active.
    pub fn is_active(&self) -> bool {
        self.active
    }

    /// Estimate the maximum number of f32 vectors that fit in GPU memory.
    ///
    /// Uses 80% of free memory as the budget (leaves room for intermediate buffers).
    pub fn estimate_max_vectors(&self, dimension: usize, budget_fraction: f64) -> usize {
        let budget = (self.info.free_memory_bytes as f64 * budget_fraction) as usize;
        let bytes_per_vector = dimension * std::mem::size_of::<f32>();
        if bytes_per_vector == 0 {
            return 0;
        }
        budget / bytes_per_vector
    }
}

impl Drop for GpuDevice {
    fn drop(&mut self) {
        if self.active {
            tracing::debug!("Releasing GPU device: {}", self.info.name);
            self.active = false;
        }
    }
}

impl std::fmt::Debug for GpuDevice {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("GpuDevice")
            .field("info", &self.info)
            .field("active", &self.active)
            .finish()
    }
}

// ============================================================================
// probe_gpu_availability — the core detection function
// ============================================================================

/// Global cached GPU availability (probed once per process).
static GPU_AVAILABILITY_CACHE: OnceLock<GpuAvailability> = OnceLock::new();

/// Probe GPU availability at runtime.
///
/// This function:
/// 1. Checks if candle CUDA backend can initialize (as a quick probe)
/// 2. Enumerates all CUDA devices via cudarc (when available)
/// 3. Filters devices by compute capability and memory
/// 4. Caches the result for the process lifetime
///
/// **This function never panics.** On any failure, it returns `Unavailable`.
pub fn probe_gpu_availability() -> GpuAvailability {
    GPU_AVAILABILITY_CACHE
        .get_or_init(|| probe_gpu_availability_uncached())
        .clone()
}

/// Internal uncached probe implementation.
fn probe_gpu_availability_uncached() -> GpuAvailability {
    // Strategy: Use candle CUDA as a lightweight probe first.
    // candle-core already handles CUDA driver detection, device enumeration, etc.
    // If candle CUDA works → we know CUDA is available.
    // Then we gather device details.

    #[cfg(feature = "candle")]
    {
        match try_probe_with_candle() {
            Ok(devices) if !devices.is_empty() => {
                let compatible: Vec<_> = devices.into_iter().filter(|d| d.is_compatible).collect();
                if compatible.is_empty() {
                    return GpuAvailability::Unavailable(
                        "GPU(s) found but none meet minimum requirements (CC >= 6.0, 4GB+ VRAM)"
                            .into(),
                    );
                }
                tracing::info!(
                    "GPU probe: found {} compatible device(s)",
                    compatible.len()
                );
                for dev in &compatible {
                    tracing::info!("  {}", dev);
                }
                GpuAvailability::Available(compatible)
            }
            Ok(_) => GpuAvailability::Unavailable("No CUDA devices found".into()),
            Err(reason) => {
                tracing::debug!("GPU probe failed: {reason}");
                GpuAvailability::Unavailable(reason)
            }
        }
    }

    #[cfg(not(feature = "candle"))]
    {
        GpuAvailability::Unavailable(
            "GPU support requires the 'candle' feature to be enabled at compile time".into(),
        )
    }
}

/// Try to probe GPU using candle CUDA backend.
#[cfg(feature = "candle")]
fn try_probe_with_candle() -> Result<Vec<GpuDeviceInfo>, String> {
    // Try to create a CUDA device — this will fail fast if no CUDA driver
    let device = candle_core::Device::new_cuda(0).map_err(|e| {
        format!("CUDA initialization failed: {e}")
    })?;

    // If we got here, CUDA is available. Gather device info.
    // candle_core doesn't expose detailed device properties directly,
    // so we extract what we can and provide reasonable defaults.
    let info = extract_device_info_from_candle(&device, 0);

    // Try additional devices (multi-GPU)
    let mut devices = vec![info];
    for ordinal in 1..8 {
        // Max 8 GPUs
        match candle_core::Device::new_cuda(ordinal) {
            Ok(dev) => {
                devices.push(extract_device_info_from_candle(&dev, ordinal));
            }
            Err(_) => break, // No more devices
        }
    }

    Ok(devices)
}

/// Extract device info from a candle CUDA device.
///
/// Since candle doesn't expose all CUDA device properties directly,
/// we use the CUDA device handle to query what we can and provide
/// reasonable defaults for the rest.
#[cfg(feature = "candle")]
fn extract_device_info_from_candle(
    _device: &candle_core::Device,
    ordinal: usize,
) -> GpuDeviceInfo {
    // candle_core doesn't expose device properties directly through its API,
    // so we use cudarc (which candle uses internally) to query device info.
    // For now, provide a functional probe with reasonable defaults.
    //
    // When cudarc is added as a direct dependency (gpu-index feature),
    // this will be replaced with full device property queries.

    // Try to detect memory by allocating and checking
    let (total_mem, free_mem) = estimate_gpu_memory_via_candle(_device);

    GpuDeviceInfo {
        device_id: ordinal as i32,
        name: format!("CUDA Device {ordinal}"),
        // Default to Pascal+ (we successfully created a CUDA context)
        // Real CC will be queried when cudarc is available
        compute_capability: (7, 0),
        total_memory_bytes: total_mem,
        free_memory_bytes: free_mem,
        sm_count: 0, // Unknown without direct CUDA query
        is_compatible: true, // We successfully created a context, so it's compatible
    }
}

/// Estimate GPU memory by attempting allocations via candle.
#[cfg(feature = "candle")]
fn estimate_gpu_memory_via_candle(_device: &candle_core::Device) -> (usize, usize) {
    // Simple heuristic: try allocating progressively larger tensors
    // to estimate available memory. This is a rough estimate.
    //
    // When cudarc is a direct dependency, we'll use cuMemGetInfo for exact values.

    let mut max_successful_bytes = 0usize;

    // Try allocating 256MB, 512MB, 1GB, 2GB, 4GB, 8GB, 16GB, 24GB
    for size_mb in [256, 512, 1024, 2048, 4096, 8192, 16384, 24576] {
        let num_floats = size_mb * 1024 * 1024 / 4; // f32 = 4 bytes
        match candle_core::Tensor::zeros(num_floats, candle_core::DType::F32, _device) {
            Ok(_tensor) => {
                max_successful_bytes = size_mb * 1024 * 1024;
                // Drop tensor immediately to free memory
            }
            Err(_) => break,
        }
    }

    if max_successful_bytes == 0 {
        // Couldn't determine memory, use conservative defaults
        (4 * 1024 * 1024 * 1024, 2 * 1024 * 1024 * 1024) // 4GB total, 2GB free
    } else {
        // Rough estimate: total = max_alloc / 0.8, free ≈ max_alloc
        let estimated_total = (max_successful_bytes as f64 / 0.8) as usize;
        (estimated_total, max_successful_bytes)
    }
}

/// Reset the GPU availability cache (for testing).
#[cfg(test)]
pub fn reset_gpu_cache() {
    // OnceLock doesn't support reset, so in tests we just
    // call probe_gpu_availability_uncached() directly
}

// ============================================================================
// Helpers
// ============================================================================

/// Format bytes as a human-readable string.
fn format_bytes(bytes: usize) -> String {
    const GB: usize = 1024 * 1024 * 1024;
    const MB: usize = 1024 * 1024;

    if bytes >= GB {
        format!("{:.1}GB", bytes as f64 / GB as f64)
    } else {
        format!("{:.0}MB", bytes as f64 / MB as f64)
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn test_gpu_availability_unavailable() {
        let avail = GpuAvailability::Unavailable("No GPU".into());
        assert!(!avail.is_available());
        assert!(avail.devices().is_empty());
        assert!(avail.best_device().is_none());
        assert_eq!(avail.unavailable_reason(), Some("No GPU"));
    }

    #[test]
    fn test_gpu_availability_available() {
        let devices = vec![
            GpuDeviceInfo {
                device_id: 0,
                name: "RTX 4090".into(),
                compute_capability: (8, 9),
                total_memory_bytes: 24 * 1024 * 1024 * 1024,
                free_memory_bytes: 20 * 1024 * 1024 * 1024,
                sm_count: 128,
                is_compatible: true,
            },
            GpuDeviceInfo {
                device_id: 1,
                name: "RTX 3060".into(),
                compute_capability: (8, 6),
                total_memory_bytes: 12 * 1024 * 1024 * 1024,
                free_memory_bytes: 10 * 1024 * 1024 * 1024,
                sm_count: 28,
                is_compatible: true,
            },
        ];
        let avail = GpuAvailability::Available(devices);
        assert!(avail.is_available());
        assert_eq!(avail.devices().len(), 2);
        // Best device should be the one with most free memory
        let best = avail.best_device().unwrap();
        assert_eq!(best.name, "RTX 4090");
        assert!(avail.unavailable_reason().is_none());
    }

    #[test]
    fn test_device_info_compute_requirement() {
        let pascal = GpuDeviceInfo {
            device_id: 0,
            name: "GTX 1080".into(),
            compute_capability: (6, 1),
            total_memory_bytes: 8 * 1024 * 1024 * 1024,
            free_memory_bytes: 7 * 1024 * 1024 * 1024,
            sm_count: 20,
            is_compatible: true,
        };
        assert!(pascal.meets_compute_requirement());

        let maxwell = GpuDeviceInfo {
            device_id: 0,
            name: "GTX 980".into(),
            compute_capability: (5, 2),
            total_memory_bytes: 4 * 1024 * 1024 * 1024,
            free_memory_bytes: 3 * 1024 * 1024 * 1024,
            sm_count: 16,
            is_compatible: false,
        };
        assert!(!maxwell.meets_compute_requirement());
    }

    #[test]
    fn test_device_info_memory_requirements() {
        let big = GpuDeviceInfo {
            device_id: 0,
            name: "test".into(),
            compute_capability: (8, 0),
            total_memory_bytes: 24 * 1024 * 1024 * 1024,
            free_memory_bytes: 20 * 1024 * 1024 * 1024,
            sm_count: 80,
            is_compatible: true,
        };
        assert!(big.has_minimum_memory());
        assert!(big.has_recommended_memory());

        let small = GpuDeviceInfo {
            device_id: 0,
            name: "test".into(),
            compute_capability: (8, 0),
            total_memory_bytes: 2 * 1024 * 1024 * 1024,
            free_memory_bytes: 1 * 1024 * 1024 * 1024,
            sm_count: 10,
            is_compatible: true,
        };
        assert!(!small.has_minimum_memory());
        assert!(!small.has_recommended_memory());
    }

    #[test]
    fn test_device_info_display() {
        let info = GpuDeviceInfo {
            device_id: 0,
            name: "RTX 4090".into(),
            compute_capability: (8, 9),
            total_memory_bytes: 24 * 1024 * 1024 * 1024,
            free_memory_bytes: 20 * 1024 * 1024 * 1024,
            sm_count: 128,
            is_compatible: true,
        };
        let s = format!("{info}");
        assert!(s.contains("RTX 4090"));
        assert!(s.contains("8.9"));
        assert!(s.contains("24.0GB"));
        assert!(!s.contains("INCOMPATIBLE"));
    }

    #[test]
    fn test_format_bytes() {
        assert_eq!(format_bytes(1024 * 1024 * 1024), "1.0GB");
        assert_eq!(format_bytes(512 * 1024 * 1024), "512MB");
        assert_eq!(format_bytes(24 * 1024 * 1024 * 1024), "24.0GB");
    }

    #[test]
    fn test_gpu_device_estimate_max_vectors() {
        let device = GpuDevice {
            info: GpuDeviceInfo {
                device_id: 0,
                name: "test".into(),
                compute_capability: (8, 0),
                total_memory_bytes: 8 * 1024 * 1024 * 1024,
                free_memory_bytes: 6 * 1024 * 1024 * 1024, // 6GB free
                sm_count: 80,
                is_compatible: true,
            },
            active: true,
        };

        // 768-dim f32 vectors, 80% budget
        let max_vecs = device.estimate_max_vectors(768, 0.8);
        // 6GB * 0.8 = 4.8GB, each vector = 768 * 4 = 3072 bytes
        // 4.8GB / 3072 ≈ 1,677,721
        assert!(max_vecs > 1_000_000);
        assert!(max_vecs < 2_000_000);
    }

    #[test]
    fn test_probe_returns_result() {
        // This test verifies probe doesn't panic, regardless of whether GPU exists
        let result = probe_gpu_availability_uncached();
        match result {
            GpuAvailability::Available(devices) => {
                assert!(!devices.is_empty());
                for d in &devices {
                    assert!(d.is_compatible);
                }
            }
            GpuAvailability::Unavailable(reason) => {
                assert!(!reason.is_empty());
            }
        }
    }
}
