ruvllm 2.1.0

//! Apple Neural Engine (ANE) Backend via Core ML
//!
//! This module provides LLM inference acceleration using Apple's Neural Engine,
//! available on M1/M2/M3/M4 chips. The ANE provides:
//!
//! - **38 TOPS** on M4 Pro (dedicated ML accelerator)
//! - **3-4x better power efficiency** vs GPU
//! - **Parallel execution** alongside GPU for hybrid pipelines
//!
//! ## When to Use ANE
//!
//! | Scenario | ANE Benefit | Recommendation |
//! |----------|-------------|----------------|
//! | Small models (<1B) | +20-40% faster | **Use ANE** |
//! | Large models (7B+) | Minimal | Use GPU |
//! | Batch inference | +50% throughput | **Use ANE** |
//! | Battery life | 3-4x better | **Use ANE** |
//! | Low latency | Higher latency | Use GPU |
//!
//! ## Architecture
//!
//! ```text
//! +-------------------+     +-------------------+
//! | GGUF Model        |---->| CoreML Converter  |
//! | (quantized)       |     | - Weights         |
//! +-------------------+     | - Topology        |
//!                           +--------+----------+
//!                                    |
//!                                    v
//!                           +--------+----------+
//!                           | Core ML Model     |
//!                           | (.mlmodel)        |
//!                           +--------+----------+
//!                                    |
//!                     +--------------+--------------+
//!                     |                             |
//!                     v                             v
//!            +--------+----------+       +----------+--------+
//!            | ANE (MLP/FFN)     |       | GPU (Attention)   |
//!            | - MatMul          |       | - Flash Attention |
//!            | - Activations     |       | - KV Cache        |
//!            +-------------------+       +-------------------+
//! ```
//!
//! ## Usage
//!
//! ```rust,ignore
//! use ruvllm::backends::CoreMLBackend;
//!
//! // Create backend with ANE preference
//! let backend = CoreMLBackend::new()?
//!     .with_compute_units(ComputeUnits::CpuAndNeuralEngine)?;
//!
//! // Load model (converts GGUF to Core ML on first load)
//! backend.load_model("path/to/model.gguf", ModelConfig::default())?;
//!
//! // Generate (uses ANE for MLP, GPU for attention)
//! let output = backend.generate("Hello", GenerateParams::default())?;
//! ```
//!
//! ## Feature Flags
//!
//! - `coreml`: Enable Core ML backend (this module)
//! - `hybrid-ane`: Enable hybrid GPU+ANE pipeline

use super::{
    DType, DeviceType, GenerateParams, GeneratedToken, LlmBackend, ModelArchitecture, ModelConfig,
    ModelInfo, Quantization, SpecialTokens, StreamEvent, TokenStream, Tokenizer,
};
use crate::error::{Result, RuvLLMError};

use std::path::{Path, PathBuf};
use std::sync::mpsc;
use std::time::Instant;

/// Compute units for Core ML inference
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ComputeUnits {
    /// CPU only (fallback)
    CpuOnly,
    /// CPU and GPU
    CpuAndGpu,
    /// CPU and Neural Engine (ANE)
    CpuAndNeuralEngine,
    /// All available compute units (CPU, GPU, ANE)
    #[default]
    All,
}

impl ComputeUnits {
    /// Get description of compute units
    pub fn description(&self) -> &'static str {
        match self {
            Self::CpuOnly => "CPU only",
            Self::CpuAndGpu => "CPU + GPU",
            Self::CpuAndNeuralEngine => "CPU + Neural Engine (ANE)",
            Self::All => "CPU + GPU + Neural Engine",
        }
    }

    /// Check if ANE is included
    pub fn uses_ane(&self) -> bool {
        matches!(self, Self::CpuAndNeuralEngine | Self::All)
    }

    /// Check if GPU is included
    pub fn uses_gpu(&self) -> bool {
        matches!(self, Self::CpuAndGpu | Self::All)
    }
}

/// ANE capability information
#[derive(Debug, Clone)]
pub struct AneCapabilities {
    /// Whether ANE is available on this device
    pub available: bool,
    /// ANE compute power in TOPS (Trillion Operations Per Second)
    pub tops: f32,
    /// Maximum supported model size in MB
    pub max_model_size_mb: usize,
    /// Supported operations
    pub supported_ops: Vec<String>,
}

impl Default for AneCapabilities {
    fn default() -> Self {
        Self::detect()
    }
}

impl AneCapabilities {
    /// Detect ANE capabilities on the current device
    pub fn detect() -> Self {
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        {
            // M4 Pro ANE specs
            Self {
                available: true,
                tops: 38.0,              // M4 Pro: 38 TOPS
                max_model_size_mb: 2048, // ~2GB models work well on ANE
                supported_ops: vec![
                    "MatMul".to_string(),
                    "Conv2D".to_string(),
                    "GELU".to_string(),
                    "SiLU".to_string(),
                    "LayerNorm".to_string(),
                    "Softmax".to_string(),
                    "Add".to_string(),
                    "Mul".to_string(),
                ],
            }
        }

        #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
        {
            Self {
                available: false,
                tops: 0.0,
                max_model_size_mb: 0,
                supported_ops: vec![],
            }
        }
    }

    /// Check if a model of given size is suitable for ANE
    pub fn is_model_suitable(&self, model_size_mb: usize) -> bool {
        self.available && model_size_mb <= self.max_model_size_mb
    }
}

// =============================================================================
// Core ML Model Handle (macOS aarch64 only with coreml feature)
// =============================================================================

/// Core ML model wrapper that holds the actual model reference
#[cfg(all(target_os = "macos", target_arch = "aarch64", feature = "coreml"))]
pub mod coreml_native {
    use super::*;
    use objc2::rc::Retained;
    use objc2::runtime::AnyObject;
    use objc2::{msg_send_id, ClassType};
    use objc2_core_ml::{
        MLComputeUnits as MLComputeUnitsObjc, MLDictionaryFeatureProvider, MLFeatureProvider,
        MLFeatureValue, MLModel, MLModelConfiguration, MLMultiArray, MLMultiArrayDataType,
        MLPredictionOptions,
    };
    use objc2_foundation::{NSArray, NSDictionary, NSNumber, NSString, NSURL};

    /// Wrapper around Core ML MLModel
    pub struct CoreMLModelHandle {
        /// The loaded Core ML model
        model: Retained<MLModel>,
        /// Path to the model file
        model_path: PathBuf,
        /// Model description from Core ML
        description: String,
        /// Input feature names
        input_names: Vec<String>,
        /// Output feature names
        output_names: Vec<String>,
        /// Vocab size detected from model (if available)
        vocab_size: Option<usize>,
        /// Hidden size detected from model (if available)
        hidden_size: Option<usize>,
    }

    // Safety: MLModel is thread-safe for predictions after loading
    // The Objective-C runtime handles thread synchronization internally
    unsafe impl Send for CoreMLModelHandle {}
    unsafe impl Sync for CoreMLModelHandle {}

    impl CoreMLModelHandle {
        /// Load a Core ML model from a compiled .mlmodelc directory or .mlmodel file
        pub fn load(path: &Path, compute_units: ComputeUnits) -> Result<Self> {
            // Validate the path exists
            if !path.exists() {
                return Err(RuvLLMError::NotFound(format!(
                    "Core ML model not found: {}",
                    path.display()
                )));
            }

            // Create NSURL from path
            let url = NSURL::from_file_path(path).ok_or_else(|| {
                RuvLLMError::CoreML(format!("Invalid model path: {}", path.display()))
            })?;

            // Create configuration with specified compute units
            let config = unsafe { MLModelConfiguration::new() };

            // Set compute units based on preference
            let ml_compute_units = match compute_units {
                ComputeUnits::CpuOnly => MLComputeUnitsObjc::CPUOnly,
                ComputeUnits::CpuAndGpu => MLComputeUnitsObjc::CPUAndGPU,
                ComputeUnits::CpuAndNeuralEngine => MLComputeUnitsObjc::CPUAndNeuralEngine,
                ComputeUnits::All => MLComputeUnitsObjc::All,
            };

            unsafe {
                config.setComputeUnits(ml_compute_units);
            }

            // Load the model synchronously
            let model =
                unsafe { MLModel::modelWithContentsOfURL_configuration_error(&url, &config) }
                    .map_err(|e| {
                        RuvLLMError::CoreML(format!(
                            "Failed to load Core ML model from {}: {}",
                            path.display(),
                            e.localizedDescription()
                        ))
                    })?;

            // Extract model info
            let (description, input_names, output_names, vocab_size, hidden_size) =
                Self::extract_model_info(&model);

            Ok(Self {
                model,
                model_path: path.to_path_buf(),
                description,
                input_names,
                output_names,
                vocab_size,
                hidden_size,
            })
        }

