use std::fs;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};

use tracing::debug;
use obipipeline::{
    Pipeline, PipelineError, PipelineSender, SharedFlatFn, Stage, WorkerPool,
    ThrottleGuard, throttle,
    make_sink, make_source, make_transform,
};

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

use crate::common::{ColBuilder, col_path_bit, col_path_int, load_meta, olm_to_sk, write_matrix_meta};
use crate::graph_pipeline::{build_graph, materialize_layer};
use crate::partition::KmerPartition;

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

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

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

pub(crate) enum SrcLayerData {
    Presence(MphfOnly, PersistentBitMatrix),
    Count(MphfOnly, PersistentCompactIntMatrix),
}

impl SrcLayerData {
    pub(crate) fn open(layer_dir: &Path, merge_mode: MergeMode) -> SKResult<Self> {
        let counts_dir = layer_dir.join("counts");
        match merge_mode {
            MergeMode::Presence => {
                if counts_dir.exists() && !layer_dir.join("presence").exists() {
                    let mphf = MphfOnly::open(layer_dir).map_err(|e| olm_to_sk(e, "merge"))?;
                    let mat = PersistentCompactIntMatrix::open(layer_dir).map_err(SKError::Io)?;
                    Ok(SrcLayerData::Count(mphf, mat))
                } else {
                    // presence dir exists, or neither exists → Implicit handled by open()
                    let mphf = MphfOnly::open(layer_dir).map_err(|e| olm_to_sk(e, "merge"))?;
                    let mat = PersistentBitMatrix::open(layer_dir).map_err(SKError::Io)?;
                    Ok(SrcLayerData::Presence(mphf, mat))
                }
            }
            MergeMode::Count => {
                let mphf = MphfOnly::open(layer_dir).map_err(|e| olm_to_sk(e, "merge"))?;
                if counts_dir.exists() {
                    let mat = PersistentCompactIntMatrix::open(layer_dir).map_err(SKError::Io)?;
                    Ok(SrcLayerData::Count(mphf, mat))
                } else {
                    // No counts → treat as implicit presence (all 1s)
                    let mat = PersistentBitMatrix::open(layer_dir).map_err(SKError::Io)?;
                    Ok(SrcLayerData::Presence(mphf, mat))
                }
            }
        }
    }

    /// 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> {
        let mut buf = vec![0u32; n_genomes];
        match self {
            SrcLayerData::Presence(mphf, mat) => mat.fill_row(mphf.index(kmer), &mut buf),
            SrcLayerData::Count(mphf, mat) => mat.fill_row(mphf.index(kmer), &mut buf),
        }
        buf
    }

    pub(crate) fn n_slots(&self) -> usize {
        match self {
            SrcLayerData::Presence(_, mat) => mat.n(),
            SrcLayerData::Count(_, mat) => mat.n(),
        }
    }

    /// MPHF lookup: returns the slot index for `kmer` (kmer must be in the domain).
    #[inline]
    pub(crate) fn slot(&self, kmer: CanonicalKmer) -> usize {
        match self {
            SrcLayerData::Presence(mphf, _) => mphf.index(kmer),
            SrcLayerData::Count(mphf, _) => mphf.index(kmer),
        }
    }

    /// Row lookup by slot index, bypassing the MPHF.
    #[inline]
    pub(crate) fn fill_row_by_slot(&self, slot: usize, n_genomes: usize) -> Vec<u32> {
        let mut buf = vec![0u32; n_genomes];
        match self {
            SrcLayerData::Presence(_, mat) => mat.fill_row(slot, &mut buf),
            SrcLayerData::Count(_, mat) => mat.fill_row(slot, &mut buf),
        }
        buf
    }

    /// Call `f` with a reference to the underlying matrix as `&dyn MatrixGroupOps`.
    pub(crate) fn with_matrix<R>(&self, f: impl FnOnce(&dyn MatrixGroupOps) -> R) -> R {
        match self {
            SrcLayerData::Presence(_, mat) => f(mat),
            SrcLayerData::Count(_, mat) => f(mat),
        }
    }
}

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

const INDEX_SUBDIR: &str = "index";

