//! Metal Device Management
//!
//! This module provides device discovery, initialization, and management
//! for Apple Metal GPUs.

use crate::error::{MetalError, Result};
use std::sync::Arc;
use tracing::{debug, info};

#[cfg(target_os = "macos")]
use objc2::rc::Retained;
#[cfg(target_os = "macos")]
use objc2::runtime::ProtocolObject;
#[cfg(target_os = "macos")]
use objc2_metal::{MTLCreateSystemDefaultDevice, MTLDevice, MTLCommandQueue, MTLCommandBuffer};

/// Information about a Metal device
#[derive(Debug, Clone)]
pub struct MetalDeviceInfo {
    /// Human-readable device name
    pub name: String,
    /// Registry ID for unique identification
    pub registry_id: u64,
    /// Whether this is a headless (compute-only) device
    pub is_headless: bool,
    /// Whether this is a low-power device
    pub is_low_power: bool,
    /// Whether this device is removable (eGPU)
    pub is_removable: bool,
    /// Recommended maximum working set size in bytes
    pub recommended_max_working_set_size: u64,
    /// Maximum buffer length in bytes
    pub max_buffer_length: usize,
    /// Maximum threads per threadgroup
    pub max_threads_per_threadgroup: usize,
}

/// Metal device wrapper providing safe access to Metal GPU
#[cfg(target_os = "macos")]
pub struct MetalDevice {
    /// The underlying Metal device
    device: Retained<ProtocolObject<dyn MTLDevice>>,
    /// Default command queue for this device
    command_queue: Retained<ProtocolObject<dyn MTLCommandQueue>>,
    /// Cached device info
    info: MetalDeviceInfo,
}

/// Stub implementation for non-macOS platforms
#[cfg(not(target_os = "macos"))]
pub struct MetalDevice {
    _private: (),
}

#[cfg(target_os = "macos")]
use objc2_metal::MTLCopyAllDevices;

#[cfg(target_os = "macos")]
impl MetalDevice {
    /// Check if Metal is available on this system
    pub fn is_available() -> bool {
        MTLCreateSystemDefaultDevice().is_some()
    }

    /// Get the number of available Metal devices
    pub fn device_count() -> usize {
        let devices = MTLCopyAllDevices();
        if devices.is_empty() {
            // Fallback: check if at least default device exists
            usize::from(MTLCreateSystemDefaultDevice().is_some())
        } else {
            devices.len()
        }
    }

    /// Get a Metal device by index
    ///
    /// # Arguments
    /// * `index` - The device index (0-based)
    ///
    /// # Returns
    /// The Metal device at the specified index, or an error if not found.
    pub fn by_index(index: usize) -> Result<Self> {
        debug!("Initializing Metal device at index {}", index);

        let devices = MTLCopyAllDevices();

        if devices.is_empty() {
            // Fallback to default device if index is 0
            if index == 0 {
                return Self::system_default();
            }
            return Err(MetalError::device_init(format!(
                "No Metal devices available (requested index {index})"
            )));
        }

        if index >= devices.len() {
            return Err(MetalError::device_init(format!(
                "Metal device {} not found (only {} devices available)",
                index, devices.len()
            )));
        }

        // Use iterator to get device at index
        let device = devices.iter().nth(index)
            .ok_or_else(|| MetalError::device_init(format!(
                "Failed to get Metal device at index {index}"
            )))?
            .clone();

        let command_queue = device.newCommandQueue()
            .ok_or_else(|| MetalError::device_init("Failed to create command queue"))?;

        let info = Self::extract_device_info(&device);

        info!("Metal device {} initialized: {} (eGPU: {})",
              index, info.name, info.is_removable);
        debug!("Device info: {:?}", info);

        Ok(Self {
            device,
            command_queue,
            info,
        })
    }

    /// Get the system default Metal device
    pub fn system_default() -> Result<Self> {
        debug!("Initializing system default Metal device");

        let device = MTLCreateSystemDefaultDevice()
            .ok_or(MetalError::NoDevice)?;

        let command_queue = device.newCommandQueue()
            .ok_or_else(|| MetalError::device_init("Failed to create command queue"))?;

        let info = Self::extract_device_info(&device);

        info!("Metal device initialized: {}", info.name);
        debug!("Device info: {:?}", info);

        Ok(Self {
            device,
            command_queue,
            info,
        })
    }

    /// Check if this device is an eGPU (external GPU)
    pub fn is_egpu(&self) -> bool {
        self.info.is_removable
    }

