use std::fs;
use std::io;
use std::path::{Path, PathBuf};

use obidebruinj::GraphDeBruijn;
use obicompactvec::{PersistentBitMatrix, PersistentBitMatrixBuilder,
                    PersistentBitVecBuilder,
                    PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
                    PersistentCompactIntVecBuilder};
use obikseq::CanonicalKmer;
use obiskio::{SKError, SKResult, UnitigFileReader};
use obilayeredmap::{IndexMode, Layer, LayeredMap, MphfOnly, OLMError};
use obilayeredmap::meta::PartitionMeta;

use crate::partition::KmerPartition;

// ── MergeMode ─────────────────────────────────────────────────────────────────

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MergeMode { Presence, Count }

// ── ColBuilder — enum dispatch to avoid trait-object boxing issues ─────────────

enum ColBuilder {
    Bit(PersistentBitVecBuilder),
    Int(PersistentCompactIntVecBuilder),
}

impl ColBuilder {
    fn set_val(&mut self, slot: usize, value: u32) {
        match self {
            ColBuilder::Bit(b) => b.set(slot, value > 0),
            ColBuilder::Int(b) => b.set(slot, value),
        }
    }

    fn close(self) -> SKResult<()> {
        match self {
            ColBuilder::Bit(b)  => b.close().map_err(SKError::Io),
            ColBuilder::Int(b)  => b.close().map_err(SKError::Io),
        }
    }
}

// ── SrcLayerData — opened source matrix for pass-2 lookup ─────────────────────

pub(crate) enum SrcLayerData {
    /// Pure set-membership layer (no data matrix): every kmer is present in all genomes.
    SetMembership,
    Presence(MphfOnly, PersistentBitMatrix),
    Count(MphfOnly, PersistentCompactIntMatrix),
}

impl SrcLayerData {
    pub(crate) fn open(layer_dir: &Path, merge_mode: MergeMode) -> SKResult<Self> {
        let presence_dir = layer_dir.join("presence");
        let counts_dir   = layer_dir.join("counts");
        match merge_mode {
            MergeMode::Presence => {
                if presence_dir.exists() {
                    let mphf = MphfOnly::open(layer_dir).map_err(olm_to_sk)?;
                    let mat  = PersistentBitMatrix::open(&presence_dir).map_err(SKError::Io)?;
                    Ok(SrcLayerData::Presence(mphf, mat))
                } else if counts_dir.exists() {
                    let mphf = MphfOnly::open(layer_dir).map_err(olm_to_sk)?;
                    let mat  = PersistentCompactIntMatrix::open(&counts_dir).map_err(SKError::Io)?;
                    Ok(SrcLayerData::Count(mphf, mat))
                } else {
                    Ok(SrcLayerData::SetMembership)
                }
            }
            MergeMode::Count => {
                if counts_dir.exists() {
                    let mphf = MphfOnly::open(layer_dir).map_err(olm_to_sk)?;
                    let mat  = PersistentCompactIntMatrix::open(&counts_dir).map_err(SKError::Io)?;
                    Ok(SrcLayerData::Count(mphf, mat))
                } else {
                    Ok(SrcLayerData::SetMembership)
                }
            }
        }
    }

    /// Return one value per source genome for `kmer`.
    /// The caller guarantees `kmer` is in the source MPHF domain.
    #[inline]
    pub(crate) fn lookup(&self, kmer: CanonicalKmer, n_genomes: usize) -> Vec<u32> {
        match self {
            SrcLayerData::SetMembership => vec![1u32; n_genomes],
            SrcLayerData::Presence(mphf, mat) => {
                mat.row(mphf.index(kmer)).iter().map(|&b| b as u32).collect()
            }
            SrcLayerData::Count(mphf, mat) => {
                mat.row(mphf.index(kmer)).iter().copied().collect()
            }
        }
    }
}

// ── helpers ───────────────────────────────────────────────────────────────────

const INDEX_SUBDIR: &str = "index";

/// Load PartitionMeta, or recover it by probing layer directories.
/// Indexes built before meta.json was introduced lack the file.
fn load_meta(dir: &Path) -> SKResult<PartitionMeta> {
    match PartitionMeta::load(dir) {
        Ok(m) => Ok(m),
        Err(e) if matches!(e, OLMError::Io(ref io_e) if io_e.kind() == std::io::ErrorKind::NotFound) => {
            let mut n = 0usize;
            while dir.join(format!("layer_{n}")).exists() { n += 1; }
            let m = PartitionMeta { n_layers: n, mode: IndexMode::default() };
            m.save(dir).map_err(olm_to_sk)?;
            Ok(m)
        }
        Err(e) => Err(olm_to_sk(e)),
    }
}

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

