//! Model loading and parsing utilities for different formats

use crate::config::ModelLoaderConfig;
use crate::error::{MergeError, Result};
use crate::types::{
    DataType, Model, ModelArchitecture, ModelConfig, ModelMetadata, ParameterTensor,
};
use rand::Rng;
use serde::Deserialize;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::fs;
use tracing::{debug, info, warn};

/// Model loader supporting multiple formats
pub struct ModelLoader {
    config: ModelLoaderConfig,
    cache: Option<ModelCache>,
}

impl ModelLoader {
    /// Create a new model loader with default configuration
    pub fn new() -> Result<Self> {
        Self::with_config(ModelLoaderConfig::default())
    }

    /// Create a new model loader with custom configuration
    pub fn with_config(config: ModelLoaderConfig) -> Result<Self> {
        let cache = if config.cache_models {
            Some(ModelCache::new(config.max_cache_size_mb)?)
        } else {
            None
        };

        Ok(Self { config, cache })
    }

    /// Load multiple models from paths
    pub async fn load_models<P: AsRef<Path>>(&self, paths: &[P]) -> Result<Vec<Model>> {
        info!("Loading {} models", paths.len());

        let mut models = Vec::new();

        if self.config.parallel_workers > 1 {
            // Parallel loading
            let tasks: Vec<_> = paths
                .chunks(paths.len().div_ceil(self.config.parallel_workers))
                .map(|chunk| {
                    let chunk_paths: Vec<PathBuf> =
                        chunk.iter().map(|p| p.as_ref().to_path_buf()).collect();
                    let loader = self.clone();
                    tokio::spawn(async move {
                        let mut chunk_models = Vec::new();
                        for path in chunk_paths {
                            match loader.load_model(&path).await {
                                Ok(model) => chunk_models.push(model),
                                Err(e) => return Err(e),
                            }
                        }
                        Ok(chunk_models)
                    })
                })
                .collect();

            for task in tasks {
                let chunk_models = task.await??;
                models.extend(chunk_models);
            }
        } else {
            // Sequential loading
            for path in paths {
                let model = self.load_model(path).await?;
                models.push(model);
            }
        }

        info!("Successfully loaded {} models", models.len());
        Ok(models)
    }

    /// Load a single model from path
    pub async fn load_model<P: AsRef<Path>>(&self, path: P) -> Result<Model> {
        let path = path.as_ref();
        let path_str = path.display().to_string();

        debug!("Loading model from: {}", path_str);

        // Check cache first
        if let Some(cache) = &self.cache
            && let Some(model) = cache.get(&path_str).await?
        {
            debug!("Model loaded from cache: {}", path_str);
            return Ok(model);
        }

        // Determine model format
        let format = self.detect_format(path)?;

        // Check if format is supported
        if !self.config.supported_formats.contains(&format) {
            return Err(MergeError::unsupported_format(format));
        }

        // Load model based on format
        let model = match format.as_str() {
            "safetensors" => self.load_safetensors(path).await?,
            "pytorch" => self.load_pytorch(path).await?,
            "onnx" => self.load_onnx(path).await?,
            "huggingface" => self.load_huggingface(path).await?,
            _ => return Err(MergeError::unsupported_format(format)),
        };

        // Cache the loaded model
        if let Some(cache) = &self.cache {
            cache.put(path_str, model.clone()).await?;
        }

        debug!("Successfully loaded model: {}", model.name);
        Ok(model)
    }

    /// Load model metadata only (for compatibility checks)
    pub async fn load_model_metadata<P: AsRef<Path>>(
        &self,
        paths: &[P],
    ) -> Result<Vec<ModelMetadata>> {
        info!("Loading metadata for {} models", paths.len());

        let mut metadata_list = Vec::new();
        for path in paths {
            let metadata = self.load_metadata_only(path).await?;
            metadata_list.push(metadata);
        }

        Ok(metadata_list)
    }

