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

use obicompactvec::{
    PersistentBitMatrix, PersistentBitMatrixBuilder,
    PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
};
use obikseq::CanonicalKmer;
use obiskio::{UnitigFileReader, UnitigFileWriter};

use crate::error::{OLMError, OLMResult};
use crate::meta::IndexMode;
use crate::mphf_layer::MphfLayer;
pub(crate) use crate::mphf_layer::UNITIGS_FILE;

const COUNTS_DIR:   &str = "counts";
const PRESENCE_DIR: &str = "presence";

// ── Trait ─────────────────────────────────────────────────────────────────────

pub trait LayerData: Sized {
    type Item;
    fn open(layer_dir: &Path) -> OLMResult<Self>;
    fn read(&self, slot: usize) -> Self::Item;
}

impl LayerData for () {
    type Item = ();
    fn open(_layer_dir: &Path) -> OLMResult<Self> { Ok(()) }
    fn read(&self, _slot: usize) {}
}

impl LayerData for PersistentCompactIntMatrix {
    type Item = Box<[u32]>;
    fn open(layer_dir: &Path) -> OLMResult<Self> {
        PersistentCompactIntMatrix::open(layer_dir).map_err(OLMError::Io)
    }
    fn read(&self, slot: usize) -> Box<[u32]> { self.row(slot) }
}

impl LayerData for PersistentBitMatrix {
    type Item = Box<[bool]>;
    fn open(layer_dir: &Path) -> OLMResult<Self> {
        PersistentBitMatrix::open(layer_dir).map_err(OLMError::Io)
    }
    fn read(&self, slot: usize) -> Box<[bool]> { self.row(slot) }
}

// ── Structures ────────────────────────────────────────────────────────────────

pub struct Layer<D: LayerData = ()> {
    mphf: MphfLayer,
    data: D,
}

pub struct Hit<T = ()> {
    pub slot: usize,
    pub data: T,
}

// ── Common read path ──────────────────────────────────────────────────────────

impl<D: LayerData> Layer<D> {
    pub fn open(path: &Path, mode: &IndexMode) -> OLMResult<Self> {
        let mphf = MphfLayer::open(path, mode)?;
        let data = D::open(path)?;
        Ok(Self { mphf, data })
    }

    pub fn query(&self, kmer: CanonicalKmer) -> Option<Hit<D::Item>> {
        self.mphf.find(kmer).map(|slot| Hit { slot, data: self.data.read(slot) })
    }

    pub fn n(&self) -> usize { self.mphf.n() }

    pub fn unitig_writer(out_dir: &Path) -> OLMResult<UnitigFileWriter> {
        MphfLayer::unitig_writer(out_dir)
    }

    /// Build `unitigs.bin.idx` and `evidence.bin` from `unitigs.bin` and
    /// `mphf.bin` already present in `layer_dir`.
    /// `block_bits` controls the `.idx` block size (2^block_bits chunks/block).
    pub fn build_exact_evidence(layer_dir: &Path, block_bits: u8) -> OLMResult<usize> {
        MphfLayer::build_exact_evidence(layer_dir, block_bits)
    }

    /// Build `fingerprint.bin` from `unitigs.bin` and `mphf.bin` already
    /// present in `layer_dir`.  `b` — fingerprint bits (1..=64); `z` — Findere
    /// consecutive k-mer parameter (≥1).
    pub fn build_approx_evidence(layer_dir: &Path, b: u8, z: u8) -> OLMResult<usize> {
        MphfLayer::build_approx_evidence(layer_dir, b, z)
    }

}

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

impl Layer<()> {
    pub fn build(out_dir: &Path, block_bits: u8, mode: &IndexMode) -> OLMResult<usize> {
        MphfLayer::build(out_dir, block_bits, mode, &mut |_, _| Ok(()))
    }

