use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obilayeredmap::LayeredStore;
use obiskio::{SKError, SKResult};

use crate::partition::KmerPartition;

const INDEX_SUBDIR: &str = "index";

fn probe_n_layers(index_dir: &std::path::Path) -> usize {
    let mut n = 0;
    while index_dir.join(format!("layer_{n}")).exists() { n += 1; }
    n
}

impl KmerPartition {
    /// Open all count matrices for partition `part`, one per layer.
    /// Layers without a `counts/` directory are skipped.
    pub fn count_store(&self, part: usize) -> SKResult<LayeredStore<PersistentCompactIntMatrix>> {
        let index_dir = self.part_dir(part).join(INDEX_SUBDIR);
        if !index_dir.exists() {
            return Ok(LayeredStore::new(vec![]));
        }
        let matrices = (0..probe_n_layers(&index_dir))
            .filter_map(|l| {
                let layer_dir = index_dir.join(format!("layer_{l}"));
                layer_dir.join("counts").exists()
                    .then(|| PersistentCompactIntMatrix::open(&layer_dir).map_err(SKError::Io))
            })
            .collect::<SKResult<Vec<_>>>()?;
        Ok(LayeredStore::new(matrices))
    }

    /// Open all presence matrices for partition `part`, one per layer.
    /// Layers without a `presence/` directory are skipped.
    pub fn presence_store(&self, part: usize) -> SKResult<LayeredStore<PersistentBitMatrix>> {
        let index_dir = self.part_dir(part).join(INDEX_SUBDIR);
        if !index_dir.exists() {
            return Ok(LayeredStore::new(vec![]));
        }
        let matrices = (0..probe_n_layers(&index_dir))
            .filter_map(|l| {
                let layer_dir = index_dir.join(format!("layer_{l}"));
                layer_dir.join("presence").exists()
                    .then(|| PersistentBitMatrix::open(&layer_dir).map_err(SKError::Io))
            })
            .collect::<SKResult<Vec<_>>>()?;
        Ok(LayeredStore::new(matrices))
    }
}