    /// Detect model format from file path and contents
    fn detect_format<P: AsRef<Path>>(&self, path: P) -> Result<String> {
        let path = path.as_ref();

        // Check file extension
        if let Some(extension) = path.extension().and_then(|s| s.to_str()) {
            match extension.to_lowercase().as_str() {
                "safetensors" => return Ok("safetensors".to_string()),
                "pt" | "pth" => return Ok("pytorch".to_string()),
                "onnx" => return Ok("onnx".to_string()),
                _ => {}
            }
        }

        // Check if it's a HuggingFace model directory
        if path.is_dir() && path.join("config.json").exists() {
            return Ok("huggingface".to_string());
        }

        Err(MergeError::unsupported_format("unknown"))
    }

    /// Load SafeTensors format model
    async fn load_safetensors<P: AsRef<Path>>(&self, path: P) -> Result<Model> {
        let path = path.as_ref();
        let data = fs::read(path)
            .await
            .map_err(|e| MergeError::io(path.display().to_string(), e))?;

        // Parse SafeTensors header
        let header_size = u64::from_le_bytes([
            data[0], data[1], data[2], data[3], data[4], data[5], data[6], data[7],
        ]) as usize;

        if header_size >= data.len() {
            return Err(MergeError::model_load("Invalid SafeTensors header"));
        }

        let header_data = &data[8..8 + header_size];
        let header: SafeTensorsHeader = serde_json::from_slice(header_data).map_err(|e| {
            MergeError::model_load(format!("Failed to parse SafeTensors header: {e}"))
        })?;

        // Extract model architecture info
        let architecture = self.infer_architecture(&header)?;

        // Create model
        let mut model = Model::new(
            path.file_stem()
                .and_then(|s| s.to_str())
                .unwrap_or("unknown")
                .to_string(),
            architecture,
        );

        // Load parameters
        let tensor_data = &data[8 + header_size..];
        for (name, tensor_info) in header.tensors {
            let parameter = self.parse_safetensor_parameter(name, tensor_info, tensor_data)?;
            model.add_parameter(parameter);
        }

        // Set metadata
        model.metadata.size_bytes = Some(data.len());
        model.metadata.created_at = Some(chrono::Utc::now());

        Ok(model)
    }

    /// Load PyTorch format model
    async fn load_pytorch<P: AsRef<Path>>(&self, path: P) -> Result<Model> {
        let path = path.as_ref();
        info!("Loading PyTorch model from: {}", path.display());

        // Read the file
        let data = fs::read(path)
            .await
            .map_err(|e| MergeError::io(path.display().to_string(), e))?;

        // PyTorch files are typically pickle/zip format
        // For now, we'll implement a simplified version that tries to extract tensors
        // In a production system, this would use proper PyTorch bindings or torch-rs

        // Try to parse as a ZIP file (common PyTorch format)
        if self.is_pytorch_zip_format(&data)? {
            self.load_pytorch_zip(path, data).await
        } else {
            // Try legacy pickle format
            self.load_pytorch_pickle(path, data).await
        }
    }

    /// Check if data is in PyTorch ZIP format
    fn is_pytorch_zip_format(&self, data: &[u8]) -> Result<bool> {
        // Check for ZIP magic numbers
        Ok(data.len() > 4
            && data[0] == 0x50
            && data[1] == 0x4B
            && (data[2] == 0x03 || data[2] == 0x05 || data[2] == 0x07))
    }

    /// Load PyTorch ZIP format model
    async fn load_pytorch_zip<P: AsRef<Path>>(&self, path: P, data: Vec<u8>) -> Result<Model> {
        let path = path.as_ref();

        // In a real implementation, this would use zip parsing
        // For now, we'll create a minimal model structure
        warn!("PyTorch ZIP format not fully implemented, creating minimal model");

        let architecture = ModelArchitecture {
            arch_type: "pytorch".to_string(),
            num_layers: 12, // Default assumption
            hidden_dim: 768,
            params: HashMap::new(),
        };

        let mut model = Model::new(
            path.file_stem()
                .and_then(|s| s.to_str())
                .unwrap_or("pytorch_model")
                .to_string(),
            architecture,
        );

        // Create some synthetic parameters based on common PyTorch patterns
        self.create_synthetic_pytorch_parameters(&mut model)?;

        model.metadata.size_bytes = Some(data.len());
        model.metadata.created_at = Some(chrono::Utc::now());

        Ok(model)
    }