        /// Extract model description and feature names from MLModel
        fn extract_model_info(
            model: &MLModel,
        ) -> (
            String,
            Vec<String>,
            Vec<String>,
            Option<usize>,
            Option<usize>,
        ) {
            unsafe {
                let desc = model.modelDescription();
                let input_desc = desc.inputDescriptionsByName();
                let output_desc = desc.outputDescriptionsByName();

                let input_count = input_desc.count();
                let output_count = output_desc.count();

                // Extract input names
                let input_names: Vec<String> = input_desc
                    .allKeys()
                    .iter()
                    .map(|key| key.to_string())
                    .collect();

                // Extract output names
                let output_names: Vec<String> = output_desc
                    .allKeys()
                    .iter()
                    .map(|key| key.to_string())
                    .collect();

                let description = format!("Inputs: {}, Outputs: {}", input_count, output_count);

                // Try to detect vocab_size and hidden_size from output feature descriptions
                // These are typically encoded in the shape of output arrays
                let vocab_size = None; // Would need to inspect output shapes
                let hidden_size = None;

                (
                    description,
                    input_names,
                    output_names,
                    vocab_size,
                    hidden_size,
                )
            }
        }

        /// Create an MLMultiArray with the given shape for token IDs (Int32)
        pub fn create_input_array(&self, token_ids: &[i32]) -> Result<Retained<MLMultiArray>> {
            let seq_len = token_ids.len();

            unsafe {
                // Create shape: [1, seq_len] for batch_size=1
                let shape_vec: Vec<Retained<NSNumber>> = vec![
                    NSNumber::new_isize(1),
                    NSNumber::new_isize(seq_len as isize),
                ];
                let shape = NSArray::from_retained_slice(&shape_vec);

                // Create MLMultiArray with Int32 data type using msg_send_id for allocation
                use objc2::rc::Allocated;
                let alloc: Allocated<MLMultiArray> = msg_send_id![MLMultiArray::class(), alloc];
                let array = MLMultiArray::initWithShape_dataType_error(
                    alloc,
                    &shape,
                    MLMultiArrayDataType::Int32,
                )
                .map_err(|e| {
                    RuvLLMError::CoreML(format!(
                        "Failed to create input MLMultiArray: {}",
                        e.localizedDescription()
                    ))
                })?;

                // Copy token IDs into the array
                let ptr = array.dataPointer().as_ptr() as *mut i32;
                for (i, &token_id) in token_ids.iter().enumerate() {
                    *ptr.add(i) = token_id;
                }

                Ok(array)
            }
        }

        /// Create an MLMultiArray with the given shape for float outputs
        pub fn create_float_array(&self, shape: &[usize]) -> Result<Retained<MLMultiArray>> {
            unsafe {
                let shape_vec: Vec<Retained<NSNumber>> = shape
                    .iter()
                    .map(|&d| NSNumber::new_isize(d as isize))
                    .collect();
                let ns_shape = NSArray::from_retained_slice(&shape_vec);

                use objc2::rc::Allocated;
                let alloc: Allocated<MLMultiArray> = msg_send_id![MLMultiArray::class(), alloc];
                let array = MLMultiArray::initWithShape_dataType_error(
                    alloc,
                    &ns_shape,
                    MLMultiArrayDataType::Float32,
                )
                .map_err(|e| {
                    RuvLLMError::CoreML(format!(
                        "Failed to create float MLMultiArray: {}",
                        e.localizedDescription()
                    ))
                })?;

                Ok(array)
            }
        }

        /// Run inference on the model with token IDs input
        ///
        /// # Arguments
        /// * `input_name` - The name of the input feature (e.g., "input_ids")
        /// * `token_ids` - The token IDs to feed to the model
        ///
        /// # Returns
        /// The raw logits output as a flattened f32 vector
        pub fn predict(&self, input_name: &str, token_ids: &[i32]) -> Result<Vec<f32>> {
            // Create input array
            let input_array = self.create_input_array(token_ids)?;

            unsafe {
                // Create NSString for input name
                let input_key = NSString::from_str(input_name);

                // Create feature value from the multi-array
                let feature_value = MLFeatureValue::featureValueWithMultiArray(&input_array);

                // Create dictionary with input feature
                // Use objc2's msg_send for dictionary creation to properly handle types
                use objc2::runtime::ProtocolObject;

                // Create NSDictionary directly with dictionaryWithObject_forKey
                // Use AnyObject as value type since initWithDictionary_error expects NSDictionary<NSString, AnyObject>
                let dict: Retained<NSDictionary<NSString, AnyObject>> = msg_send_id![NSDictionary::<NSString, AnyObject>::class(), dictionaryWithObject: &*feature_value, forKey: &*input_key];

                // Create feature provider using msg_send_id for allocation
                use objc2::rc::Allocated;
                let alloc: Allocated<MLDictionaryFeatureProvider> =
                    msg_send_id![MLDictionaryFeatureProvider::class(), alloc];
                let provider = MLDictionaryFeatureProvider::initWithDictionary_error(alloc, &*dict)
                    .map_err(|e| {
                        RuvLLMError::CoreML(format!(
                            "Failed to create feature provider: {}",
                            e.localizedDescription()
                        ))
                    })?;

                // Create prediction options
                let options = MLPredictionOptions::new();

                // Run prediction - cast provider to protocol object
                let provider_ref = ProtocolObject::from_ref(&*provider);
                let output = self
                    .model
                    .predictionFromFeatures_options_error(provider_ref, &options)
                    .map_err(|e| {
                        RuvLLMError::CoreML(format!(
                            "Prediction failed: {}",
                            e.localizedDescription()
                        ))
                    })?;

                // Get the output feature value (assume first output is logits)
                let output_name = self
                    .output_names
                    .first()
                    .ok_or_else(|| RuvLLMError::CoreML("No output features found".to_string()))?;

                let output_key = NSString::from_str(output_name);
                // Use MLFeatureProvider protocol method
                let output_value = MLFeatureProvider::featureValueForName(&*output, &output_key)
                    .ok_or_else(|| {
                        RuvLLMError::CoreML(format!("Output feature '{}' not found", output_name))
                    })?;

                // Get the multi-array from the output
                let output_array = output_value.multiArrayValue().ok_or_else(|| {
                    RuvLLMError::CoreML("Output is not a multi-array".to_string())
                })?;

                // Extract data from the output array
                let count = output_array.count() as usize;
                let ptr = output_array.dataPointer().as_ptr() as *const f32;
                let logits: Vec<f32> = (0..count).map(|i| *ptr.add(i)).collect();

                Ok(logits)
            }
        }

