use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};

use obicompactvec::PersistentCompactIntMatrix;
use obikseq::CanonicalKmer;
use obiskio::{UnitigFileWriter, DEFAULT_BLOCK_BITS};

use crate::error::OLMResult;
use crate::layer::{Hit, Layer, LayerData};
use crate::meta::{IndexMode, PartitionMeta};

/// Layered kmer index for a single partition.
///
/// Each layer covers a disjoint kmer set. Queries probe layers in order;
/// the first match wins. Adding a dataset appends a new layer without
/// rebuilding existing ones.
pub struct LayeredMap<D: LayerData = ()> {
    root:   PathBuf,
    meta:   PartitionMeta,
    layers: Vec<Layer<D>>,
}

// ── Common methods ────────────────────────────────────────────────────────────

impl<D: LayerData> LayeredMap<D> {
    /// Open an existing layered index at `root`.
    /// The mode is read once from `PartitionMeta` and applied to all layers.
    pub fn open(root: &Path) -> OLMResult<Self> {
        let meta = PartitionMeta::load(root)?;
        let layers = (0..meta.n_layers)
            .map(|i| Layer::<D>::open(&layer_dir(root, i), &meta.mode))
            .collect::<OLMResult<Vec<_>>>()?;
        Ok(Self { root: root.to_owned(), meta, layers })
    }

    /// Create a new, empty layered index at `root` with the given mode.
    pub fn create(root: &Path, mode: IndexMode) -> OLMResult<Self> {
        fs::create_dir_all(root)?;
        let meta = PartitionMeta::new(mode);
        meta.save(root)?;
        Ok(Self { root: root.to_owned(), meta, layers: Vec::new() })
    }

    pub fn n_layers(&self) -> usize { self.layers.len() }
    pub fn layer(&self, i: usize) -> &Layer<D> { &self.layers[i] }
    pub fn mode(&self) -> &IndexMode { &self.meta.mode }

    /// Query `kmer` across all layers. Returns `(layer_index, Hit)` on match.
    pub fn query(&self, kmer: CanonicalKmer) -> Option<(usize, Hit<D::Item>)> {
        self.layers
            .iter()
            .enumerate()
            .find_map(|(i, layer)| layer.query(kmer).map(|hit| (i, hit)))
    }

    pub fn next_layer_writer(&self) -> OLMResult<UnitigFileWriter> {
        let dir = layer_dir(&self.root, self.layers.len());
        Layer::<D>::unitig_writer(&dir)
    }

    fn append_layer(&mut self) -> OLMResult<()> {
        let i = self.layers.len();
        let dir = layer_dir(&self.root, i);
        self.layers.push(Layer::<D>::open(&dir, &self.meta.mode)?);
        self.meta.n_layers = self.layers.len();
        self.meta.save(&self.root)?;
        Ok(())
    }
}

// ── Mode 1 — set membership ───────────────────────────────────────────────────

impl LayeredMap<()> {
    pub fn push_layer(&mut self) -> OLMResult<usize> {
        let i = self.layers.len();
        let dir = layer_dir(&self.root, i);
        Layer::<()>::build(&dir, DEFAULT_BLOCK_BITS, &self.meta.mode)?;
        self.append_layer()?;
        Ok(i)
    }
}

// ── Mode 2 — count matrix ─────────────────────────────────────────────────────

impl LayeredMap<PersistentCompactIntMatrix> {
    pub fn push_layer(&mut self, count_of: impl Fn(CanonicalKmer) -> u32) -> OLMResult<usize> {
        let i = self.layers.len();
        let dir = layer_dir(&self.root, i);
        Layer::<PersistentCompactIntMatrix>::build(&dir, DEFAULT_BLOCK_BITS, &self.meta.mode, count_of)?;
        self.append_layer()?;
        Ok(i)
    }

    pub fn push_layer_from_map(&mut self, counts: &HashMap<CanonicalKmer, u32>) -> OLMResult<usize> {
        self.push_layer(|kmer| counts.get(&kmer).copied().unwrap_or(0))
    }
}

fn layer_dir(root: &Path, i: usize) -> PathBuf {
    root.join(format!("layer_{i}"))
}

#[cfg(test)]
#[path = "tests/map.rs"]
mod tests;