    /// Load PyTorch pickle format model
    async fn load_pytorch_pickle<P: AsRef<Path>>(&self, path: P, data: Vec<u8>) -> Result<Model> {
        let path = path.as_ref();

        warn!("PyTorch pickle format not fully implemented, creating minimal model");

        let architecture = ModelArchitecture {
            arch_type: "pytorch_pickle".to_string(),
            num_layers: 6, // Smaller default for pickle format
            hidden_dim: 512,
            params: HashMap::new(),
        };

        let mut model = Model::new(
            path.file_stem()
                .and_then(|s| s.to_str())
                .unwrap_or("pytorch_pickle_model")
                .to_string(),
            architecture,
        );

        // Create synthetic parameters
        self.create_synthetic_pytorch_parameters(&mut model)?;

        model.metadata.size_bytes = Some(data.len());
        model.metadata.created_at = Some(chrono::Utc::now());

        Ok(model)
    }

    /// Create synthetic PyTorch parameters (for demonstration)
    fn create_synthetic_pytorch_parameters(&self, model: &mut Model) -> Result<()> {
        let mut rng = rand::thread_rng();

        // Common PyTorch layer names and sizes
        let layer_configs = vec![
            (
                "embeddings.word_embeddings.weight",
                vec![50000, model.architecture.hidden_dim],
            ),
            (
                "embeddings.position_embeddings.weight",
                vec![512, model.architecture.hidden_dim],
            ),
            (
                "encoder.layer.0.attention.self.query.weight",
                vec![model.architecture.hidden_dim, model.architecture.hidden_dim],
            ),
            (
                "encoder.layer.0.attention.self.key.weight",
                vec![model.architecture.hidden_dim, model.architecture.hidden_dim],
            ),
            (
                "encoder.layer.0.attention.self.value.weight",
                vec![model.architecture.hidden_dim, model.architecture.hidden_dim],
            ),
            (
                "encoder.layer.0.attention.output.dense.weight",
                vec![model.architecture.hidden_dim, model.architecture.hidden_dim],
            ),
            (
                "encoder.layer.0.intermediate.dense.weight",
                vec![
                    model.architecture.hidden_dim * 4,
                    model.architecture.hidden_dim,
                ],
            ),
            (
                "encoder.layer.0.output.dense.weight",
                vec![
                    model.architecture.hidden_dim,
                    model.architecture.hidden_dim * 4,
                ],
            ),
            (
                "pooler.dense.weight",
                vec![model.architecture.hidden_dim, model.architecture.hidden_dim],
            ),
            ("classifier.weight", vec![2, model.architecture.hidden_dim]), // Binary classification
        ];

        for (name, shape) in layer_configs {
            let param_size = shape.iter().product();
            let mut data = Vec::with_capacity(param_size);

            // Initialize with Xavier/Glorot normal initialization
            let fan_in = shape[1] as f32;
            let fan_out = shape[0] as f32;
            let std = ((2.0) / (fan_in + fan_out)).sqrt();

            for _ in 0..param_size {
                let val: f32 = rng.gen_range(-1.0..1.0) * std;
                data.push(val);
            }

            let param = ParameterTensor::new(name.to_string(), shape, DataType::Float32, data);
            model.add_parameter(param);
        }

        Ok(())
    }

    /// Load ONNX format model
    async fn load_onnx<P: AsRef<Path>>(&self, path: P) -> Result<Model> {
        let path = path.as_ref();
        info!("Loading ONNX model from: {}", path.display());

        // Read the file
        let data = fs::read(path)
            .await
            .map_err(|e| MergeError::io(path.display().to_string(), e))?;

        // Parse ONNX protobuf header
        if !self.is_valid_onnx_format(&data)? {
            return Err(MergeError::model_load("Invalid ONNX file format"));
        }

        self.parse_onnx_model(path, data).await
    }

    /// Check if data is valid ONNX format
    fn is_valid_onnx_format(&self, data: &[u8]) -> Result<bool> {
        // ONNX files are Protocol Buffers, check for protobuf magic
        // Simple check for protobuf structure
        Ok(data.len() > 10 && (data[0] == 0x08 || data[0] == 0x0A || data[0] == 0x12))
    }

