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
}

/// Load a model dir splitting tensors by dtype: float tensors (F32/F16/BF16)
/// become `Tensor`s on `device`; U8 tensors (gpt-oss MXFP4 `_blocks`/`_scales`,
/// which are not an xserv Tensor dtype) are returned as raw `(bytes, shape)`.
/// One pass over the shards (the 13GB MXFP4 file is read once).
pub fn load_model_dir_split(
    dir: &Path, device: Device,
) -> (HashMap<String, Tensor>, HashMap<String, (Vec<u8>, Vec<usize>)>) {
    let mut files: Vec<std::path::PathBuf> = Vec::new();
    let single = dir.join("model.safetensors");
    if single.exists() {
        files.push(single);
    } else {
        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());
        files.extend(entries.into_iter().map(|e| e.path()));
    }
    assert!(!files.is_empty(), "no safetensors files in {}", dir.display());

    let mut floats = HashMap::new();
    let mut u8s = HashMap::new();
    for path in &files {
        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 {}: {e}", path.display()));
        for (name, view) in st.tensors() {
            let shape: Vec<usize> = view.shape().to_vec();
            let raw = view.data();
            match view.dtype() {
                safetensors::Dtype::F32 => { floats.insert(name, make_tensor(raw, &shape, DType::F32).to_device(device)); }
                safetensors::Dtype::F16 => { floats.insert(name, make_tensor(raw, &shape, DType::F16).to_device(device)); }
                safetensors::Dtype::BF16 => { floats.insert(name, make_tensor(raw, &shape, DType::BF16).to_device(device)); }
                safetensors::Dtype::U8 => { u8s.insert(name, (raw.to_vec(), shape)); }
                other => eprintln!("load_model_dir_split: skipping {name}: dtype {other:?}"),
            }
        }
    }
    (floats, u8s)
}

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)
        }
    }
}