        /// Extract embeddings from the model (hidden states)
        ///
        /// # Arguments
        /// * `input_name` - The name of the input feature
        /// * `token_ids` - The token IDs to feed to the model
        /// * `embedding_output_name` - The name of the embedding output feature (optional)
        ///
        /// # Returns
        /// The embedding vector (last token's hidden state, or pooled output)
        pub fn get_embeddings(
            &self,
            input_name: &str,
            token_ids: &[i32],
            embedding_output_name: Option<&str>,
        ) -> Result<Vec<f32>> {
            let input_array = self.create_input_array(token_ids)?;

            unsafe {
                use objc2::rc::Allocated;
                use objc2::runtime::ProtocolObject;

                let input_key = NSString::from_str(input_name);
                let feature_value = MLFeatureValue::featureValueWithMultiArray(&input_array);

                // Create NSDictionary directly with dictionaryWithObject_forKey
                // Use AnyObject as value type since initWithDictionary_error expects NSDictionary<NSString, AnyObject>
                let dict: Retained<NSDictionary<NSString, AnyObject>> = msg_send_id![NSDictionary::<NSString, AnyObject>::class(), dictionaryWithObject: &*feature_value, forKey: &*input_key];

                // Create feature provider using msg_send_id for allocation
                let alloc: Allocated<MLDictionaryFeatureProvider> =
                    msg_send_id![MLDictionaryFeatureProvider::class(), alloc];
                let provider = MLDictionaryFeatureProvider::initWithDictionary_error(alloc, &*dict)
                    .map_err(|e| {
                        RuvLLMError::CoreML(format!(
                            "Failed to create feature provider: {}",
                            e.localizedDescription()
                        ))
                    })?;

                let options = MLPredictionOptions::new();
                // Run prediction - cast provider to protocol object
                let provider_ref = ProtocolObject::from_ref(&*provider);
                let output = self
                    .model
                    .predictionFromFeatures_options_error(provider_ref, &options)
                    .map_err(|e| {
                        RuvLLMError::CoreML(format!(
                            "Prediction failed: {}",
                            e.localizedDescription()
                        ))
                    })?;

                // Try to find embeddings output - use specified name or fall back to common patterns
                let embedding_name = embedding_output_name.map(String::from).or_else(|| {
                    // Common names for embedding outputs
                    for name in &self.output_names {
                        let lower = name.to_lowercase();
                        if lower.contains("embed")
                            || lower.contains("hidden")
                            || lower.contains("pooled")
                            || lower.contains("last_hidden")
                        {
                            return Some(name.clone());
                        }
                    }
                    // Fall back to first output if no match
                    self.output_names.first().cloned()
                });

                let output_name = embedding_name.ok_or_else(|| {
                    RuvLLMError::CoreML("No embedding output found in model".to_string())
                })?;

                let output_key = NSString::from_str(&output_name);
                // Use MLFeatureProvider protocol method
                let output_value = MLFeatureProvider::featureValueForName(&*output, &output_key)
                    .ok_or_else(|| {
                        RuvLLMError::CoreML(format!("Embedding output '{}' not found", output_name))
                    })?;

                let output_array = output_value.multiArrayValue().ok_or_else(|| {
                    RuvLLMError::CoreML("Embedding output is not a multi-array".to_string())
                })?;

                // For embeddings, we typically want the last token's hidden state
                // Shape is usually [batch, seq_len, hidden_dim] - we take [0, -1, :]
                let count = output_array.count() as usize;
                let ptr = output_array.dataPointer().as_ptr() as *const f32;

                // Get shape to extract last token embedding
                let shape_count = output_array.shape().count() as usize;
                if shape_count >= 3 {
                    // Shape: [batch, seq_len, hidden_dim]
                    let shape_arr = output_array.shape();
                    let seq_len = shape_arr.objectAtIndex(1).intValue() as usize;
                    let hidden_dim = shape_arr.objectAtIndex(2).intValue() as usize;

                    // Extract last token's embedding
                    let last_token_start = (seq_len - 1) * hidden_dim;
                    let embeddings: Vec<f32> = (0..hidden_dim)
                        .map(|i| *ptr.add(last_token_start + i))
                        .collect();

                    Ok(embeddings)
                } else {
                    // Flat or pooled output - return all
                    let embeddings: Vec<f32> = (0..count).map(|i| *ptr.add(i)).collect();
                    Ok(embeddings)
                }
            }
        }

        /// Get the underlying MLModel reference
        pub fn model(&self) -> &MLModel {
            &self.model
        }

        /// Get the model path
        pub fn path(&self) -> &Path {
            &self.model_path
        }

        /// Get model description string
        pub fn description(&self) -> &str {
            &self.description
        }

        /// Get input feature names
        pub fn input_names(&self) -> &[String] {
            &self.input_names
        }

        /// Get output feature names
        pub fn output_names(&self) -> &[String] {
            &self.output_names
        }

        /// Get the number of input features
        pub fn num_inputs(&self) -> usize {
            self.input_names.len()
        }

        /// Get the number of output features
        pub fn num_outputs(&self) -> usize {
            self.output_names.len()
        }
    }

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

#[cfg(all(target_os = "macos", target_arch = "aarch64", feature = "coreml"))]
pub use coreml_native::CoreMLModelHandle;

// =============================================================================
// Core ML Stream Iterator (for generate_stream)
// =============================================================================

/// Iterator for streaming Core ML token generation
#[cfg(all(
    target_os = "macos",
    target_arch = "aarch64",
    feature = "coreml",
    feature = "candle"
))]
pub struct CoreMLStreamIterator<'a> {
    model_handle: &'a CoreMLModelHandle,
    tokenizer: &'a crate::tokenizer::RuvTokenizer,
    input_ids: Vec<i32>,
    max_tokens: usize,
    temperature: f32,
    top_p: f32,
    input_feature_name: String,
    eos_token_id: u32,
    vocab_size: usize,
    generated_count: usize,
    finished: bool,
}

#[cfg(all(
    target_os = "macos",
    target_arch = "aarch64",
    feature = "coreml",
    feature = "candle"
))]
impl<'a> CoreMLStreamIterator<'a> {
    /// Create a new streaming iterator
    pub fn new(
        model_handle: &'a CoreMLModelHandle,
        tokenizer: &'a crate::tokenizer::RuvTokenizer,
        input_ids: Vec<i32>,
        max_tokens: usize,
        temperature: f32,
        top_p: f32,
        input_feature_name: String,
        eos_token_id: u32,
        vocab_size: usize,
    ) -> Self {
        Self {
            model_handle,
            tokenizer,
            input_ids,
            max_tokens,
            temperature,
            top_p,
            input_feature_name,
            eos_token_id,
            vocab_size,
            generated_count: 0,
            finished: false,
        }
    }

    /// Sample a token from logits
    fn sample_token(&self, logits: &[f32]) -> Result<u32> {
        use rand::Rng;

        if logits.is_empty() {
            return Err(RuvLLMError::Generation("Empty logits".to_string()));
        }

        // Apply temperature
        let scaled_logits: Vec<f32> = if self.temperature > 0.0 && self.temperature != 1.0 {
            logits.iter().map(|&x| x / self.temperature).collect()
        } else {
            logits.to_vec()
        };

        // Softmax
        let max_logit = scaled_logits
            .iter()
            .cloned()
            .fold(f32::NEG_INFINITY, f32::max);
        let exp_logits: Vec<f32> = scaled_logits
            .iter()
            .map(|&x| (x - max_logit).exp())
            .collect();
        let sum_exp: f32 = exp_logits.iter().sum();
        let probs: Vec<f32> = exp_logits.iter().map(|&x| x / sum_exp).collect();

        // Top-p sampling
        if self.top_p < 1.0 {
            let mut indexed_probs: Vec<(usize, f32)> = probs.iter().copied().enumerate().collect();
            indexed_probs
                .sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

            let mut cumsum = 0.0;
            let mut cutoff_idx = indexed_probs.len();
            for (i, (_, p)) in indexed_probs.iter().enumerate() {
                cumsum += p;
                if cumsum >= self.top_p {
                    cutoff_idx = i + 1;
                    break;
                }
            }

            let filtered: Vec<(usize, f32)> = indexed_probs[..cutoff_idx].to_vec();
            let filter_sum: f32 = filtered.iter().map(|(_, p)| p).sum();
            let normalized: Vec<(usize, f32)> = filtered
                .into_iter()
                .map(|(i, p)| (i, p / filter_sum))
                .collect();

            let mut rng = rand::thread_rng();
            let r: f32 = rng.gen();
            let mut cumsum = 0.0;
            for (idx, p) in &normalized {
                cumsum += p;
                if r < cumsum {
                    return Ok(*idx as u32);
                }
            }
            return Ok(normalized.last().map(|(i, _)| *i as u32).unwrap_or(0));
        }

        // Regular sampling
        let mut rng = rand::thread_rng();
        let r: f32 = rng.gen();
        let mut cumsum = 0.0;
        for (idx, &p) in probs.iter().enumerate() {
            cumsum += p;
            if r < cumsum {
                return Ok(idx as u32);
            }
        }

        Ok(probs
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .map(|(i, _)| i as u32)
            .unwrap_or(0))
    }
}

#[cfg(all(
    target_os = "macos",
    target_arch = "aarch64",
    feature = "coreml",
    feature = "candle"
))]
impl<'a> Iterator for CoreMLStreamIterator<'a> {
    type Item = Result<GeneratedToken>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.finished || self.generated_count >= self.max_tokens {
            return None;
        }

