use half::{bf16, f16};
use safetensors::SafeTensors;
use std::collections::HashMap;
use std::path::Path;
use xserv_tensor::{DType, Device, Tensor};

pub fn load_safetensors(path: &Path, device: Device) -> HashMap<String, Tensor> {
    let data = std::fs::read(path)
        .unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()));
    let st = SafeTensors::deserialize(&data)
        .unwrap_or_else(|e| panic!("failed to parse safetensors {}: {e}", path.display()));

    let mut tensors = HashMap::new();

    for (name, view) in st.tensors() {
        let shape: Vec<usize> = view.shape().to_vec();
        let raw_bytes = view.data();
        let dtype = match view.dtype() {
            safetensors::Dtype::F32 => DType::F32,
            safetensors::Dtype::F16 => DType::F16,
            safetensors::Dtype::BF16 => DType::BF16,
            other => {
                eprintln!("skipping tensor {name}: unsupported dtype {other:?}");
                continue;
            }
        };

        let tensor = make_tensor(raw_bytes, &shape, dtype);
        let tensor = tensor.to_device(device);
        tensors.insert(name.to_string(), tensor);
    }

    tensors
}

/// Load from a directory containing model.safetensors (or sharded files) + config.json.
pub fn load_model_dir(dir: &Path, device: Device) -> HashMap<String, Tensor> {
    let single = dir.join("model.safetensors");
    if single.exists() {
        return load_safetensors(&single, device);
    }

    // Try sharded: model-00001-of-NNNNN.safetensors
    let mut all_tensors = HashMap::new();
    let mut entries: Vec<_> = std::fs::read_dir(dir)
        .unwrap()
        .filter_map(|e| e.ok())
        .filter(|e| {
            e.path()
                .file_name()
                .map(|f| f.to_string_lossy().ends_with(".safetensors"))
                .unwrap_or(false)
        })
        .collect();
    entries.sort_by_key(|e| e.file_name());

    for entry in entries {
        let tensors = load_safetensors(&entry.path(), device);
        all_tensors.extend(tensors);
    }

    assert!(!all_tensors.is_empty(), "no safetensors files found in {}", dir.display());
    all_tensors
}

fn make_tensor(raw_bytes: &[u8], shape: &[usize], dtype: DType) -> Tensor {
    match dtype {
        DType::F32 => {
            let floats: &[f32] = unsafe {
                std::slice::from_raw_parts(raw_bytes.as_ptr() as *const f32, raw_bytes.len() / 4)
            };
            Tensor::from_slice(floats, shape)
        }
        DType::F16 => {
            let halfs: &[f16] = unsafe {
                std::slice::from_raw_parts(raw_bytes.as_ptr() as *const f16, raw_bytes.len() / 2)
            };
            Tensor::from_slice(halfs, shape)
        }
        DType::BF16 => {
            let bfs: &[bf16] = unsafe {
                std::slice::from_raw_parts(raw_bytes.as_ptr() as *const bf16, raw_bytes.len() / 2)
            };
            Tensor::from_slice(bfs, shape)
        }
    }
}