fn col_path_bit(dir: &Path, col: usize) -> PathBuf {
    dir.join(format!("col_{col:06}.pbiv"))
}

fn col_path_int(dir: &Path, col: usize) -> PathBuf {
    dir.join(format!("col_{col:06}.pciv"))
}

fn write_matrix_meta(dir: &Path, n: usize, n_cols: usize) -> io::Result<()> {
    fs::write(dir.join("meta.json"), format!("{{\"n\":{n},\"n_cols\":{n_cols}}}\n"))
}

// ── KmerPartition::merge_partition ────────────────────────────────────────────

impl KmerPartition {
    /// Merge `sources` into destination partition `i`.
    ///
    /// Each entry in `sources` is `(partition, n_genomes)` where `n_genomes` is
    /// the number of genome columns that source contributes. A merged index
    /// contributes more than one. The total new columns added to the destination
    /// is `sum(n_genomes)`.
    ///
    /// `n_dst_genomes` is the number of genome columns already in the destination
    /// matrices (copied from source_0 before this call).
    pub fn merge_partition(
        &self,
        i: usize,
        sources: &[(&KmerPartition, usize)],
        mode: MergeMode,
        n_dst_genomes: usize,
        block_bits: u8,
        evidence: &IndexMode,
    ) -> SKResult<()> {
        let dst_index_dir = self.part_dir(i).join(INDEX_SUBDIR);
        if !dst_index_dir.exists() {
            return Ok(());
        }

        load_meta(&dst_index_dir)?;  // ensure meta.json exists before LayeredMap::open
        let dst_map = LayeredMap::<()>::open(&dst_index_dir).map_err(olm_to_sk)?;
        let n_dst_layers = dst_map.n_layers();
        let n_src_total: usize = sources.iter().map(|(_, n)| *n).sum();

        // First merge in presence mode: init presence matrices on existing layers
        // (all slots true — every kmer in those layers belongs to genome_0).
        if n_dst_genomes == 1 && mode == MergeMode::Presence {
            for l in 0..n_dst_layers {
                let layer_dir = dst_index_dir.join(format!("layer_{l}"));
                Layer::<()>::init_presence_matrix(&layer_dir, dst_map.layer(l).n())
                    .map_err(olm_to_sk)?;
            }
        }

        // ── Pass 1: classify kmers, build new-kmer de Bruijn graph ───────────
        let mut g = GraphDeBruijn::new();
        let mut any_new = false;

        for (src, _) in sources.iter() {
            let src_index_dir = src.part_dir(i).join(INDEX_SUBDIR);
            if !src_index_dir.exists() { continue; }
            let src_meta = load_meta(&src_index_dir)?;

            for l in 0..src_meta.n_layers {
                let unitigs_path = src_index_dir
                    .join(format!("layer_{l}")).join("unitigs.bin");
                let reader = UnitigFileReader::open_sequential(&unitigs_path)?;
                for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
                    if dst_map.query(kmer).is_none() {
                        g.push(kmer);
                        any_new = true;
                    }
                }
            }
        }

        // Build new layer from de Bruijn graph if there are new kmers.
        let new_layer_idx = n_dst_layers;
        let new_layer_dir = dst_index_dir.join(format!("layer_{new_layer_idx}"));

        let n_new = if any_new {
            g.compute_degrees();
            fs::create_dir_all(&new_layer_dir)?;
            let mut uw = Layer::<()>::unitig_writer(&new_layer_dir).map_err(olm_to_sk)?;
            for unitig in g.iter_unitig() {
                uw.write(&unitig)?;
            }
            uw.close()?;
            Layer::<()>::build(&new_layer_dir, block_bits, evidence).map_err(olm_to_sk)?;
            g.len()
        } else {
            0
        };
        drop(g);

        let new_mphf = if any_new {
            Some(MphfOnly::open(&new_layer_dir).map_err(olm_to_sk)?)
        } else {
            None
        };