        // Run inference
        let logits = match self
            .model_handle
            .predict(&self.input_feature_name, &self.input_ids)
        {
            Ok(l) => l,
            Err(e) => {
                self.finished = true;
                return Some(Err(e));
            }
        };

        // Get last token logits
        let last_token_logits = if logits.len() >= self.vocab_size {
            &logits[logits.len() - self.vocab_size..]
        } else {
            &logits
        };

        // Sample next token
        let next_token = match self.sample_token(last_token_logits) {
            Ok(t) => t,
            Err(e) => {
                self.finished = true;
                return Some(Err(e));
            }
        };

        // Check for EOS
        if next_token == self.eos_token_id {
            self.finished = true;
            return None;
        }

        // Decode the token
        let text = self.tokenizer.decode(&[next_token]).unwrap_or_default();

        // Add to sequence
        self.input_ids.push(next_token as i32);
        self.generated_count += 1;

        Some(Ok(GeneratedToken {
            id: next_token,
            text,
            logprob: None,
            is_special: false,
        }))
    }
}

// Safety: The iterator holds references to CoreMLModelHandle and RuvTokenizer which are Send+Sync
#[cfg(all(
    target_os = "macos",
    target_arch = "aarch64",
    feature = "coreml",
    feature = "candle"
))]
unsafe impl<'a> Send for CoreMLStreamIterator<'a> {}

// =============================================================================
// Core ML Backend Implementation
// =============================================================================

/// Core ML backend for Apple Neural Engine acceleration
#[cfg(feature = "coreml")]
pub struct CoreMLBackend {
    /// Compute units preference
    compute_units: ComputeUnits,
    /// ANE capabilities
    ane_caps: AneCapabilities,
    /// Cache directory for converted models
    cache_dir: PathBuf,
    /// Model info
    model_info: Option<ModelInfo>,
    /// Whether model is loaded
    loaded: bool,
    /// The loaded Core ML model handle (only on macOS aarch64)
    #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
    model_handle: Option<CoreMLModelHandle>,
    /// The tokenizer for encoding/decoding text
    #[cfg(feature = "candle")]
    tokenizer: Option<crate::tokenizer::RuvTokenizer>,
    /// Input feature name for the model (e.g., "input_ids")
    input_feature_name: String,
    /// EOS token ID for stopping generation
    eos_token_id: u32,
    /// Vocab size
    vocab_size: usize,
}

#[cfg(feature = "coreml")]
impl std::fmt::Debug for CoreMLBackend {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CoreMLBackend")
            .field("compute_units", &self.compute_units)
            .field("ane_caps", &self.ane_caps)
            .field("cache_dir", &self.cache_dir)
            .field("model_info", &self.model_info)
            .field("loaded", &self.loaded)
            .field("input_feature_name", &self.input_feature_name)
            .field("eos_token_id", &self.eos_token_id)
            .field("vocab_size", &self.vocab_size)
            .finish()
    }
}

// Implement Send + Sync for CoreMLBackend
#[cfg(feature = "coreml")]
unsafe impl Send for CoreMLBackend {}
#[cfg(feature = "coreml")]
unsafe impl Sync for CoreMLBackend {}

#[cfg(feature = "coreml")]
impl Default for CoreMLBackend {
    fn default() -> Self {
        Self {
            compute_units: ComputeUnits::All,
            ane_caps: AneCapabilities::detect(),
            cache_dir: dirs::cache_dir()
                .unwrap_or_else(|| PathBuf::from("."))
                .join("ruvllm")
                .join("coreml"),
            model_info: None,
            loaded: false,
            #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
            model_handle: None,
            #[cfg(feature = "candle")]
            tokenizer: None,
            input_feature_name: "input_ids".to_string(),
            eos_token_id: 2,   // Common default EOS token
            vocab_size: 32000, // Common default vocab size
        }
    }
}

#[cfg(feature = "coreml")]
impl CoreMLBackend {
    /// Create a new Core ML backend
    pub fn new() -> Result<Self> {
        let caps = AneCapabilities::detect();

        if !caps.available {
            return Err(RuvLLMError::Config(
                "Apple Neural Engine not available on this device".to_string(),
            ));
        }

        let cache_dir = dirs::cache_dir()
            .unwrap_or_else(|| PathBuf::from("."))
            .join("ruvllm")
            .join("coreml");

        std::fs::create_dir_all(&cache_dir).map_err(|e| {
            RuvLLMError::Storage(format!("Failed to create Core ML cache directory: {}", e))
        })?;

        Ok(Self {
            compute_units: ComputeUnits::All,
            ane_caps: caps,
            cache_dir,
            model_info: None,
            loaded: false,
            #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
            model_handle: None,
            #[cfg(feature = "candle")]
            tokenizer: None,
            input_feature_name: "input_ids".to_string(),
            eos_token_id: 2,   // Common default EOS token
            vocab_size: 32000, // Common default vocab size
        })
    }

    /// Set the tokenizer for encoding/decoding text
    #[cfg(feature = "candle")]
    pub fn with_tokenizer(mut self, tokenizer: crate::tokenizer::RuvTokenizer) -> Self {
        self.eos_token_id = tokenizer.eos_token_id();
        self.vocab_size = tokenizer.vocab_size();
        self.tokenizer = Some(tokenizer);
        self
    }

    /// Set the input feature name for the model
    pub fn with_input_feature_name(mut self, name: impl Into<String>) -> Self {
        self.input_feature_name = name.into();
        self
    }

    /// Set the EOS token ID
    pub fn with_eos_token_id(mut self, eos_token_id: u32) -> Self {
        self.eos_token_id = eos_token_id;
        self
    }

    /// Set the vocab size
    pub fn with_vocab_size(mut self, vocab_size: usize) -> Self {
        self.vocab_size = vocab_size;
        self
    }

    /// Load tokenizer from HuggingFace Hub or local path
    #[cfg(feature = "candle")]
    pub fn load_tokenizer(&mut self, model_id_or_path: &str) -> Result<()> {
        let tokenizer = if std::path::Path::new(model_id_or_path).exists() {
            crate::tokenizer::RuvTokenizer::from_file(std::path::Path::new(model_id_or_path))?
        } else {
            crate::tokenizer::RuvTokenizer::from_pretrained(model_id_or_path)?
        };

        self.eos_token_id = tokenizer.eos_token_id();
        self.vocab_size = tokenizer.vocab_size();
        self.tokenizer = Some(tokenizer);
        Ok(())
    }

    /// Set compute units preference
    pub fn with_compute_units(mut self, units: ComputeUnits) -> Self {
        self.compute_units = units;
        self
    }

    /// Get ANE capabilities
    pub fn ane_capabilities(&self) -> &AneCapabilities {
        &self.ane_caps
    }

    /// Check if model is suitable for ANE acceleration
    pub fn is_model_ane_suitable(&self, model_size_mb: usize) -> bool {
        self.ane_caps.is_model_suitable(model_size_mb)
    }

    /// Get the Core ML model cache path for a given model
    fn get_coreml_cache_path(&self, model_path: &Path) -> PathBuf {
        let model_name = model_path
            .file_stem()
            .and_then(|s| s.to_str())
            .unwrap_or("model");
        self.cache_dir.join(format!("{}.mlmodelc", model_name))
    }

    /// Convert GGUF model to Core ML format
    ///
    /// Note: Full implementation would use coremltools or a Rust Core ML converter
    fn convert_to_coreml(&self, _gguf_path: &Path, _output_path: &Path) -> Result<()> {
        // TODO: Implement GGUF to Core ML conversion
        // This would involve:
        // 1. Parse GGUF weights and architecture
        // 2. Build Core ML model specification
        // 3. Compile to .mlmodelc
        //
        // For now, return a placeholder error
        Err(RuvLLMError::NotImplemented(
            "GGUF to Core ML conversion not yet implemented. \
             Use `coremltools` Python package to convert models, or \
             use pre-converted Core ML models."
                .to_string(),
        ))
    }

