use std::path::Path;

use obicompactvec::{
    FilterMask, eval_filter_mask,
    PersistentBitMatrixBuilder, PersistentBitVecBuilder,
    PersistentCompactIntMatrixBuilder, PersistentCompactIntVecBuilder,
};
use obidebruinj::GraphDeBruijn;
use obikseq::CanonicalKmer;
use obilayeredmap::meta::PartitionMeta;
use obilayeredmap::{IndexMode, MphfLayer};
use obiskio::{SKError, SKResult, UnitigFileReader};

use crate::common::{load_meta, olm_to_sk};
use crate::filter::KmerFilter;
use crate::graph_pipeline::materialize_layer;
use crate::merge_layer::{MergeMode, SrcLayerData};
use crate::partition::KmerPartition;

const INDEX_SUBDIR: &str = "index";

// ── Builders — pair matrix builder + column builders for one mode ─────────────

enum Builders {
    Presence(PersistentBitMatrixBuilder, Vec<PersistentBitVecBuilder>),
    Count(PersistentCompactIntMatrixBuilder, Vec<PersistentCompactIntVecBuilder>),
}

impl Builders {
    fn new(mode: MergeMode, n: usize, dir: &Path, n_genomes: usize) -> SKResult<Self> {
        match mode {
            MergeMode::Presence => {
                let mut mat = PersistentBitMatrixBuilder::new(n, dir).map_err(SKError::Io)?;
                let mut cols = Vec::with_capacity(n_genomes);
                for _ in 0..n_genomes { cols.push(mat.add_col().map_err(SKError::Io)?); }
                Ok(Builders::Presence(mat, cols))
            }
            MergeMode::Count => {
                let mut mat = PersistentCompactIntMatrixBuilder::new(n, dir).map_err(SKError::Io)?;
                let mut cols = Vec::with_capacity(n_genomes);
                for _ in 0..n_genomes { cols.push(mat.add_col().map_err(SKError::Io)?); }
                Ok(Builders::Count(mat, cols))
            }
        }
    }

    fn set_val(&mut self, col: usize, slot: usize, value: u32) {
        match self {
            Builders::Presence(_, cols) => cols[col].set(slot, value > 0),
            Builders::Count(_, cols)    => cols[col].set(slot, value),
        }
    }

    fn close(self) -> SKResult<()> {
        match self {
            Builders::Presence(mat, cols) => {
                for b in cols { b.close().map_err(SKError::Io)?; }
                mat.close().map_err(SKError::Io)
            }
            Builders::Count(mat, cols) => {
                for b in cols { b.close().map_err(SKError::Io)?; }
                mat.close().map_err(SKError::Io)
            }
        }
    }
}

// ── try_compute_combined_mask ─────────────────────────────────────────────────

/// Build a per-slot `TempBitVec` mask from `filters` using column operations
/// on the source matrix — no per-kmer MPHF lookup or row read needed.
///
/// Returns `Some(mask)` when every filter in `filters` can express itself as
/// a [`FilterMask`] expression.  Returns `None` when any filter requires
/// row-level inspection (fall back to `passes_all`).
fn try_compute_combined_mask(
    filters: &[Box<dyn KmerFilter>],
    src_data: &SrcLayerData,
    n_genomes: usize,
) -> SKResult<Option<obicompactvec::TempBitVec>> {
    if filters.is_empty() {
        return Ok(None);
    }
    let mut exprs: Vec<FilterMask> = Vec::with_capacity(filters.len());
    for f in filters {
        match f.column_mask_expr(n_genomes) {
            Some(expr) => exprs.push(expr),
            None => return Ok(None),
        }
    }
    let combined = FilterMask::And(exprs);
    let n = src_data.n_slots();
    let mask = src_data
        .with_matrix(|mat| eval_filter_mask(&combined, mat, n))
        .map_err(SKError::Io)?;
    Ok(Some(mask))
}

// ── iter_src_kmers_masked (pass 1) ────────────────────────────────────────────

/// Iterate all passing kmers in `src_index_dir`, yielding only the kmer value.
///
/// When all filters can be expressed as column operations, a per-slot mask is
/// computed once per layer and used for O(1) slot-check per kmer instead of a
/// full row read.  Falls back to row-level `passes_all` otherwise.
fn iter_src_kmers_masked(
    src_index_dir: &Path,
    mode: MergeMode,
    n_genomes: usize,
    filters: &[Box<dyn KmerFilter>],
    mut cb: impl FnMut(CanonicalKmer),
) -> SKResult<()> {
    let src_meta = load_meta(src_index_dir, "rebuild")?;
    for l in 0..src_meta.n_layers {
        let src_layer_dir = src_index_dir.join(format!("layer_{l}"));
        let unitigs_path = src_layer_dir.join("unitigs.bin");
        if !unitigs_path.exists() { continue; }

        let src_data = SrcLayerData::open(&src_layer_dir, mode)?;
        let mask = try_compute_combined_mask(filters, &src_data, n_genomes)?;
        let reader = UnitigFileReader::open_sequential(&unitigs_path)?;

        for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
            let slot = src_data.slot(kmer);
            let passes = match &mask {
                Some(m) => m.get(slot),
                None => {
                    let row = src_data.fill_row_by_slot(slot, n_genomes);
                    filters.iter().all(|f| f.passes(&row, n_genomes))
                }
            };
            if passes { cb(kmer); }
        }
    }
    Ok(())
}