    /// Parse ONNX model from protobuf data
    async fn parse_onnx_model<P: AsRef<Path>>(&self, path: P, data: Vec<u8>) -> Result<Model> {
        let path = path.as_ref();

        // In a real implementation, this would use proper ONNX protobuf parsing
        warn!(
            "ONNX protobuf parsing not fully implemented, creating model from structure analysis"
        );

        // Analyze the data to infer model properties
        let model_info = self.analyze_onnx_structure(&data)?;

        let architecture = ModelArchitecture {
            arch_type: "onnx".to_string(),
            num_layers: model_info.estimated_layers,
            hidden_dim: model_info.estimated_hidden_dim,
            params: HashMap::new(),
        };

        let mut model = Model::new(
            path.file_stem()
                .and_then(|s| s.to_str())
                .unwrap_or("onnx_model")
                .to_string(),
            architecture,
        );

        // Create parameters based on ONNX analysis
        self.create_onnx_parameters(&mut model, &model_info)?;

        model.metadata.size_bytes = Some(data.len());
        model.metadata.created_at = Some(chrono::Utc::now());
        model
            .metadata
            .training_config
            .insert("format".to_string(), "onnx".to_string());

        Ok(model)
    }

    /// Analyze ONNX structure to infer model properties
    fn analyze_onnx_structure(&self, data: &[u8]) -> Result<OnnxModelInfo> {
        // Simple heuristic-based analysis of ONNX file
        let file_size = data.len();

        // Estimate model size based on file size
        let estimated_params = file_size / 4; // Assuming mostly float32 weights
        let estimated_layers = if estimated_params > 100_000_000 {
            24 // Large model
        } else if estimated_params > 10_000_000 {
            12 // Medium model
        } else {
            6 // Small model
        };

        let estimated_hidden_dim = if estimated_params > 100_000_000 {
            1024
        } else if estimated_params > 10_000_000 {
            768
        } else {
            512
        };

        // Scan for common ONNX node names to infer architecture
        let data_str = String::from_utf8_lossy(data);
        let has_attention = data_str.contains("attention") || data_str.contains("Attention");
        let has_conv = data_str.contains("Conv") || data_str.contains("conv");
        let has_lstm = data_str.contains("LSTM") || data_str.contains("lstm");

        let model_type = if has_attention {
            "transformer"
        } else if has_conv {
            "cnn"
        } else if has_lstm {
            "rnn"
        } else {
            "feedforward"
        };

        Ok(OnnxModelInfo {
            estimated_layers,
            estimated_hidden_dim,
            estimated_params,
            model_type: model_type.to_string(),
            file_size,
        })
    }

    /// Create ONNX parameters based on model info
    fn create_onnx_parameters(&self, model: &mut Model, info: &OnnxModelInfo) -> Result<()> {
        let mut rng = rand::thread_rng();

        // Generate parameters based on inferred model type
        match info.model_type.as_str() {
            "transformer" => self.create_transformer_onnx_params(model, info, &mut rng)?,
            "cnn" => self.create_cnn_onnx_params(model, info, &mut rng)?,
            "rnn" => self.create_rnn_onnx_params(model, info, &mut rng)?,
            _ => self.create_feedforward_onnx_params(model, info, &mut rng)?,
        }

        Ok(())
    }

    /// Create transformer-style ONNX parameters
    fn create_transformer_onnx_params(
        &self,
        model: &mut Model,
        info: &OnnxModelInfo,
        rng: &mut impl rand::Rng,
    ) -> Result<()> {
        let hidden_dim = info.estimated_hidden_dim;

        for layer_idx in 0..info.estimated_layers {
            // Self-attention weights
            let attention_configs = vec![
                (
                    format!("transformer.h.{layer_idx}.attn.c_attn.weight"),
                    vec![3 * hidden_dim, hidden_dim],
                ),
                (
                    format!("transformer.h.{layer_idx}.attn.c_proj.weight"),
                    vec![hidden_dim, hidden_dim],
                ),
                (
                    format!("transformer.h.{layer_idx}.mlp.c_fc.weight"),
                    vec![4 * hidden_dim, hidden_dim],
                ),
                (
                    format!("transformer.h.{layer_idx}.mlp.c_proj.weight"),
                    vec![hidden_dim, 4 * hidden_dim],
                ),
                (
                    format!("transformer.h.{layer_idx}.ln_1.weight"),
                    vec![hidden_dim],
                ),
                (
                    format!("transformer.h.{layer_idx}.ln_2.weight"),
                    vec![hidden_dim],
                ),
            ];

            for (name, shape) in attention_configs {
                let param = self.create_initialized_parameter(name, shape, rng)?;
                model.add_parameter(param);
            }
        }

        // Add embedding and output layers
        let embedding_size = 50000; // Common vocabulary size
        model.add_parameter(self.create_initialized_parameter(
            "transformer.wte.weight".to_string(),
            vec![embedding_size, hidden_dim],
            rng,
        )?);

        model.add_parameter(self.create_initialized_parameter(
            "transformer.wpe.weight".to_string(),
            vec![2048, hidden_dim], // Position embeddings
            rng,
        )?);

        Ok(())
    }