    /// Validate that a path points to a valid Core ML model
    fn validate_coreml_path(path: &Path) -> Result<()> {
        if !path.exists() {
            return Err(RuvLLMError::NotFound(format!(
                "Model path does not exist: {}",
                path.display()
            )));
        }

        let extension = path.extension().and_then(|e| e.to_str());
        match extension {
            Some("mlmodelc") => {
                // Compiled model - check if it's a directory with valid contents
                if !path.is_dir() {
                    return Err(RuvLLMError::CoreML(
                        ".mlmodelc should be a directory (compiled Core ML model)".to_string(),
                    ));
                }
                // Check for model.mil file or coremldata.bin (Core ML compiled model markers)
                let model_mil = path.join("model.mil");
                let coreml_data = path.join("coremldata.bin");
                let weights = path.join("weights");
                if !model_mil.exists() && !coreml_data.exists() && !weights.exists() {
                    return Err(RuvLLMError::CoreML(format!(
                        "Invalid .mlmodelc directory: missing expected files at {}",
                        path.display()
                    )));
                }
            }
            Some("mlmodel") => {
                // Uncompiled model - single file
                if !path.is_file() {
                    return Err(RuvLLMError::CoreML(".mlmodel should be a file".to_string()));
                }
            }
            Some("mlpackage") => {
                // ML Package format - directory with specific structure
                if !path.is_dir() {
                    return Err(RuvLLMError::CoreML(
                        ".mlpackage should be a directory".to_string(),
                    ));
                }
            }
            _ => {
                return Err(RuvLLMError::CoreML(format!(
                    "Unsupported Core ML model format. Expected .mlmodel, .mlmodelc, or .mlpackage: {}",
                    path.display()
                )));
            }
        }

        Ok(())
    }

    /// Get the loaded model handle (macOS aarch64 only)
    #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
    pub fn model_handle(&self) -> Option<&CoreMLModelHandle> {
        self.model_handle.as_ref()
    }

    /// Get the current compute units setting
    pub fn compute_units(&self) -> ComputeUnits {
        self.compute_units
    }

    /// Sample a token from logits using temperature and top-p sampling
    #[cfg(all(target_os = "macos", target_arch = "aarch64", feature = "candle"))]
    fn sample_token(&self, logits: &[f32], temperature: f32, top_p: f32) -> Result<u32> {
        use rand::Rng;

        if logits.is_empty() {
            return Err(RuvLLMError::Generation("Empty logits".to_string()));
        }

        // Apply temperature
        let scaled_logits: Vec<f32> = if temperature > 0.0 && temperature != 1.0 {
            logits.iter().map(|&x| x / temperature).collect()
        } else {
            logits.to_vec()
        };

        // Softmax to get probabilities
        let max_logit = scaled_logits
            .iter()
            .cloned()
            .fold(f32::NEG_INFINITY, f32::max);
        let exp_logits: Vec<f32> = scaled_logits
            .iter()
            .map(|&x| (x - max_logit).exp())
            .collect();
        let sum_exp: f32 = exp_logits.iter().sum();
        let probs: Vec<f32> = exp_logits.iter().map(|&x| x / sum_exp).collect();

        // Top-p (nucleus) sampling
        if top_p < 1.0 {
            let mut indexed_probs: Vec<(usize, f32)> = probs.iter().copied().enumerate().collect();
            indexed_probs
                .sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

            let mut cumsum = 0.0;
            let mut cutoff_idx = indexed_probs.len();
            for (i, (_, p)) in indexed_probs.iter().enumerate() {
                cumsum += p;
                if cumsum >= top_p {
                    cutoff_idx = i + 1;
                    break;
                }
            }

            let filtered: Vec<(usize, f32)> = indexed_probs[..cutoff_idx].to_vec();
            let filter_sum: f32 = filtered.iter().map(|(_, p)| p).sum();
            let normalized: Vec<(usize, f32)> = filtered
                .into_iter()
                .map(|(i, p)| (i, p / filter_sum))
                .collect();

            // Sample from filtered distribution
            let mut rng = rand::thread_rng();
            let r: f32 = rng.gen();
            let mut cumsum = 0.0;
            for (idx, p) in &normalized {
                cumsum += p;
                if r < cumsum {
                    return Ok(*idx as u32);
                }
            }
            // Fallback to last token in filtered set
            return Ok(normalized.last().map(|(i, _)| *i as u32).unwrap_or(0));
        }

        // Regular sampling from full distribution
        let mut rng = rand::thread_rng();
        let r: f32 = rng.gen();
        let mut cumsum = 0.0;
        for (idx, &p) in probs.iter().enumerate() {
            cumsum += p;
            if r < cumsum {
                return Ok(idx as u32);
            }
        }

        // Fallback to argmax
        Ok(probs
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .map(|(i, _)| i as u32)
            .unwrap_or(0))
    }
}

#[cfg(feature = "coreml")]
impl LlmBackend for CoreMLBackend {
    fn load_model(&mut self, model_id: &str, config: ModelConfig) -> Result<()> {
        let path = Path::new(model_id);

        // Check if it's already a Core ML model
        let extension = path.extension().and_then(|e| e.to_str());

        if matches!(extension, Some("mlmodelc" | "mlmodel" | "mlpackage")) {
            // Validate the Core ML model path
            Self::validate_coreml_path(path)?;

            #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
            {
                // Load the Core ML model using objc2-core-ml
                let handle = CoreMLModelHandle::load(path, self.compute_units)?;

                // Extract model info from the handle
                let input_names = handle.input_names();
                let output_names = handle.output_names();

                tracing::info!(
                    "Loaded Core ML model: {} (inputs: {:?}, outputs: {:?})",
                    path.display(),
                    input_names,
                    output_names
                );

                // Calculate model size from file/directory
                let memory_usage = if path.is_dir() {
                    // For directories, estimate by walking contents
                    walkdir_size(path).unwrap_or(0)
                } else {
                    std::fs::metadata(path)
                        .map(|m| m.len() as usize)
                        .unwrap_or(0)
                };

                // Store model info
                self.model_info = Some(ModelInfo {
                    name: path
                        .file_stem()
                        .and_then(|s| s.to_str())
                        .unwrap_or("unknown")
                        .to_string(),
                    architecture: config.architecture,
                    num_parameters: 0, // Would need to inspect model for this
                    vocab_size: config.vocab_size.unwrap_or(32000),
                    hidden_size: config.hidden_size.unwrap_or(4096),
                    num_layers: config.num_layers.unwrap_or(32),
                    max_context_length: config.max_sequence_length,
                    quantization: config.quantization,
                    memory_usage,
                });

                self.model_handle = Some(handle);
                self.loaded = true;
                return Ok(());
            }

            #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
            {
                return Err(RuvLLMError::Config(
                    "Core ML model loading is only supported on macOS aarch64 (Apple Silicon)"
                        .to_string(),
                ));
            }
        }

        // Check if it's a GGUF model that needs conversion
        if matches!(extension, Some("gguf")) {
            let coreml_path = self.get_coreml_cache_path(path);

            if !coreml_path.exists() {
                // Need to convert
                self.convert_to_coreml(path, &coreml_path)?;
            }

            // Recursively load the converted model
            return self.load_model(coreml_path.to_str().unwrap(), config);
        }

        Err(RuvLLMError::NotFound(format!(
            "Unsupported model format. Expected .mlmodel, .mlmodelc, .mlpackage, or .gguf: {}",
            model_id
        )))
    }

    fn generate(&self, prompt: &str, params: GenerateParams) -> Result<String> {
        if !self.loaded {
            return Err(RuvLLMError::InvalidOperation("No model loaded".to_string()));
        }

        #[cfg(all(target_os = "macos", target_arch = "aarch64", feature = "candle"))]
        {
            let model_handle = self.model_handle.as_ref().ok_or_else(|| {
                RuvLLMError::InvalidOperation("Model handle not initialized".to_string())
            })?;

            let tokenizer = self.tokenizer.as_ref().ok_or_else(|| {
                RuvLLMError::Config(
                    "Tokenizer not loaded. Call load_tokenizer() or use with_tokenizer() first."
                        .to_string(),
                )
            })?;

            // Encode the prompt
            let mut input_ids: Vec<i32> = tokenizer
                .encode(prompt)?
                .into_iter()
                .map(|t| t as i32)
                .collect();

            let max_tokens = params.max_tokens;
            let temperature = params.temperature;
            let top_p = params.top_p;
            let _start_time = Instant::now();

            let mut generated_tokens: Vec<u32> = Vec::with_capacity(max_tokens);

            // Autoregressive generation loop
            for _ in 0..max_tokens {
                // Run inference
                let logits = model_handle.predict(&self.input_feature_name, &input_ids)?;

                // Get logits for the last position (shape: [batch, seq, vocab] -> last token)
                let vocab_size = self.vocab_size;
                let last_token_logits = if logits.len() >= vocab_size {
                    &logits[logits.len() - vocab_size..]
                } else {
                    &logits
                };

                // Apply temperature and sample
                let next_token = self.sample_token(last_token_logits, temperature, top_p)?;

                // Check for EOS token
                if next_token == self.eos_token_id {
                    break;
                }

                // Add token to sequence
                generated_tokens.push(next_token);
                input_ids.push(next_token as i32);
            }

            // Decode generated tokens
            let output = tokenizer.decode(&generated_tokens)?;
            return Ok(output);
        }

        #[cfg(not(all(target_os = "macos", target_arch = "aarch64", feature = "candle")))]
        {
            let _ = (prompt, params);
            Err(RuvLLMError::Config(
                "Core ML inference requires macOS aarch64 with candle feature enabled".to_string(),
            ))
        }
    }