    /// Extract device information
    fn extract_device_info(device: &ProtocolObject<dyn MTLDevice>) -> MetalDeviceInfo {
        MetalDeviceInfo {
            name: device.name().to_string(),
            registry_id: device.registryID(),
            is_headless: device.isHeadless(),
            is_low_power: device.isLowPower(),
            is_removable: device.isRemovable(),
            recommended_max_working_set_size: device.recommendedMaxWorkingSetSize(),
            max_buffer_length: device.maxBufferLength(),
            max_threads_per_threadgroup: 1024, // Metal maximum
        }
    }

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

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

    /// Get the underlying MTLDevice
    pub fn mtl_device(&self) -> &ProtocolObject<dyn MTLDevice> {
        &self.device
    }

    /// Get the retained MTLDevice
    pub fn mtl_device_retained(&self) -> &Retained<ProtocolObject<dyn MTLDevice>> {
        &self.device
    }

    /// Get the command queue
    pub fn command_queue(&self) -> &ProtocolObject<dyn MTLCommandQueue> {
        &self.command_queue
    }

    /// Get the retained command queue
    pub fn command_queue_retained(&self) -> &Retained<ProtocolObject<dyn MTLCommandQueue>> {
        &self.command_queue
    }

    /// Create a new command queue
    pub fn new_command_queue(&self) -> Result<Retained<ProtocolObject<dyn MTLCommandQueue>>> {
        self.device.newCommandQueue()
            .ok_or_else(|| MetalError::device_init("Failed to create command queue"))
    }

    /// Get maximum threads per threadgroup
    pub fn max_threads_per_threadgroup(&self) -> usize {
        self.info.max_threads_per_threadgroup
    }

    /// Get maximum buffer length
    pub fn max_buffer_length(&self) -> usize {
        self.info.max_buffer_length
    }

    /// Check if unified memory is supported (always true on Apple Silicon)
    pub fn supports_unified_memory(&self) -> bool {
        // Apple Silicon uses unified memory architecture
        true
    }

    /// Synchronize all pending operations
    pub fn synchronize(&self) -> Result<()> {
        // Create a command buffer and wait for it to complete
        let cmd_buffer = self.command_queue.commandBuffer()
            .ok_or_else(|| MetalError::CommandBuffer("Failed to create sync command buffer".into()))?;
        cmd_buffer.commit();
        cmd_buffer.waitUntilCompleted();
        Ok(())
    }
}

#[cfg(not(target_os = "macos"))]
impl MetalDevice {
    /// Check if Metal is available on this system (always false on non-macOS)
    pub fn is_available() -> bool {
        false
    }

    /// Get the number of available Metal devices (always 0 on non-macOS)
    pub fn device_count() -> usize {
        0
    }

    /// Get a Metal device by index (not available on non-macOS)
    pub fn by_index(_index: usize) -> Result<Self> {
        Err(MetalError::NotAvailable)
    }

    /// Get the system default Metal device (not available on non-macOS)
    pub fn system_default() -> Result<Self> {
        Err(MetalError::NotAvailable)
    }

    /// Check if this device is an eGPU
    pub fn is_egpu(&self) -> bool {
        false
    }

    /// Get the device name
    pub fn name(&self) -> &str {
        "Metal not available"
    }

    /// Get device information
    pub fn info(&self) -> MetalDeviceInfo {
        MetalDeviceInfo {
            name: "Metal not available".to_string(),
            registry_id: 0,
            is_headless: false,
            is_low_power: false,
            is_removable: false,
            recommended_max_working_set_size: 0,
            max_buffer_length: 0,
            max_threads_per_threadgroup: 0,
        }
    }

    /// Get maximum threads per threadgroup
    pub fn max_threads_per_threadgroup(&self) -> usize {
        0
    }

    /// Get maximum buffer length
    pub fn max_buffer_length(&self) -> usize {
        0
    }

    /// Check if unified memory is supported
    pub fn supports_unified_memory(&self) -> bool {
        false
    }

    /// Synchronize all pending operations
    pub fn synchronize(&self) -> Result<()> {
        Err(MetalError::NotAvailable)
    }
}

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

/// Thread-safe wrapper for MetalDevice
pub type SharedMetalDevice = Arc<MetalDevice>;

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

    #[test]
    fn test_metal_availability() {
        let available = MetalDevice::is_available();
        println!("Metal available: {}", available);

        #[cfg(not(target_os = "macos"))]
        assert!(!available);
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_device_creation() {
        if MetalDevice::is_available() {
            let device = MetalDevice::system_default().expect("Failed to create device");
            println!("Device name: {}", device.name());
            assert!(!device.name().is_empty());
        }
    }
}