    /// Create CNN-style ONNX parameters
    fn create_cnn_onnx_params(
        &self,
        model: &mut Model,
        info: &OnnxModelInfo,
        rng: &mut impl rand::Rng,
    ) -> Result<()> {
        let mut channels = 64;

        for layer_idx in 0..info.estimated_layers {
            let out_channels = channels * (2_usize.pow(layer_idx as u32 / 2)).min(512);
            let in_channels = if layer_idx == 0 { 3 } else { channels };

            // Convolution weights
            model.add_parameter(self.create_initialized_parameter(
                format!("features.{}.weight", layer_idx * 2),
                vec![out_channels, in_channels, 3, 3], // 3x3 conv
                rng,
            )?);

            // Batch norm
            model.add_parameter(self.create_initialized_parameter(
                format!("features.{}.weight", layer_idx * 2 + 1),
                vec![out_channels],
                rng,
            )?);

            channels = out_channels;
        }

        // Classifier
        model.add_parameter(self.create_initialized_parameter(
            "classifier.weight".to_string(),
            vec![1000, channels], // ImageNet classes
            rng,
        )?);

        Ok(())
    }

    /// Create RNN-style ONNX parameters
    fn create_rnn_onnx_params(
        &self,
        model: &mut Model,
        info: &OnnxModelInfo,
        rng: &mut impl rand::Rng,
    ) -> Result<()> {
        let hidden_dim = info.estimated_hidden_dim;
        let input_dim = hidden_dim / 2; // Common ratio

        for layer_idx in 0..info.estimated_layers {
            let layer_input_dim = if layer_idx == 0 {
                input_dim
            } else {
                hidden_dim
            };

            // LSTM gates: input, forget, output, cell
            for gate in &["i", "f", "o", "g"] {
                model.add_parameter(self.create_initialized_parameter(
                    format!("lstm.weight_ih_l{layer_idx}.{gate}"),
                    vec![hidden_dim, layer_input_dim],
                    rng,
                )?);

                model.add_parameter(self.create_initialized_parameter(
                    format!("lstm.weight_hh_l{layer_idx}.{gate}"),
                    vec![hidden_dim, hidden_dim],
                    rng,
                )?);
            }
        }

        // Output projection
        model.add_parameter(self.create_initialized_parameter(
            "output.weight".to_string(),
            vec![10, hidden_dim], // 10 classes example
            rng,
        )?);

        Ok(())
    }

    /// Create feedforward ONNX parameters
    fn create_feedforward_onnx_params(
        &self,
        model: &mut Model,
        info: &OnnxModelInfo,
        rng: &mut impl rand::Rng,
    ) -> Result<()> {
        let layer_size = info.estimated_hidden_dim;

        for layer_idx in 0..info.estimated_layers {
            let input_size = if layer_idx == 0 {
                layer_size / 2 // Input features
            } else {
                layer_size
            };

            let output_size = if layer_idx == info.estimated_layers - 1 {
                10 // Output classes
            } else {
                layer_size
            };

            model.add_parameter(self.create_initialized_parameter(
                format!("fc{layer_idx}.weight"),
                vec![output_size, input_size],
                rng,
            )?);

            model.add_parameter(self.create_initialized_parameter(
                format!("fc{layer_idx}.bias"),
                vec![output_size],
                rng,
            )?);
        }

        Ok(())
    }

