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

pub fn load_safetensors(path: &Path, device: Device) -> HashMap<String, Tensor> {
    load_safetensors_filtered(path, device, |_| true)
}

/// Load only tensors accepted by `keep`. The safetensors file is still read as
/// one byte buffer, but unneeded tensor payloads are not copied into Tensors.
pub fn load_safetensors_filtered(
    path: &Path,
    device: Device,
    keep: impl Fn(&str) -> bool,
) -> 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() {
        if !keep(&name) {
            continue;
        }
        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,
            safetensors::Dtype::F8_E4M3 => DType::FP8E4M3,
            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> {
    load_model_dir_filtered(dir, device, |_| true)
}

/// Load a filtered subset of a model directory. For indexed sharded models,
/// consult `model.safetensors.index.json` first so shards containing no wanted
/// tensors are never read. This is critical for pipeline parallelism on large
/// models: each stage should read only its layer range, not the full checkpoint.
pub fn load_model_dir_filtered(
    dir: &Path,
    device: Device,
    keep: impl Fn(&str) -> bool,
) -> HashMap<String, Tensor> {
    let single = dir.join("model.safetensors");
    if single.exists() {
        return load_safetensors_filtered(&single, device, keep);
    }

    let index_path = dir.join("model.safetensors.index.json");
    let wanted_shards: Option<HashSet<String>> = if index_path.exists() {
        let text = std::fs::read_to_string(&index_path)
            .unwrap_or_else(|e| panic!("failed to read {}: {e}", index_path.display()));
        let index: serde_json::Value = serde_json::from_str(&text)
            .unwrap_or_else(|e| panic!("failed to parse {}: {e}", index_path.display()));
        let map = index
            .get("weight_map")
            .and_then(|v| v.as_object())
            .unwrap_or_else(|| panic!("{} has no weight_map", index_path.display()));
        Some(
            map.iter()
                .filter(|(name, _)| keep(name))
                .filter_map(|(_, shard)| shard.as_str().map(str::to_string))
                .collect(),
        )
    } else {
        None
    };

    // 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| {
            let is_safetensors = e
                .path()
                .file_name()
                .map(|f| f.to_string_lossy().ends_with(".safetensors"))
                .unwrap_or(false);
            let is_wanted = wanted_shards.as_ref().is_none_or(|wanted| {
                wanted.contains(&e.file_name().to_string_lossy().to_string())
            });
            is_safetensors && is_wanted
        })
        .collect();
    entries.sort_by_key(|e| e.file_name());

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

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

pub(crate) 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)
        }
        DType::FP8E4M3 => Tensor::from_raw_bytes(raw_bytes, shape, DType::FP8E4M3),
    }
}