    fn generate_stream(
        &self,
        prompt: &str,
        params: GenerateParams,
    ) -> Result<Box<dyn Iterator<Item = Result<GeneratedToken>> + Send + '_>> {
        #[cfg(all(target_os = "macos", target_arch = "aarch64", feature = "candle"))]
        {
            if !self.loaded {
                return Err(RuvLLMError::InvalidOperation("No model loaded".to_string()));
            }

            let model_handle = self.model_handle.as_ref().ok_or_else(|| {
                RuvLLMError::InvalidOperation("Model handle not initialized".to_string())
            })?;

            let tokenizer = self.tokenizer.as_ref().ok_or_else(|| {
                RuvLLMError::Config(
                    "Tokenizer not loaded. Call load_tokenizer() or use with_tokenizer() first."
                        .to_string(),
                )
            })?;

            // Encode the prompt
            let input_ids: Vec<i32> = tokenizer
                .encode(prompt)?
                .into_iter()
                .map(|t| t as i32)
                .collect();

            let max_tokens = params.max_tokens;
            let temperature = params.temperature;
            let top_p = params.top_p;

            // Clone necessary data for the iterator
            let input_feature_name = self.input_feature_name.clone();
            let eos_token_id = self.eos_token_id;
            let vocab_size = self.vocab_size;

            // Generate tokens in iterator fashion
            let iter = CoreMLStreamIterator::new(
                model_handle,
                tokenizer,
                input_ids,
                max_tokens,
                temperature,
                top_p,
                input_feature_name,
                eos_token_id,
                vocab_size,
            );

            return Ok(Box::new(iter));
        }

        #[cfg(not(all(target_os = "macos", target_arch = "aarch64", feature = "candle")))]
        {
            let _ = (prompt, params);
            Err(RuvLLMError::Config(
                "Core ML streaming requires macOS aarch64 with candle feature enabled".to_string(),
            ))
        }
    }

    fn generate_stream_v2(&self, prompt: &str, params: GenerateParams) -> Result<TokenStream> {
        #[cfg(all(target_os = "macos", target_arch = "aarch64", feature = "candle"))]
        {
            if !self.loaded {
                return Err(RuvLLMError::InvalidOperation("No model loaded".to_string()));
            }

            let model_handle = self.model_handle.as_ref().ok_or_else(|| {
                RuvLLMError::InvalidOperation("Model handle not initialized".to_string())
            })?;

            let tokenizer = self.tokenizer.as_ref().ok_or_else(|| {
                RuvLLMError::Config(
                    "Tokenizer not loaded. Call load_tokenizer() or use with_tokenizer() first."
                        .to_string(),
                )
            })?;

            // Encode the prompt
            let mut input_ids: Vec<i32> = tokenizer
                .encode(prompt)?
                .into_iter()
                .map(|t| t as i32)
                .collect();

            let max_tokens = params.max_tokens;
            let temperature = params.temperature;
            let top_p = params.top_p;
            let start_time = Instant::now();

            // Create a channel for streaming
            let (tx, rx) = mpsc::channel::<StreamEvent>();

            // Generate tokens (no start event - StreamEvent doesn't have Start variant)
            let mut generated_count = 0;
            for _step in 0..max_tokens {
                let logits = match model_handle.predict(&self.input_feature_name, &input_ids) {
                    Ok(l) => l,
                    Err(e) => {
                        let _ = tx.send(StreamEvent::Error(e.to_string()));
                        break;
                    }
                };

                let last_token_logits = if logits.len() >= self.vocab_size {
                    &logits[logits.len() - self.vocab_size..]
                } else {
                    &logits
                };

                let next_token = match self.sample_token(last_token_logits, temperature, top_p) {
                    Ok(t) => t,
                    Err(e) => {
                        let _ = tx.send(StreamEvent::Error(e.to_string()));
                        break;
                    }
                };

                // Check for EOS
                if next_token == self.eos_token_id {
                    break;
                }

                // Decode the token
                let text = tokenizer.decode(&[next_token]).unwrap_or_default();

                // Send token event
                let _ = tx.send(StreamEvent::Token(GeneratedToken {
                    id: next_token,
                    text,
                    logprob: None,
                    is_special: next_token == self.eos_token_id,
                }));

                input_ids.push(next_token as i32);
                generated_count += 1;
            }

            // Send done event
            let elapsed = start_time.elapsed();
            let tokens_per_sec = generated_count as f64 / elapsed.as_secs_f64();
            let _ = tx.send(StreamEvent::Done {
                total_tokens: input_ids.len(),
                duration_ms: elapsed.as_millis() as u64,
                tokens_per_second: tokens_per_sec,
            });

            // Return the stream wrapped in TokenStream
            return Ok(TokenStream::new(rx));
        }

        #[cfg(not(all(target_os = "macos", target_arch = "aarch64", feature = "candle")))]
        {
            let _ = (prompt, params);
            Err(RuvLLMError::Config(
                "Core ML streaming requires macOS aarch64 with candle feature enabled".to_string(),
            ))
        }
    }

    fn get_embeddings(&self, text: &str) -> Result<Vec<f32>> {
        #[cfg(all(target_os = "macos", target_arch = "aarch64", feature = "candle"))]
        {
            if !self.loaded {
                return Err(RuvLLMError::InvalidOperation("No model loaded".to_string()));
            }

            let model_handle = self.model_handle.as_ref().ok_or_else(|| {
                RuvLLMError::InvalidOperation("Model handle not initialized".to_string())
            })?;

            let tokenizer = self.tokenizer.as_ref().ok_or_else(|| {
                RuvLLMError::Config(
                    "Tokenizer not loaded. Call load_tokenizer() or use with_tokenizer() first."
                        .to_string(),
                )
            })?;

            // Encode the text
            let token_ids: Vec<i32> = tokenizer
                .encode(text)?
                .into_iter()
                .map(|t| t as i32)
                .collect();

            // Get embeddings from the model
            let embeddings = model_handle.get_embeddings(
                &self.input_feature_name,
                &token_ids,
                None, // Use auto-detection for embedding output name
            )?;

            return Ok(embeddings);
        }

        #[cfg(not(all(target_os = "macos", target_arch = "aarch64", feature = "candle")))]
        {
            let _ = text;
            Err(RuvLLMError::Config(
                "Core ML embeddings require macOS aarch64 with candle feature enabled".to_string(),
            ))
        }
    }

    fn tokenizer(&self) -> Option<&dyn Tokenizer> {
        #[cfg(feature = "candle")]
        {
            self.tokenizer.as_ref().map(|t| t as &dyn Tokenizer)
        }
        #[cfg(not(feature = "candle"))]
        {
            None
        }
    }

    fn is_model_loaded(&self) -> bool {
        self.loaded
    }

    fn model_info(&self) -> Option<ModelInfo> {
        self.model_info.clone()
    }

    fn unload_model(&mut self) {
        self.loaded = false;
        self.model_info = None;
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        {
            self.model_handle = None;
        }
    }
}