// ── iter_src_layers (pass 2) ──────────────────────────────────────────────────

/// Iterate all passing kmers in `src_index_dir`, yielding `(kmer, row)`.
///
/// When the slot mask is available, skips the row read for filtered-out slots.
fn iter_src_layers(
    src_index_dir: &Path,
    mode: MergeMode,
    n_genomes: usize,
    filters: &[Box<dyn KmerFilter>],
    mut cb: impl FnMut(CanonicalKmer, Box<[u32]>),
) -> SKResult<()> {
    let src_meta = load_meta(src_index_dir, "rebuild")?;
    for l in 0..src_meta.n_layers {
        let src_layer_dir = src_index_dir.join(format!("layer_{l}"));
        let unitigs_path = src_layer_dir.join("unitigs.bin");
        if !unitigs_path.exists() { continue; }

        let src_data = SrcLayerData::open(&src_layer_dir, mode)?;
        let mask = try_compute_combined_mask(filters, &src_data, n_genomes)?;
        let reader = UnitigFileReader::open_sequential(&unitigs_path)?;

        for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
            let slot = src_data.slot(kmer);
            if let Some(ref m) = mask {
                if !m.get(slot) { continue; }
                let row = src_data.fill_row_by_slot(slot, n_genomes);
                cb(kmer, row.into_boxed_slice());
            } else {
                let row = src_data.fill_row_by_slot(slot, n_genomes);
                if filters.iter().all(|f| f.passes(&row, n_genomes)) {
                    cb(kmer, row.into_boxed_slice());
                }
            }
        }
    }
    Ok(())
}

// ── KmerPartition::rebuild_partition ─────────────────────────────────────────

impl KmerPartition {
    /// Rebuild partition `i` from `src` into `self` (an empty destination partition).
    ///
    /// Only k-mers whose per-genome row passes all `filters` are written.
    /// The output is a single-layer index — regardless of how many layers the
    /// source has.
    ///
    /// `n_genomes` is the number of genome columns in the source (and destination).
    pub fn rebuild_partition(
        &self,
        src: &KmerPartition,
        i: usize,
        filters: &[Box<dyn KmerFilter>],
        mode: MergeMode,
        n_genomes: usize,
        block_bits: u8,
    ) -> SKResult<()> {
        let src_index_dir = src.part_dir(i).join(INDEX_SUBDIR);
        if !src_index_dir.exists() {
            return Ok(());
        }

        let src_meta = load_meta(&src_index_dir, "rebuild")?;
        if src_meta.n_layers == 0 {
            return Ok(());
        }

        // ── Pass 1: collect filtered kmers into de Bruijn graph ───────────────
        let mut g = GraphDeBruijn::new();
        iter_src_kmers_masked(&src_index_dir, mode, n_genomes, filters, |kmer| {
            g.push(kmer);
        })?;

        if g.len() == 0 {
            return Ok(());
        }

        // ── Build MPHF in dst layer_0 ─────────────────────────────────────────
        let dst_index_dir = self.part_dir(i).join(INDEX_SUBDIR);
        let dst_layer_dir = dst_index_dir.join("layer_0");

        let n_new = materialize_layer(g, &dst_layer_dir, block_bits, &IndexMode::Exact)?;
        let dst_mphf = MphfLayer::open(&dst_layer_dir, &IndexMode::Exact)
            .map_err(|e| olm_to_sk(e, "rebuild"))?;

        // ── Prepare matrix builders (one column per genome) ───────────────────
        let data_dir = match mode {
            MergeMode::Presence => dst_layer_dir.join("presence"),
            MergeMode::Count    => dst_layer_dir.join("counts"),
        };
        std::fs::create_dir_all(&data_dir)?;
        let mut builders = Builders::new(mode, n_new, &data_dir, n_genomes)?;

        // ── Pass 2: fill builders ─────────────────────────────────────────────
        iter_src_layers(&src_index_dir, mode, n_genomes, filters, |kmer, row| {
            if let Some(slot) = dst_mphf.find(kmer) {
                for (col, &value) in row.iter().enumerate() {
                    builders.set_val(col, slot, value);
                }
            }
        })?;

        // ── Close builders and write metadata ─────────────────────────────────
        builders.close()?;

        PartitionMeta {
            n_layers: 1,
            mode: IndexMode::Exact,
        }
        .save(&dst_index_dir)
        .map_err(|e| olm_to_sk(e, "rebuild"))?;

        Ok(())
    }
}