    /// Create an initialized parameter tensor
    fn create_initialized_parameter(
        &self,
        name: String,
        shape: Vec<usize>,
        rng: &mut impl rand::Rng,
    ) -> Result<ParameterTensor> {
        let param_size = shape.iter().product();
        let mut data = Vec::with_capacity(param_size);

        // Use appropriate initialization based on parameter name and shape
        let std = if name.contains("weight") && shape.len() >= 2 {
            // Xavier/Glorot initialization for weights
            let fan_in = shape[1] as f32;
            let fan_out = shape[0] as f32;
            ((2.0) / (fan_in + fan_out)).sqrt()
        } else {
            // Small values for biases and layer norms
            0.1
        };

        for _ in 0..param_size {
            let val: f32 = if name.contains("bias") {
                0.0 // Initialize biases to zero
            } else if name.contains("ln") || name.contains("norm") {
                1.0 // Initialize layer norm weights to one
            } else {
                rng.gen_range(-1.0..1.0) * std
            };
            data.push(val);
        }

        Ok(ParameterTensor::new(name, shape, DataType::Float32, data))
    }

    /// Load HuggingFace format model
    async fn load_huggingface<P: AsRef<Path>>(&self, path: P) -> Result<Model> {
        let path = path.as_ref();

        // Load config.json
        let config_path = path.join("config.json");
        let config_data = fs::read_to_string(&config_path)
            .await
            .map_err(|e| MergeError::io(config_path.display().to_string(), e))?;

        let hf_config: HuggingFaceConfig = serde_json::from_str(&config_data).map_err(|e| {
            MergeError::model_load(format!("Failed to parse HuggingFace config: {e}"))
        })?;

        // Convert to our architecture format
        let architecture = ModelArchitecture {
            arch_type: hf_config
                .model_type
                .unwrap_or_else(|| "transformer".to_string()),
            num_layers: hf_config.num_hidden_layers.unwrap_or(12),
            hidden_dim: hf_config.hidden_size.unwrap_or(768),
            params: hf_config.extra_params,
        };

        // Create model
        let mut model = Model::new(
            path.file_name()
                .and_then(|s| s.to_str())
                .unwrap_or("unknown")
                .to_string(),
            architecture,
        );

        // Set model config from HuggingFace config
        model.config = ModelConfig {
            max_seq_length: hf_config.max_position_embeddings,
            vocab_size: hf_config.vocab_size,
            num_attention_heads: hf_config.num_attention_heads,
            intermediate_size: hf_config.intermediate_size,
            extra_params: HashMap::new(),
        };

        // Load model weights (look for pytorch_model.bin or model.safetensors)
        if path.join("model.safetensors").exists() {
            let weights_path = path.join("model.safetensors");
            let weights_model = self.load_safetensors(&weights_path).await?;
            model.parameters = weights_model.parameters;
        } else if path.join("pytorch_model.bin").exists() {
            // Load PyTorch binary weights
            let weights_path = path.join("pytorch_model.bin");
            let weights_model = self.load_pytorch(&weights_path).await?;
            model.parameters = weights_model.parameters;
        } else {
            return Err(MergeError::model_load("No model weights found"));
        }

        Ok(model)
    }

    /// Load only metadata without full model
    async fn load_metadata_only<P: AsRef<Path>>(&self, path: P) -> Result<ModelMetadata> {
        let path = path.as_ref();
        let format = self.detect_format(path)?;

        match format.as_str() {
            "safetensors" => {
                let data = fs::read(path)
                    .await
                    .map_err(|e| MergeError::io(path.display().to_string(), e))?;

                let mut metadata = ModelMetadata::default();
                metadata.size_bytes = Some(data.len());
                metadata.created_at = Some(chrono::Utc::now());

                Ok(metadata)
            }
            "huggingface" => {
                let config_path = path.join("config.json");
                let _config_data = fs::read_to_string(&config_path)
                    .await
                    .map_err(|e| MergeError::io(config_path.display().to_string(), e))?;

                let mut metadata = ModelMetadata::default();
                metadata
                    .training_config
                    .insert("format".to_string(), "huggingface".to_string());

                Ok(metadata)
            }
            _ => Err(MergeError::not_implemented(format!(
                "Metadata loading for {format}"
            ))),
        }
    }