// ── 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<usize> {
        let dst_index_dir = self.part_dir(i).join(INDEX_SUBDIR);
        if !dst_index_dir.exists() {
            return Ok(0);
        }

        load_meta(&dst_index_dir, "merge")?; // ensure meta.json exists before LayeredMap::open
        let dst_map = Arc::new(LayeredMap::<()>::open(&dst_index_dir).map_err(|e| olm_to_sk(e, "merge"))?);
        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(|e| olm_to_sk(e, "merge"))?;
            }
        }

        // ── Pass 1: pipeline — parallel file read + dst_map filter + graph fill ─
        //
        // Source  : list of unitigs.bin paths (one per source × layer)
        // Flat    : open file, emit Vec<CanonicalKmer> batches (BeeGFS parallel I/O)
        // Transform: filter via dst_map.query() — thread-safe, LayeredMap<()>: Sync
        // Sink    : push new kmers into GraphDeBruijn (single thread, no locks needed)

        // Collect file paths (propagates load_meta errors before the pipeline starts)
        let mut unitig_paths: Vec<PathBuf> = Vec::new();
        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, "merge")?;
            for l in 0..src_meta.n_layers {
                let p = src_index_dir.join(format!("layer_{l}")).join("unitigs.bin");
                if p.exists() {
                    unitig_paths.push(p);
                }
            }
        }

        let n_src_layers = unitig_paths.len();
        debug!("partition {i}: de Bruijn graph build start — {n_src_layers} source layer(s)");

        const BATCH: usize = 4096;
        let n_workers = rayon::current_num_threads().min(16).max(4);
        // At most 2 files open simultaneously: keeps n_workers-2 workers free
        // for the Transform stage. Each open file monopolises one worker for the
        // full duration of its read, so this must stay well below n_workers.
        let max_open = 2;

        let dst_filter = Arc::clone(&dst_map);

        let g = build_graph(
            unitig_paths.into_iter(),
            move |path: PathBuf, emit: &mut dyn FnMut(Vec<CanonicalKmer>)| -> SKResult<()> {
                let reader = UnitigFileReader::open_sequential(&path)?;
                let mut batch: Vec<CanonicalKmer> = Vec::with_capacity(BATCH);
                for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
                    batch.push(kmer);
                    if batch.len() == BATCH {
                        emit(std::mem::replace(&mut batch, Vec::with_capacity(BATCH)));
                    }
                }
                if !batch.is_empty() {
                    emit(batch);
                }
                Ok(())
            },
            move |kmer| dst_filter.query(kmer).is_none(),
            n_workers,
            max_open,
        )?;

        let any_new = g.len() > 0;
        debug!("partition {i}: de Bruijn graph done — {} new kmers", g.len());

        // 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 {
            debug!("partition {i}: unitig traversal start — {} nodes", g.len());
            let n_nodes = materialize_layer(g, &new_layer_dir, block_bits, evidence)?;
            debug!("partition {i}: MPHF build done");
            n_nodes
        } else {
            drop(g);
            0
        };

        let t_open = std::time::Instant::now();
        let new_mphf: Option<Arc<MphfOnly>> = if any_new {
            Some(Arc::new(MphfOnly::open(&new_layer_dir).map_err(|e| olm_to_sk(e, "merge"))?))
        } else {
            None
        };
        debug!("partition {i}: MPHF open in {:.3}s", t_open.elapsed().as_secs_f64());

        // ── 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 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![]
        };

        let t_builders = std::time::Instant::now();
        // Builders for existing layers: n_src_total per layer.
        // Columns n_dst_genomes .. n_dst_genomes + n_src_total - 1.
        let 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<_>>()?;

        debug!("partition {i}: builders ready in {:.3}s", t_builders.elapsed().as_secs_f64());

        // ── Pass 2: fill builders (pipeline) ─────────────────────────────────
        let t_pass2 = std::time::Instant::now();
        // Collect source items before the pipeline so load_meta errors propagate
        // via ? before any worker thread is spawned.
        let mut pass2_items: Vec<(usize, usize, PathBuf)> = Vec::new();
        {
            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, "merge")?;
                for l in 0..src_meta.n_layers {
                    let src_layer_dir = src_index_dir.join(format!("layer_{l}"));
                    if src_layer_dir.join("unitigs.bin").exists() {
                        pass2_items.push((col_offset, *src_n, src_layer_dir));
                    }
                }
                col_offset += src_n;
            }
        }

        enum Pass2Data {
            SrcLayer((usize, usize, PathBuf, ThrottleGuard)),
            RawBatch((usize, usize, Arc<SrcLayerData>, Vec<CanonicalKmer>)),
            WriteBatch(Vec<(Option<usize>, usize, usize, u32)>),
        }

        let exist_locked: Vec<Vec<Arc<Mutex<ColBuilder>>>> = exist_builders
            .into_iter()
            .map(|layer| layer.into_iter().map(|b| Arc::new(Mutex::new(b))).collect())
            .collect();
        let new_locked: Vec<Arc<Mutex<ColBuilder>>> = new_src_builders
            .into_iter()
            .map(|b| Arc::new(Mutex::new(b)))
            .collect();
        let exist_sink: Vec<Vec<Arc<Mutex<ColBuilder>>>> = exist_locked.iter()
            .map(|layer| layer.iter().map(Arc::clone).collect())
            .collect();
        let new_sink: Vec<Arc<Mutex<ColBuilder>>> = new_locked.iter().map(Arc::clone).collect();
        let dst_map_t2     = Arc::clone(&dst_map);
        let new_mphf_t2    = new_mphf.clone();
        let pass2_err: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
        let err_cap2       = Arc::clone(&pass2_err);

        let capacity = 2;
        let throttled_pass2 = throttle(pass2_items.into_iter(), max_open);

        let pipeline2 = Pipeline::new(
            make_source!(Pass2Data, throttled_pass2.map(|t| {
                let (col_offset, src_n, src_layer_dir) = t.item;
                (col_offset, src_n, src_layer_dir, t.guard)
            }), SrcLayer),
            vec![
                Stage::Flat(Arc::new(
                    move |data: Pass2Data,
                          push:  &PipelineSender<Result<Pass2Data, PipelineError>>,
                          delta: &PipelineSender<isize>|
                    {
                        if let Pass2Data::SrcLayer((col_offset, src_n, src_layer_dir, _guard)) = data {
                            // _guard dropped at end of block, releasing the slot.
                            let reader = match UnitigFileReader::open_sequential(
                                &src_layer_dir.join("unitigs.bin"),
                            ) {
                                Ok(r)  => r,
                                Err(e) => {
                                    *err_cap2.lock().unwrap() = Some(e.to_string());
                                    delta.send(-1).ok();
                                    return;
                                }
                            };
                            let src_data = match SrcLayerData::open(&src_layer_dir, mode) {
                                Ok(d)  => Arc::new(d),
                                Err(e) => {
                                    *err_cap2.lock().unwrap() = Some(e.to_string());
                                    delta.send(-1).ok();
                                    return;
                                }
                            };
                            const BATCH: usize = 4096;
                            let mut batch: Vec<CanonicalKmer> = Vec::with_capacity(BATCH);
                            let mut count: isize = 0;
                            for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
                                batch.push(kmer);
                                if batch.len() == BATCH {
                                    let b = std::mem::replace(&mut batch, Vec::with_capacity(BATCH));
                                    push.send(Ok(Pass2Data::RawBatch((
                                        col_offset, src_n, Arc::clone(&src_data), b,
                                    )))).ok();
                                    count += 1;
                                }
                            }
                            if !batch.is_empty() {
                                push.send(Ok(Pass2Data::RawBatch((
                                    col_offset, src_n, src_data, batch,
                                )))).ok();
                                count += 1;
                            }
                            delta.send(count - 1).ok();
                        }
                    }
                ) as SharedFlatFn<Pass2Data>),
                make_transform!(Pass2Data, {
                    move |(col_offset, src_n, src_data, kmers): (usize, usize, Arc<SrcLayerData>, Vec<CanonicalKmer>)|
                        -> Vec<(Option<usize>, usize, usize, u32)>
                    {
                        let mut ops: Vec<(Option<usize>, usize, usize, u32)> = Vec::new();
                        for kmer in kmers {
                            let values = src_data.lookup(kmer, src_n);
                            if let Some((dst_layer, hit)) = dst_map_t2.query(kmer) {
                                for (g, val) in values.into_iter().enumerate() {
                                    ops.push((Some(dst_layer), col_offset + g, hit.slot, val));
                                }
                            } else if let Some(ref mphf) = new_mphf_t2 {
                                let slot = mphf.index(kmer);
                                for (g, val) in values.into_iter().enumerate() {
                                    ops.push((None, col_offset + g, slot, val));
                                }
                            }
                        }
                        ops
                    }
                }, RawBatch, WriteBatch),
            ],
            make_sink!(Pass2Data, {
                move |ops: Vec<(Option<usize>, usize, usize, u32)>| {
                    for (layer_opt, col, slot, val) in ops {
                        match layer_opt {
                            Some(l) => exist_sink[l][col].lock().unwrap().set_val(slot, val),
                            None    => new_sink[col].lock().unwrap().set_val(slot, val),
                        }
                    }
                }
            }, WriteBatch),
        );

        WorkerPool::new(pipeline2, n_workers, capacity).run();
        debug!("partition {i}: pass2 pipeline done in {:.3}s", t_pass2.elapsed().as_secs_f64());

        if let Some(msg) = Arc::try_unwrap(pass2_err)
            .unwrap_or_else(|_| panic!("pass2: pass2_err not uniquely owned"))
            .into_inner()
            .unwrap_or_else(|e| e.into_inner())
        {
            return Err(SKError::InvalidData { context: "merge pass2", detail: msg });
        }

        let t_close = std::time::Instant::now();
        // ── Close builders and update metadata ────────────────────────────────
        for (l, builders) in exist_locked.into_iter().enumerate() {
            let layer_dir = dst_index_dir.join(format!("layer_{l}"));
            for b in builders {
                Arc::try_unwrap(b)
                    .unwrap_or_else(|_| panic!("pass2: exist_builder[{l}] not uniquely owned"))
                    .into_inner()
                    .unwrap_or_else(|e| e.into_inner())
                    .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_locked {
            Arc::try_unwrap(b)
                .unwrap_or_else(|_| panic!("pass2: new_builder not uniquely owned"))
                .into_inner()
                .unwrap_or_else(|e| e.into_inner())
                .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(|e| olm_to_sk(e, "merge"))?;
            part_meta.n_layers = new_layer_idx + 1;
            part_meta.save(&dst_index_dir).map_err(|e| olm_to_sk(e, "merge"))?;
        }

        debug!("partition {i}: builders closed in {:.3}s", t_close.elapsed().as_secs_f64());

        Ok(n_new)
    }
}