        // ── Prepare matrix directories for the new layer ──────────────────────
        // Absent columns (dst genomes) are written via append_column (all-zero/false).
        // Source-genome columns are created as mutable builders for pass 2.
        let mut new_src_builders: Vec<ColBuilder> = if any_new {
            let data_dir = match mode {
                MergeMode::Presence => new_layer_dir.join("presence"),
                MergeMode::Count    => new_layer_dir.join("counts"),
            };
            fs::create_dir_all(&data_dir)?;
            match mode {
                MergeMode::Presence => {
                    PersistentBitMatrixBuilder::new(n_new, &data_dir)
                        .map_err(SKError::Io)?.close().map_err(SKError::Io)?;
                    for _ in 0..n_dst_genomes {
                        PersistentBitMatrix::append_column(&data_dir, |_| false)
                            .map_err(SKError::Io)?;
                    }
                    (0..n_src_total).map(|g| -> SKResult<ColBuilder> {
                        let b = PersistentBitVecBuilder::new(
                            n_new, &col_path_bit(&data_dir, n_dst_genomes + g))?;
                        Ok(ColBuilder::Bit(b))
                    }).collect::<SKResult<_>>()?
                }
                MergeMode::Count => {
                    PersistentCompactIntMatrixBuilder::new(n_new, &data_dir)
                        .map_err(SKError::Io)?.close().map_err(SKError::Io)?;
                    for _ in 0..n_dst_genomes {
                        PersistentCompactIntMatrix::append_column(&data_dir, |_| 0)
                            .map_err(SKError::Io)?;
                    }
                    (0..n_src_total).map(|g| -> SKResult<ColBuilder> {
                        let b = PersistentCompactIntVecBuilder::new(
                            n_new, &col_path_int(&data_dir, n_dst_genomes + g))?;
                        Ok(ColBuilder::Int(b))
                    }).collect::<SKResult<_>>()?
                }
            }
        } else {
            vec![]
        };

        // Builders for existing layers: n_src_total per layer.
        // Columns n_dst_genomes .. n_dst_genomes + n_src_total - 1.
        let mut exist_builders: Vec<Vec<ColBuilder>> = (0..n_dst_layers)
            .map(|l| {
                let layer_dir = dst_index_dir.join(format!("layer_{l}"));
                let n = dst_map.layer(l).n();
                (0..n_src_total).map(|src_g| -> SKResult<ColBuilder> {
                    match mode {
                        MergeMode::Presence => {
                            let data_dir = layer_dir.join("presence");
                            let b = PersistentBitVecBuilder::new(
                                n, &col_path_bit(&data_dir, n_dst_genomes + src_g))?;
                            Ok(ColBuilder::Bit(b))
                        }
                        MergeMode::Count => {
                            let data_dir = layer_dir.join("counts");
                            let b = PersistentCompactIntVecBuilder::new(
                                n, &col_path_int(&data_dir, n_dst_genomes + src_g))?;
                            Ok(ColBuilder::Int(b))
                        }
                    }
                }).collect::<SKResult<_>>()
            })
            .collect::<SKResult<_>>()?;

        // ── Pass 2: fill builders ─────────────────────────────────────────────
        let mut col_offset = 0usize;
        for (src, src_n) in sources.iter() {
            let src_index_dir = src.part_dir(i).join(INDEX_SUBDIR);
            if !src_index_dir.exists() { col_offset += src_n; continue; }
            let src_meta = load_meta(&src_index_dir)?;

            for l in 0..src_meta.n_layers {
                let src_layer_dir = src_index_dir.join(format!("layer_{l}"));
                let reader = UnitigFileReader::open_sequential(&src_layer_dir.join("unitigs.bin"))?;
                let src_data = SrcLayerData::open(&src_layer_dir, mode)?;

                for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
                    let values = src_data.lookup(kmer, *src_n);
                    for (g, &value) in values.iter().enumerate() {
                        let builder_idx = col_offset + g;
                        if let Some((dst_layer, hit)) = dst_map.query(kmer) {
                            exist_builders[dst_layer][builder_idx].set_val(hit.slot, value);
                        } else if let Some(ref mphf) = new_mphf {
                            let slot = mphf.index(kmer);
                            new_src_builders[builder_idx].set_val(slot, value);
                        }
                    }
                }
            }
            col_offset += src_n;
        }

        // ── Close builders and update metadata ────────────────────────────────
        for (l, builders) in exist_builders.into_iter().enumerate() {
            let layer_dir = dst_index_dir.join(format!("layer_{l}"));
            for b in builders { b.close()?; }
            let n = dst_map.layer(l).n();
            let data_dir = match mode {
                MergeMode::Presence => layer_dir.join("presence"),
                MergeMode::Count    => layer_dir.join("counts"),
            };
            write_matrix_meta(&data_dir, n, n_dst_genomes + n_src_total).map_err(SKError::Io)?;
        }

        for b in new_src_builders { b.close()?; }
        if any_new {
            let data_dir = match mode {
                MergeMode::Presence => new_layer_dir.join("presence"),
                MergeMode::Count    => new_layer_dir.join("counts"),
            };
            write_matrix_meta(&data_dir, n_new, n_dst_genomes + n_src_total).map_err(SKError::Io)?;

            let mut part_meta = PartitionMeta::load(&dst_index_dir).map_err(olm_to_sk)?;
            part_meta.n_layers = new_layer_idx + 1;
            part_meta.save(&dst_index_dir).map_err(olm_to_sk)?;
        }

        Ok(())
    }
}