/// Calculate directory size recursively (for .mlmodelc directories)
#[cfg(feature = "coreml")]
fn walkdir_size(path: &Path) -> std::io::Result<usize> {
    let mut total = 0;
    if path.is_dir() {
        for entry in std::fs::read_dir(path)? {
            let entry = entry?;
            let path = entry.path();
            if path.is_dir() {
                total += walkdir_size(&path)?;
            } else {
                total += std::fs::metadata(&path)?.len() as usize;
            }
        }
    } else {
        total = std::fs::metadata(path)?.len() as usize;
    }
    Ok(total)
}

/// Stub implementation when coreml feature is not enabled
#[cfg(not(feature = "coreml"))]
#[derive(Debug)]
pub struct CoreMLBackend;

#[cfg(not(feature = "coreml"))]
impl CoreMLBackend {
    pub fn new() -> Result<Self> {
        Err(RuvLLMError::Config(
            "Core ML feature not enabled. Enable with `coreml` feature flag.".to_string(),
        ))
    }
}

#[cfg(not(feature = "coreml"))]
impl LlmBackend for CoreMLBackend {
    fn load_model(&mut self, _model_id: &str, _config: ModelConfig) -> Result<()> {
        Err(RuvLLMError::Config(
            "Core ML feature not enabled".to_string(),
        ))
    }

    fn generate(&self, _prompt: &str, _params: GenerateParams) -> Result<String> {
        Err(RuvLLMError::Config(
            "Core ML feature not enabled".to_string(),
        ))
    }

