use crate::tensor::{Tensor, Shape, DType};
use crate::hardware::device::DeviceType;
use safetensors::SafeTensors;
use memmap2::MmapOptions;
use std::fs::File;
use std::path::Path;
use std::collections::HashMap;

pub struct ModelLoader;

impl ModelLoader {
    pub fn load_safetensors(path: &str) -> Result<HashMap<String, Tensor>, String> {
        let path = Path::new(path);
        let file = File::open(path).map_err(|e| e.to_string())?;
        
        // Memory map the file for zero-copy reading
        let mmap = unsafe { MmapOptions::new().map(&file).map_err(|e| e.to_string())? };
        
        let tensors = SafeTensors::deserialize(&mmap).map_err(|e| e.to_string())?;
        let mut loaded_tensors = HashMap::new();

        for (name, view) in tensors.tensors() {
            let shape = Shape::new(view.shape().to_vec());
            
            let dtype = match view.dtype() {
                safetensors::Dtype::F32 => DType::F32,
                // TODO: Add support for other types
                _ => return Err(format!("Unsupported dtype in tensor {}", name)),
            };

            // Create a new Ionian Tensor
            let mut tensor = Tensor::new(shape, dtype, DeviceType::Cpu)?;
            
            // Copy data from SafeTensor view to our buffer
            // Note: This involves a copy. Future optimization could use the mmap directly if aligned.
            let data_bytes = view.data();
            
            // We need to cast bytes to F32. 
            // In a real implementation, we should handle endianness, but assuming Little Endian (standard).
            let data_f32: Vec<f32> = data_bytes
                .chunks_exact(4)
                .map(|b| f32::from_le_bytes(b.try_into().unwrap()))
                .collect();
                
            tensor.write_data(&data_f32)?;
            loaded_tensors.insert(name, tensor);
        }

        Ok(loaded_tensors)
    }
}