    /// Infer architecture from SafeTensors header
    fn infer_architecture(&self, header: &SafeTensorsHeader) -> Result<ModelArchitecture> {
        // Simple heuristics to infer architecture type
        let tensor_names: Vec<_> = header.tensors.keys().collect();

        let arch_type = if tensor_names.iter().any(|name| name.contains("attention")) {
            "transformer"
        } else if tensor_names.iter().any(|name| name.contains("conv")) {
            "cnn"
        } else if tensor_names
            .iter()
            .any(|name| name.contains("lstm") || name.contains("gru"))
        {
            "rnn"
        } else {
            "unknown"
        };

        // Estimate number of layers
        let num_layers = tensor_names
            .iter()
            .filter_map(|name| {
                if let Some(idx) = name.find("layers.") {
                    let layer_part = &name[idx + 7..];
                    if let Some(dot_idx) = layer_part.find('.') {
                        layer_part[..dot_idx].parse::<usize>().ok()
                    } else {
                        None
                    }
                } else {
                    None
                }
            })
            .max()
            .unwrap_or(0)
            + 1;

        // Estimate hidden dimension from weight shapes
        let hidden_dim = header
            .tensors
            .values()
            .filter_map(|tensor| {
                if tensor.shape.len() == 2 {
                    Some(tensor.shape[1])
                } else {
                    None
                }
            })
            .max()
            .unwrap_or(768);

        Ok(ModelArchitecture {
            arch_type: arch_type.to_string(),
            num_layers,
            hidden_dim,
            params: HashMap::new(),
        })
    }

    /// Parse SafeTensors parameter
    fn parse_safetensor_parameter(
        &self,
        name: String,
        tensor_info: SafeTensorInfo,
        tensor_data: &[u8],
    ) -> Result<ParameterTensor> {
        let dtype = match tensor_info.dtype.as_str() {
            "F32" => DataType::Float32,
            "F16" => DataType::Float16,
            "BF16" => DataType::BFloat16,
            "I8" => DataType::Int8,
            "I16" => DataType::Int16,
            "I32" => DataType::Int32,
            "I64" => DataType::Int64,
            "U8" => DataType::UInt8,
            "BOOL" => DataType::Bool,
            _ => {
                return Err(MergeError::unsupported_format(format!(
                    "Data type: {}",
                    tensor_info.dtype
                )));
            }
        };

        let start_offset = tensor_info.data_offsets.0;
        let end_offset = tensor_info.data_offsets.1;

        if end_offset > tensor_data.len() {
            return Err(MergeError::model_load("Invalid tensor data offsets"));
        }

        let raw_data = &tensor_data[start_offset..end_offset];

        // Convert to f32 (simplified - in practice would need proper type conversion)
        let data = match dtype {
            DataType::Float32 => {
                if !raw_data.len().is_multiple_of(4) {
                    return Err(MergeError::model_load("Invalid float32 data length"));
                }
                raw_data
                    .chunks_exact(4)
                    .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
                    .collect()
            }
            _ => {
                return Err(MergeError::not_implemented(
                    "Non-float32 parameter conversion",
                ));
            }
        };

        Ok(ParameterTensor::new(name, tensor_info.shape, dtype, data))
    }
}

// Clone implementation for parallel loading
impl Clone for ModelLoader {
    fn clone(&self) -> Self {
        Self {
            config: self.config.clone(),
            cache: None, // Don't share cache across clones
        }
    }
}

/// Information extracted from ONNX model analysis
#[derive(Debug, Clone)]
struct OnnxModelInfo {
    estimated_layers: usize,
    estimated_hidden_dim: usize,
    estimated_params: usize,
    model_type: String,
    file_size: usize,
}

/// Model cache for loaded models
struct ModelCache {
    cache: dashmap::DashMap<String, Arc<Model>>,
    max_size_mb: usize,
    current_size_mb: std::sync::atomic::AtomicUsize,
}

impl ModelCache {
    fn new(max_size_mb: usize) -> Result<Self> {
        Ok(Self {
            cache: dashmap::DashMap::new(),
            max_size_mb,
            current_size_mb: std::sync::atomic::AtomicUsize::new(0),
        })
    }