    fn generate_stream(
        &self,
        _prompt: &str,
        _params: GenerateParams,
    ) -> Result<Box<dyn Iterator<Item = Result<GeneratedToken>> + Send + '_>> {
        Err(RuvLLMError::Config(
            "Core ML feature not enabled".to_string(),
        ))
    }

    fn generate_stream_v2(&self, _prompt: &str, _params: GenerateParams) -> Result<TokenStream> {
        Err(RuvLLMError::Config(
            "Core ML feature not enabled".to_string(),
        ))
    }

    fn get_embeddings(&self, _text: &str) -> Result<Vec<f32>> {
        Err(RuvLLMError::Config(
            "Core ML feature not enabled".to_string(),
        ))
    }

    fn tokenizer(&self) -> Option<&dyn Tokenizer> {
        None
    }

    fn is_model_loaded(&self) -> bool {
        false
    }

    fn model_info(&self) -> Option<ModelInfo> {
        None
    }

    fn unload_model(&mut self) {
        // No-op when feature not enabled
    }
}

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

    // ============================================================================
    // ComputeUnits Tests
    // ============================================================================

    #[test]
    fn test_compute_units_default() {
        let units = ComputeUnits::default();
        assert_eq!(units, ComputeUnits::All);
    }

    #[test]
    fn test_compute_units_uses_ane() {
        assert!(ComputeUnits::CpuAndNeuralEngine.uses_ane());
        assert!(ComputeUnits::All.uses_ane());
        assert!(!ComputeUnits::CpuOnly.uses_ane());
        assert!(!ComputeUnits::CpuAndGpu.uses_ane());
    }

    #[test]
    fn test_compute_units_uses_gpu() {
        assert!(ComputeUnits::CpuAndGpu.uses_gpu());
        assert!(ComputeUnits::All.uses_gpu());
        assert!(!ComputeUnits::CpuOnly.uses_gpu());
        assert!(!ComputeUnits::CpuAndNeuralEngine.uses_gpu());
    }

    #[test]
    fn test_compute_units_description() {
        assert_eq!(ComputeUnits::CpuOnly.description(), "CPU only");
        assert_eq!(ComputeUnits::CpuAndGpu.description(), "CPU + GPU");
        assert_eq!(
            ComputeUnits::CpuAndNeuralEngine.description(),
            "CPU + Neural Engine (ANE)"
        );
        assert_eq!(ComputeUnits::All.description(), "CPU + GPU + Neural Engine");
    }

    #[test]
    fn test_compute_units_clone() {
        let units = ComputeUnits::CpuAndNeuralEngine;
        let cloned = units.clone();
        assert_eq!(units, cloned);
    }

    #[test]
    fn test_compute_units_copy() {
        let units = ComputeUnits::All;
        let copied: ComputeUnits = units; // Copy semantics
        assert_eq!(units, copied);
    }

    #[test]
    fn test_compute_units_debug() {
        let debug_str = format!("{:?}", ComputeUnits::CpuAndNeuralEngine);
        assert!(debug_str.contains("CpuAndNeuralEngine"));
    }

    #[test]
    fn test_compute_units_eq() {
        assert_eq!(ComputeUnits::CpuOnly, ComputeUnits::CpuOnly);
        assert_ne!(ComputeUnits::CpuOnly, ComputeUnits::CpuAndGpu);
        assert_ne!(ComputeUnits::All, ComputeUnits::CpuAndNeuralEngine);
    }

    // ============================================================================
    // AneCapabilities Tests
    // ============================================================================

    #[test]
    fn test_ane_capabilities_detect() {
        let caps = AneCapabilities::detect();

        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        {
            assert!(caps.available);
            assert!(caps.tops > 0.0);
            assert!(!caps.supported_ops.is_empty());
        }

        #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
        {
            assert!(!caps.available);
        }
    }

    #[test]
    fn test_ane_capabilities_default() {
        let caps = AneCapabilities::default();
        // Default calls detect(), so behavior is platform-dependent
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        {
            assert!(caps.available);
        }
        #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
        {
            assert!(!caps.available);
        }
    }

    #[test]
    fn test_ane_capabilities_model_suitability() {
        let caps = AneCapabilities {
            available: true,
            tops: 38.0,
            max_model_size_mb: 2048,
            supported_ops: vec!["MatMul".to_string()],
        };

        assert!(caps.is_model_suitable(1000)); // 1GB model - fits
        assert!(caps.is_model_suitable(2048)); // 2GB model - at limit
        assert!(!caps.is_model_suitable(4096)); // 4GB model - too large
        assert!(caps.is_model_suitable(0)); // Edge case: 0 size
        assert!(caps.is_model_suitable(1)); // Edge case: tiny model
    }

    #[test]
    fn test_ane_capabilities_unavailable_device() {
        let caps = AneCapabilities {
            available: false,
            tops: 0.0,
            max_model_size_mb: 0,
            supported_ops: vec![],
        };

        // When ANE is unavailable, no model is suitable
        assert!(!caps.is_model_suitable(100));
        assert!(!caps.is_model_suitable(0));
    }

    #[test]
    fn test_ane_capabilities_clone() {
        let caps = AneCapabilities {
            available: true,
            tops: 38.0,
            max_model_size_mb: 2048,
            supported_ops: vec!["MatMul".to_string(), "GELU".to_string()],
        };
        let cloned = caps.clone();

        assert_eq!(caps.available, cloned.available);
        assert_eq!(caps.tops, cloned.tops);
        assert_eq!(caps.max_model_size_mb, cloned.max_model_size_mb);
        assert_eq!(caps.supported_ops, cloned.supported_ops);
    }

    #[test]
    fn test_ane_capabilities_debug() {
        let caps = AneCapabilities::detect();
        let debug_str = format!("{:?}", caps);
        assert!(debug_str.contains("AneCapabilities"));
        assert!(debug_str.contains("available"));
        assert!(debug_str.contains("tops"));
    }

    #[test]
    fn test_ane_capabilities_supported_ops() {
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        {
            let caps = AneCapabilities::detect();
            // Verify expected operations are supported
            assert!(caps.supported_ops.contains(&"MatMul".to_string()));
            assert!(caps.supported_ops.contains(&"GELU".to_string()));
            assert!(caps.supported_ops.contains(&"SiLU".to_string()));
            assert!(caps.supported_ops.contains(&"LayerNorm".to_string()));
            assert!(caps.supported_ops.contains(&"Softmax".to_string()));
        }
    }

    #[test]
    fn test_ane_capabilities_tops_reasonable() {
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        {
            let caps = AneCapabilities::detect();
            // M1 Pro starts at 11 TOPS, M4 Pro is 38 TOPS
            // Should be in reasonable range
            assert!(caps.tops >= 10.0);
            assert!(caps.tops <= 50.0);
        }
    }

    // ============================================================================
    // CoreMLBackend Tests (Feature-gated)
    // ============================================================================

    #[cfg(feature = "coreml")]
    mod coreml_backend_tests {
        use super::*;

        #[test]
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        fn test_coreml_backend_new_on_apple_silicon() {
            let backend = CoreMLBackend::new();
            assert!(backend.is_ok());

            let backend = backend.unwrap();
            assert!(!backend.is_model_loaded());
            assert!(backend.model_info().is_none());
        }

        #[test]
        #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
        fn test_coreml_backend_new_on_non_apple_silicon() {
            let backend = CoreMLBackend::new();
            assert!(backend.is_err());

            let err = backend.unwrap_err();
            assert!(err
                .to_string()
                .contains("Apple Neural Engine not available"));
        }

        #[test]
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        fn test_coreml_backend_with_compute_units() {
            let backend = CoreMLBackend::new()
                .unwrap()
                .with_compute_units(ComputeUnits::CpuAndNeuralEngine);

            assert_eq!(backend.compute_units(), ComputeUnits::CpuAndNeuralEngine);
        }

        #[test]
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        fn test_coreml_backend_ane_capabilities() {
            let backend = CoreMLBackend::new().unwrap();
            let caps = backend.ane_capabilities();

            assert!(caps.available);
            assert!(caps.tops > 0.0);
        }

        #[test]
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        fn test_coreml_backend_is_model_ane_suitable() {
            let backend = CoreMLBackend::new().unwrap();

            assert!(backend.is_model_ane_suitable(1000)); // 1GB
            assert!(backend.is_model_ane_suitable(2048)); // 2GB
            assert!(!backend.is_model_ane_suitable(5000)); // 5GB too large
        }

        #[test]
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        fn test_coreml_backend_unsupported_format() {
            let mut backend = CoreMLBackend::new().unwrap();

            // Try loading a file with unsupported extension
            let result = backend.load_model("model.safetensors", ModelConfig::default());
            assert!(result.is_err());

            let err = result.unwrap_err();
            assert!(err.to_string().contains("Unsupported model format"));
        }

        #[test]
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        fn test_coreml_backend_gguf_conversion_not_implemented() {
            let mut backend = CoreMLBackend::new().unwrap();

            // Try loading a GGUF file (conversion not implemented)
            let result = backend.load_model("/nonexistent/model.gguf", ModelConfig::default());
            assert!(result.is_err());

            let err = result.unwrap_err();
            assert!(
                err.to_string().contains("not yet implemented")
                    || err.to_string().contains("conversion")
            );
        }

        #[test]
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        fn test_coreml_backend_generate_requires_loaded_model() {
            let backend = CoreMLBackend::new().unwrap();

            let result = backend.generate("Hello", GenerateParams::default());
            assert!(result.is_err());

            let err = result.unwrap_err();
            assert!(err.to_string().contains("No model loaded"));
        }

        #[test]
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        fn test_coreml_backend_unload_model() {
            let mut backend = CoreMLBackend::new().unwrap();

            // Even without a model loaded, unload should be safe
            backend.unload_model();
            assert!(!backend.is_model_loaded());
        }

        #[test]
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        fn test_coreml_backend_tokenizer_not_available() {
            let backend = CoreMLBackend::new().unwrap();
            assert!(backend.tokenizer().is_none());
        }

        #[test]
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        fn test_coreml_backend_generate_stream_requires_model() {
            let backend = CoreMLBackend::new().unwrap();

            let result = backend.generate_stream("Hello", GenerateParams::default());
            assert!(result.is_err());

            // When no model is loaded, should return appropriate error
            match result {
                Err(err) => {
                    let msg = err.to_string();
                    // Should fail because either no model loaded or tokenizer not available
                    assert!(
                        msg.contains("No model loaded")
                            || msg.contains("Tokenizer")
                            || msg.contains("requires macOS aarch64"),
                        "Unexpected error: {}",
                        msg
                    );
                }
                Ok(_) => panic!("Expected error when no model loaded, got Ok"),
            }
        }

        #[test]
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        fn test_coreml_backend_get_embeddings_requires_model() {
            let backend = CoreMLBackend::new().unwrap();

            let result = backend.get_embeddings("Test text");
            assert!(result.is_err());

            let err = result.unwrap_err();
            let msg = err.to_string();
            // Should fail because either no model loaded or tokenizer not available
            assert!(
                msg.contains("No model loaded")
                    || msg.contains("Tokenizer")
                    || msg.contains("requires macOS aarch64"),
                "Unexpected error: {}",
                msg
            );
        }

        #[test]
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        fn test_coreml_backend_cache_directory() {
            let backend = CoreMLBackend::new().unwrap();

            // Cache dir should exist after backend creation
            assert!(backend.cache_dir.to_str().unwrap().contains("coreml"));
        }

        #[test]
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        fn test_coreml_backend_validate_path_nonexistent() {
            let result =
                CoreMLBackend::validate_coreml_path(Path::new("/nonexistent/model.mlmodel"));
            assert!(result.is_err());
            assert!(result.unwrap_err().to_string().contains("does not exist"));
        }

        #[test]
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        fn test_coreml_backend_validate_path_wrong_extension() {
            // Create a temp file with wrong extension
            let temp_dir = std::env::temp_dir();
            let temp_file = temp_dir.join("test_model.txt");
            std::fs::write(&temp_file, "test").unwrap();

            let result = CoreMLBackend::validate_coreml_path(&temp_file);
            assert!(result.is_err());
            assert!(result
                .unwrap_err()
                .to_string()
                .contains("Unsupported Core ML model format"));

            std::fs::remove_file(temp_file).ok();
        }
    }

    // ============================================================================
    // Stub Backend Tests (No Feature)
    // ============================================================================

    #[cfg(not(feature = "coreml"))]
    mod stub_backend_tests {
        use super::*;

        #[test]
        fn test_stub_backend_new_returns_error() {
            let result = CoreMLBackend::new();
            assert!(result.is_err());

            let err = result.unwrap_err();
            assert!(err.to_string().contains("feature not enabled"));
        }
    }

    // ============================================================================
    // LlmBackend Trait Implementation Tests
    // ============================================================================

    #[test]
    fn test_backend_trait_bounds() {
        // Verify CoreMLBackend implements Send + Sync (required by LlmBackend)
        fn assert_send_sync<T: Send + Sync>() {}

        #[cfg(feature = "coreml")]
        assert_send_sync::<CoreMLBackend>();
    }

    // ============================================================================
    // Edge Cases and Boundary Tests
    // ============================================================================

    #[test]
    fn test_model_suitability_boundary_values() {
        let caps = AneCapabilities {
            available: true,
            tops: 38.0,
            max_model_size_mb: 2048,
            supported_ops: vec!["MatMul".to_string()],
        };

        // At boundary
        assert!(caps.is_model_suitable(2048));
        // Just over boundary
        assert!(!caps.is_model_suitable(2049));
        // Just under boundary
        assert!(caps.is_model_suitable(2047));
    }

    #[test]
    fn test_compute_units_all_variants() {
        // Exhaustive test of all variants
        let variants = [
            ComputeUnits::CpuOnly,
            ComputeUnits::CpuAndGpu,
            ComputeUnits::CpuAndNeuralEngine,
            ComputeUnits::All,
        ];

        for variant in &variants {
            // Should not panic
            let _ = variant.description();
            let _ = variant.uses_ane();
            let _ = variant.uses_gpu();
            let _ = format!("{:?}", variant);
        }
    }

    #[test]
    fn test_ane_capabilities_empty_ops() {
        let caps = AneCapabilities {
            available: true,
            tops: 38.0,
            max_model_size_mb: 2048,
            supported_ops: vec![], // Empty ops list
        };

        // Should still work for suitability check
        assert!(caps.is_model_suitable(1000));
    }

    #[test]
    fn test_ane_capabilities_max_tops_value() {
        let caps = AneCapabilities {
            available: true,
            tops: f32::MAX,
            max_model_size_mb: usize::MAX,
            supported_ops: vec!["MatMul".to_string()],
        };

        // Should handle extreme values
        assert!(caps.available);
        assert!(caps.is_model_suitable(usize::MAX - 1));
    }

    #[test]
    fn test_ane_capabilities_zero_values() {
        let caps = AneCapabilities {
            available: true, // Available but with zero specs
            tops: 0.0,
            max_model_size_mb: 0,
            supported_ops: vec![],
        };

        // Model of size 0 should fit, size 1 should not
        assert!(caps.is_model_suitable(0));
        assert!(!caps.is_model_suitable(1));
    }
}