use std::path::Path;

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

use crate::partition::KmerPartition;

const INDEX_SUBDIR: &str = "index";

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

// ── per-layer query handle ────────────────────────────────────────────────────

enum QueryLayer {
    Presence(MphfLayer, PersistentBitMatrix),
    Count(MphfLayer, PersistentCompactIntMatrix),
}

impl QueryLayer {
    fn open(layer_dir: &Path, with_counts: bool, mode: &IndexMode) -> SKResult<Self> {
        let mphf        = MphfLayer::open(layer_dir, mode).map_err(olm_to_sk)?;
        let counts_dir  = layer_dir.join("counts");
        let presence_dir = layer_dir.join("presence");

        if with_counts && counts_dir.exists() {
            let mat = PersistentCompactIntMatrix::open(layer_dir).map_err(SKError::Io)?;
            Ok(QueryLayer::Count(mphf, mat))
        } else if presence_dir.exists() || !counts_dir.exists() {
            // presence mode, or no matrix at all → Implicit handled inside open()
            let mat = PersistentBitMatrix::open(layer_dir).map_err(SKError::Io)?;
            Ok(QueryLayer::Presence(mphf, mat))
        } else {
            // counts exist but not presence — count layer, no presence requested
            let mat = PersistentCompactIntMatrix::open(layer_dir).map_err(SKError::Io)?;
            Ok(QueryLayer::Count(mphf, mat))
        }
    }

    /// Write per-genome values into `buf` if `kmer` is indexed; returns true on hit.
    fn find_into(&self, kmer: CanonicalKmer, n_genomes: usize, buf: &mut [u32]) -> bool {
        match self {
            QueryLayer::Presence(mphf, mat) => {
                if let Some(slot) = mphf.find(kmer) {
                    mat.fill_row(slot, &mut buf[..n_genomes]);
                    true
                } else {
                    false
                }
            }
            QueryLayer::Count(mphf, mat) => {
                if let Some(slot) = mphf.find(kmer) {
                    mat.fill_row(slot, &mut buf[..n_genomes]);
                    true
                } else {
                    false
                }
            }
        }
    }
}

// ── KmerPartition::query_partition* ──────────────────────────────────────────

impl KmerPartition {
    /// Query a single partition, calling `on_hit(sk_idx, kmer_idx, row)` for
    /// every found k-mer without allocating intermediate result vectors.
    pub fn query_partition_with<F>(
        &self,
        part_idx: usize,
        superkmers: &[&RoutableSuperKmer],
        _k: usize,
        n_genomes: usize,
        with_counts: bool,
        mut on_hit: F,
    ) -> SKResult<()>
    where
        F: FnMut(usize, usize, &[u32]),
    {
        if superkmers.is_empty() {
            return Ok(());
        }

        let index_dir = self.part_dir(part_idx).join(INDEX_SUBDIR);
        if !index_dir.exists() {
            return Ok(());
        }

        let meta = PartitionMeta::load(&index_dir).map_err(olm_to_sk)?;
        let layers: Vec<QueryLayer> = (0..meta.n_layers)
            .map(|i| QueryLayer::open(&index_dir.join(format!("layer_{i}")), with_counts, &meta.mode))
            .collect::<SKResult<_>>()?;

        let mut buf = vec![0u32; n_genomes];

        for (sk_idx, rsk) in superkmers.iter().enumerate() {
            for (kmer_idx, kmer) in rsk.superkmer().iter_canonical_kmers().enumerate() {
                for layer in &layers {
                    if layer.find_into(kmer, n_genomes, &mut buf) {
                        on_hit(sk_idx, kmer_idx, &buf);
                        buf.fill(0);
                        break;
                    }
                }
            }
        }

        Ok(())
    }

    /// Query a single partition for a slice of super-kmers, returning per-kmer rows.
    /// Prefer [`query_partition_with`] to avoid per-kmer heap allocations.
    pub fn query_partition(
        &self,
        part_idx: usize,
        superkmers: &[&RoutableSuperKmer],
        _k: usize,
        n_genomes: usize,
        with_counts: bool,
    ) -> SKResult<Vec<Vec<Option<Box<[u32]>>>>> {
        if superkmers.is_empty() {
            return Ok(Vec::new());
        }

        let index_dir = self.part_dir(part_idx).join(INDEX_SUBDIR);

        if !index_dir.exists() {
            return Ok(superkmers
                .iter()
                .map(|rsk| vec![None; rsk.seql()])
                .collect());
        }

        let meta = PartitionMeta::load(&index_dir).map_err(olm_to_sk)?;
        let layers: Vec<QueryLayer> = (0..meta.n_layers)
            .map(|i| QueryLayer::open(&index_dir.join(format!("layer_{i}")), with_counts, &meta.mode))
            .collect::<SKResult<_>>()?;

        let mut buf = vec![0u32; n_genomes];
        Ok(superkmers
            .iter()
            .map(|rsk| {
                rsk.superkmer()
                    .iter_canonical_kmers()
                    .map(|kmer| {
                        for layer in &layers {
                            if layer.find_into(kmer, n_genomes, &mut buf) {
                                let row: Box<[u32]> = buf[..n_genomes].into();
                                buf.fill(0);
                                return Some(row);
                            }
                        }
                        None
                    })
                    .collect()
            })
            .collect())
    }
}
