use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obikseq::CanonicalKmer;
use obiskio::{SKError, SKResult, UnitigFileReader};
use obilayeredmap::{IndexMode, MphfLayer, OLMError};
use obilayeredmap::meta::PartitionMeta;

use crate::filter::{KmerFilter, passes_all};
use crate::partition::KmerPartition;

const INDEX_SUBDIR: &str = "index";

fn olm_to_sk(e: OLMError) -> SKError {
    match e {
        OLMError::Io(e) => SKError::Io(e),
        other => SKError::InvalidData { context: "dump", detail: other.to_string() },
    }
}

impl KmerPartition {
    /// Iterate all indexed kmers in partition `part`, calling `cb(kmer, row)` for each
    /// kmer that passes every filter in `filters`.
    ///
    /// `use_counts = true`  → reads count columns (u32 values per genome).
    /// `use_counts = false` → reads presence columns, converted to 0/1 u32.
    ///
    /// If no data matrix exists for a layer (pure set-membership, single genome),
    /// a row of `n_genomes` ones is emitted for every kmer in that layer — unless
    /// the filter rejects it, in which case the whole layer is skipped.
    /// Like [`iter_partition_kmers`] but the callback returns `false` to stop early.
    /// Returns `Ok(true)` if all kmers were visited, `Ok(false)` if the callback halted.
    pub fn iter_partition_kmers(
        &self,
        part: usize,
        use_counts: bool,
        n_genomes: usize,
        filters: &[Box<dyn KmerFilter>],
        mut cb: impl FnMut(CanonicalKmer, Box<[u32]>) -> bool,
    ) -> SKResult<bool> {
        let index_dir = self.part_dir(part).join(INDEX_SUBDIR);
        if !index_dir.exists() {
            return Ok(true);
        }

        let index_mode = PartitionMeta::load(&index_dir)
            .map(|m| m.mode)
            .unwrap_or(IndexMode::Exact);

        let mut l = 0;
        loop {
            let layer_dir = index_dir.join(format!("layer_{l}"));
            if !layer_dir.exists() { break; }
            l += 1;
            let mphf   = MphfLayer::open(&layer_dir, &index_mode).map_err(olm_to_sk)?;
            let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))?;

            let counts_dir   = layer_dir.join("counts");
            let presence_dir = layer_dir.join("presence");

            let cont = if use_counts && counts_dir.exists() {
                let mat = PersistentCompactIntMatrix::open(&layer_dir).map_err(SKError::Io)?;
                let mut cont = true;
                for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
                    if let Some(slot) = mphf.find(kmer) {
                        let row = mat.row(slot);
                        if passes_all(filters, &row, n_genomes) {
                            cont = cb(kmer, row);
                            if !cont { break; }
                        }
                    }
                }
                cont
            } else if !use_counts && presence_dir.exists() {
                let mat = PersistentBitMatrix::open(&layer_dir).map_err(SKError::Io)?;
                let mut cont = true;
                for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
                    if let Some(slot) = mphf.find(kmer) {
                        let row: Box<[u32]> = mat.row(slot).iter().map(|&b| b as u32).collect();
                        if passes_all(filters, &row, n_genomes) {
                            cont = cb(kmer, row);
                            if !cont { break; }
                        }
                    }
                }
                cont
            } else {
                // No data matrix: implicit presence — all values are 1.
                // The filter result is identical for every kmer, so evaluate once.
                let all_present: Box<[u32]> = vec![1u32; n_genomes].into();
                let mut cont = true;
                if passes_all(filters, &all_present, n_genomes) {
                    for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
                        if mphf.find(kmer).is_some() {
                            cont = cb(kmer, all_present.clone());
                            if !cont { break; }
                        }
                    }
                }
                cont
            };

            if !cont { return Ok(false); }
        }

        Ok(true)
    }

    /// Like [`iter_partition_kmers`] but the callback also receives `(partition, layer)`
    /// indices, enabling debug output that identifies where each kmer was stored.
    /// Returns `Ok(true)` if all kmers were visited, `Ok(false)` if the callback halted.
    pub fn iter_partition_kmers_located(
        &self,
        part: usize,
        use_counts: bool,
        n_genomes: usize,
        filters: &[Box<dyn KmerFilter>],
        mut cb: impl FnMut(usize, usize, CanonicalKmer, Box<[u32]>) -> bool,
    ) -> SKResult<bool> {
        let index_dir = self.part_dir(part).join(INDEX_SUBDIR);
        if !index_dir.exists() {
            return Ok(true);
        }

        let index_mode = PartitionMeta::load(&index_dir)
            .map(|m| m.mode)
            .unwrap_or(IndexMode::Exact);

        let mut layer = 0;
        loop {
            let layer_dir = index_dir.join(format!("layer_{layer}"));
            if !layer_dir.exists() { break; }
            let mphf   = MphfLayer::open(&layer_dir, &index_mode).map_err(olm_to_sk)?;
            let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))?;

            let counts_dir   = layer_dir.join("counts");
            let presence_dir = layer_dir.join("presence");

            let cont = if use_counts && counts_dir.exists() {
                let mat = PersistentCompactIntMatrix::open(&layer_dir).map_err(SKError::Io)?;
                let mut cont = true;
                for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
                    if let Some(slot) = mphf.find(kmer) {
                        let row = mat.row(slot);
                        if passes_all(filters, &row, n_genomes) {
                            cont = cb(part, layer, kmer, row);
                            if !cont { break; }
                        }
                    }
                }
                cont
            } else if !use_counts && presence_dir.exists() {
                let mat = PersistentBitMatrix::open(&layer_dir).map_err(SKError::Io)?;
                let mut cont = true;
                for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
                    if let Some(slot) = mphf.find(kmer) {
                        let row: Box<[u32]> = mat.row(slot).iter().map(|&b| b as u32).collect();
                        if passes_all(filters, &row, n_genomes) {
                            cont = cb(part, layer, kmer, row);
                            if !cont { break; }
                        }
                    }
                }
                cont
            } else {
                let all_present: Box<[u32]> = vec![1u32; n_genomes].into();
                let mut cont = true;
                if passes_all(filters, &all_present, n_genomes) {
                    for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
                        if mphf.find(kmer).is_some() {
                            cont = cb(part, layer, kmer, all_present.clone());
                            if !cont { break; }
                        }
                    }
                }
                cont
            };

            if !cont { return Ok(false); }
            layer += 1;
        }

        Ok(true)
    }
}