    /// Create a presence matrix for a set-membership layer (first merge).
    pub fn init_presence_matrix(layer_dir: &Path, n_kmers: usize) -> OLMResult<()> {
        let presence_dir = layer_dir.join(PRESENCE_DIR);
        fs::create_dir_all(&presence_dir).map_err(OLMError::Io)?;
        let mut mb = PersistentBitMatrixBuilder::new(n_kmers, &presence_dir).map_err(OLMError::Io)?;
        mb.add_col_ones().map_err(OLMError::Io)?.close().map_err(OLMError::Io)?;
        mb.close().map_err(OLMError::Io)
    }
}

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

impl Layer<PersistentCompactIntMatrix> {
    pub fn build(
        out_dir: &Path,
        block_bits: u8,
        mode: &IndexMode,
        count_of: impl Fn(CanonicalKmer) -> u32,
    ) -> OLMResult<usize> {
        let n = UnitigFileReader::open_sequential(&out_dir.join(UNITIGS_FILE))?.n_kmers();
        let counts_dir = out_dir.join(COUNTS_DIR);
        let mut mb = PersistentCompactIntMatrixBuilder::new(n, &counts_dir)
            .map_err(OLMError::Io)?;
        let mut col = mb.add_col().map_err(OLMError::Io)?;
        let n_built = MphfLayer::build(out_dir, block_bits, mode, &mut |slot, kmer| {
            col.set(slot, count_of(kmer));
            Ok(())
        })?;
        col.close().map_err(OLMError::Io)?;
        mb.close().map_err(OLMError::Io)?;
        Ok(n_built)
    }

    pub fn build_from_map(
        out_dir: &Path,
        block_bits: u8,
        mode: &IndexMode,
        counts: &HashMap<CanonicalKmer, u32>,
    ) -> OLMResult<usize> {
        Self::build(out_dir, block_bits, mode, |kmer| counts.get(&kmer).copied().unwrap_or(0))
    }
}

// ── Mode 2 — count matrix column append ──────────────────────────────────────

impl Layer<PersistentCompactIntMatrix> {
    pub fn append_genome_column(
        layer_dir: &Path,
        value_of: impl Fn(usize) -> u32,
    ) -> OLMResult<()> {
        PersistentCompactIntMatrix::append_column(&layer_dir.join(COUNTS_DIR), value_of)
            .map_err(OLMError::Io)
    }
}

// ── Mode 3 — presence/absence matrix ─────────────────────────────────────────

impl Layer<PersistentBitMatrix> {
    pub fn append_genome_column(
        layer_dir: &Path,
        value_of: impl Fn(usize) -> bool,
    ) -> OLMResult<()> {
        PersistentBitMatrix::append_column(&layer_dir.join(PRESENCE_DIR), value_of)
            .map_err(OLMError::Io)
    }

    pub fn build_presence(
        out_dir: &Path,
        block_bits: u8,
        mode: &IndexMode,
        n_genomes: usize,
        present_in: impl Fn(CanonicalKmer, usize) -> bool,
    ) -> OLMResult<usize> {
        let n = UnitigFileReader::open_sequential(&out_dir.join(UNITIGS_FILE))?.n_kmers();
        let presence_dir = out_dir.join(PRESENCE_DIR);
        let mut mb = PersistentBitMatrixBuilder::new(n, &presence_dir).map_err(OLMError::Io)?;
        let mut cols: Vec<_> = (0..n_genomes)
            .map(|_| mb.add_col().map_err(OLMError::Io))
            .collect::<OLMResult<_>>()?;
        let n_built = MphfLayer::build(out_dir, block_bits, mode, &mut |slot, kmer| {
            for (g, col) in cols.iter_mut().enumerate() {
                col.set(slot, present_in(kmer, g));
            }
            Ok(())
        })?;
        for col in cols {
            col.close().map_err(OLMError::Io)?;
        }
        mb.close().map_err(OLMError::Io)?;
        Ok(n_built)
    }
}

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