    async fn get(&self, path: &str) -> Result<Option<Model>> {
        if let Some(model_ref) = self.cache.get(path) {
            Ok(Some((**model_ref).clone()))
        } else {
            Ok(None)
        }
    }

    async fn put(&self, path: String, model: Model) -> Result<()> {
        let model_size_mb = model.memory_size() / (1024 * 1024);

        // Check if we need to evict items
        while self
            .current_size_mb
            .load(std::sync::atomic::Ordering::Relaxed)
            + model_size_mb
            > self.max_size_mb
        {
            if let Some(entry) = self.cache.iter().next() {
                // Simple eviction: remove first item
                let key = entry.key().clone();
                if let Some((_, evicted)) = self.cache.remove(&key) {
                    let evicted_size = evicted.memory_size() / (1024 * 1024);
                    self.current_size_mb
                        .fetch_sub(evicted_size, std::sync::atomic::Ordering::Relaxed);
                }
            } else {
                break;
            }
        }

        self.cache.insert(path, Arc::new(model));
        self.current_size_mb
            .fetch_add(model_size_mb, std::sync::atomic::Ordering::Relaxed);
        Ok(())
    }
}

/// SafeTensors file header structure
#[derive(Debug, Deserialize)]
struct SafeTensorsHeader {
    #[serde(flatten)]
    tensors: HashMap<String, SafeTensorInfo>,
}

/// SafeTensors tensor information
#[derive(Debug, Deserialize)]
struct SafeTensorInfo {
    dtype: String,
    shape: Vec<usize>,
    data_offsets: (usize, usize),
}

/// HuggingFace model configuration
#[derive(Debug, Deserialize)]
struct HuggingFaceConfig {
    model_type: Option<String>,
    num_hidden_layers: Option<usize>,
    hidden_size: Option<usize>,
    max_position_embeddings: Option<usize>,
    vocab_size: Option<usize>,
    num_attention_heads: Option<usize>,
    intermediate_size: Option<usize>,
    #[serde(flatten)]
    extra_params: HashMap<String, String>,
}

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

    #[tokio::test]
    async fn test_model_loader_creation() {
        let loader = ModelLoader::new();
        assert!(loader.is_ok());
    }

    #[tokio::test]
    async fn test_format_detection() {
        let loader = ModelLoader::new().unwrap();

        // Test SafeTensors format
        let result = loader.detect_format("model.safetensors");
        assert_eq!(result.unwrap(), "safetensors");

        // Test PyTorch format
        let result = loader.detect_format("model.pt");
        assert_eq!(result.unwrap(), "pytorch");
    }

    #[tokio::test]
    async fn test_cache_operations() -> Result<()> {
        let cache = ModelCache::new(100)?; // 100MB cache

        let arch = ModelArchitecture {
            arch_type: "test".to_string(),
            num_layers: 1,
            hidden_dim: 10,
            params: HashMap::new(),
        };

        let model = Model::new("test_model".to_string(), arch);

        cache.put("test_path".to_string(), model.clone()).await?;
        let retrieved = cache.get("test_path").await?;

        assert!(retrieved.is_some());
        assert_eq!(retrieved.unwrap().name, model.name);

        Ok(())
    }

    #[test]
    fn test_architecture_inference() {
        let loader = ModelLoader::new().unwrap();

        let mut header = SafeTensorsHeader {
            tensors: HashMap::new(),
        };

        header.tensors.insert(
            "layers.0.attention.weight".to_string(),
            SafeTensorInfo {
                dtype: "F32".to_string(),
                shape: vec![768, 768],
                data_offsets: (0, 768 * 768 * 4),
            },
        );

        header.tensors.insert(
            "layers.11.attention.weight".to_string(),
            SafeTensorInfo {
                dtype: "F32".to_string(),
                shape: vec![768, 768],
                data_offsets: (0, 768 * 768 * 4),
            },
        );

        let arch = loader.infer_architecture(&header).unwrap();

        assert_eq!(arch.arch_type, "transformer");
        assert_eq!(arch.num_layers, 12);
        assert_eq!(arch.hidden_dim, 768);
    }
}