use std::collections::HashMap;
use std::path::Path;

use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obikseq::CanonicalKmer;
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))
        }
    }

    /// MPHF lookup only — no matrix access. `Some(slot)` on hit.
    fn find_slot(&self, kmer: CanonicalKmer) -> Option<usize> {
        match self {
            QueryLayer::Presence(mphf, _) | QueryLayer::Count(mphf, _) => mphf.find(kmer),
        }
    }

    /// Number of genome columns this layer's matrix actually has. Bounds
    /// column-major iteration — usually equal to the index's `n_genomes`, but
    /// `PersistentBitMatrix::Implicit` (the documented mono-genome fast path)
    /// always reports exactly `1`, regardless of the index's real genome
    /// count, so callers must use this rather than assuming `n_genomes`.
    fn n_cols(&self) -> usize {
        match self {
            QueryLayer::Presence(_, mat) => mat.n_cols(),
            QueryLayer::Count(_, mat) => mat.n_cols(),
        }
    }

    /// Column-major point lookup: value for genome column `g` at `slot`.
    /// `g` must be `< self.n_cols()`; `slot` must come from [`find_slot`] on
    /// this same layer.
    fn col_value(&self, g: usize, slot: usize) -> u32 {
        match self {
            QueryLayer::Presence(_, mat) => mat.get(g, slot),
            QueryLayer::Count(_, mat) => mat.col_view(g).get(slot),
        }
    }
}

// ── KmerDesc — one occurrence of a k-mer in the query batch ──────────────────

/// Describes one occurrence of a (deduplicated) k-mer in the query batch:
/// which sequence it came from, and its absolute s-mer position within it.
#[derive(Debug, Clone, Copy)]
pub struct KmerDesc {
    pub seq_idx: u32,
    pub pos: u32,
}

/// Aggregate counters for one `query_partition_with` call — feeds the
/// dedup-ratio and column-scan logging in `obikmer::cmd::query` (occurrences
/// vs. unique k-mers is the whole justification for k-mer-level
/// dereplication; columns scanned / `get()` calls quantify the column-major
/// fetch's locality claim).
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub struct QueryStats {
    /// Distinct canonical k-mers queried in this partition.
    pub n_unique_kmers: usize,
    /// Total `MphfLayer::find` calls issued (a k-mer tried against more than
    /// one layer before a hit, or against all layers on a miss, counts once
    /// per layer attempted).
    pub n_mphf_calls: usize,
    /// Distinct canonical k-mers that matched some layer.
    pub n_hits: usize,
    /// Total genome columns scanned across all hit layers (sum of
    /// `layer.n_cols()` over layers with at least one hit).
    pub n_columns_scanned: usize,
    /// Total `col_value` calls issued during the column-major fetch pass
    /// (`n_columns_scanned` × hits-per-layer, summed over layers).
    pub n_col_get_calls: usize,
}

impl std::ops::AddAssign for QueryStats {
    fn add_assign(&mut self, other: Self) {
        self.n_unique_kmers += other.n_unique_kmers;
        self.n_mphf_calls += other.n_mphf_calls;
        self.n_hits += other.n_hits;
        self.n_columns_scanned += other.n_columns_scanned;
        self.n_col_get_calls += other.n_col_get_calls;
    }
}

// ── QueryHit — one event delivered to query_partition_with's callback ───────

/// One event from [`KmerPartition::query_partition_with`]'s two-stage query:
/// a `Found` event once per hit k-mer (stage 1, MPHF-only — mark the k-mer as
/// indexed regardless of any genome's value), then a `Value` event per
/// `(hit k-mer, genome)` pair with a nonzero matrix value (stage 2,
/// column-major fetch). Carried as one enum, not two separate callbacks, so
/// the caller only needs one `FnMut` closure — passing two closures that each
/// need to mutably borrow the same accumulator does not borrow-check.
pub enum QueryHit<'a> {
    Found(&'a [KmerDesc]),
    Value(&'a [KmerDesc], usize, u32),
}

// ── KmerPartition::query_partition_with ──────────────────────────────────────

impl KmerPartition {
    /// Query a single partition for a pre-deduplicated map of canonical
    /// k-mers → their occurrences (`seq_idx`, `pos`) in the query batch.
    ///
    /// Two stages:
    /// 1. **MPHF-only pass**: for each unique k-mer, try each layer's MPHF in
    ///    turn (stopping at the first hit) and bucket confirmed hits by
    ///    `(layer, slot)`. Emits one `QueryHit::Found` per hit k-mer. This
    ///    stage's cost is independent of the index's genome count.
    /// 2. **Column-major fetch**: for each layer with at least one hit, walk
    ///    its matrix **column by column** (genome by genome) — for each
    ///    genome, scan the slots bucketed in stage 1 and look up their value.
    ///    Emits one `QueryHit::Value` per nonzero `(k-mer, genome)` pair.
    ///    Total lookups are the same as a row-major pass (`n_hits × n_cols`
    ///    in the worst case); the win is memory locality — both persistent
    ///    matrix formats are column-oriented on disk (one `mmap`'d region per
    ///    genome), so scanning one column at a time touches far fewer
    ///    distinct mmap regions than fetching one full row per hit.
    pub fn query_partition_with<F>(
        &self,
        part_idx: usize,
        kmers: &HashMap<CanonicalKmer, Vec<KmerDesc>>,
        n_genomes: usize,
        with_counts: bool,
        mut on_event: F,
    ) -> SKResult<QueryStats>
    where
        F: FnMut(QueryHit),
    {
        let mut stats = QueryStats::default();

        if kmers.is_empty() {
            return Ok(stats);
        }

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

        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<_>>()?;

        // ── Stage 1: MPHF-only pass, bucket hits by (layer_idx, slot) ────────
        let mut by_layer: Vec<HashMap<usize, &Vec<KmerDesc>>> =
            (0..layers.len()).map(|_| HashMap::new()).collect();

        for (kmer, descs) in kmers {
            stats.n_unique_kmers += 1;
            for (layer_idx, layer) in layers.iter().enumerate() {
                stats.n_mphf_calls += 1;
                if let Some(slot) = layer.find_slot(*kmer) {
                    by_layer[layer_idx].insert(slot, descs);
                    on_event(QueryHit::Found(descs));
                    stats.n_hits += 1;
                    break;
                }
            }
        }

        // ── Stage 2: column-major fetch, per layer ───────────────────────────
        for (layer_idx, slots) in by_layer.iter().enumerate() {
            if slots.is_empty() {
                continue;
            }
            let layer = &layers[layer_idx];
            let n_cols = layer.n_cols().min(n_genomes);
            stats.n_columns_scanned += n_cols;

            for g in 0..n_cols {
                for (&slot, descs) in slots {
                    stats.n_col_get_calls += 1;
                    let v = layer.col_value(g, slot);
                    if v != 0 {
                        on_event(QueryHit::Value(descs, g, v));
                    }
                }
            }
        }

        Ok(stats)
    }
}

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