use std::io::Write;
use std::sync::atomic::{AtomicUsize, Ordering};

use rayon::prelude::*;

use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use obikpartitionner::KmerFilter;

impl KmerIndex {
    /// Write a CSV table of all indexed kmers to `out`.
    ///
    /// Columns: `kmer`, then one column per genome (in index order).
    /// Values are counts (u32) when `use_counts = true`, otherwise 0/1.
    ///
    /// `force_presence` overrides `with_counts`: even if the index stores counts,
    /// the output uses 0/1 presence columns.
    ///
    /// Partitions are scanned in parallel; each partition buffers its output locally
    /// before the main thread writes the chunks in partition order.
    ///
    /// The caller must have set the global kmer length (`obikseq::set_k`) before
    /// calling this method.
    pub fn dump<W: Write, F: Fn() + Send + Sync>(
        &self,
        out: &mut W,
        force_presence: bool,
        debug: bool,
        head: Option<usize>,
        filters: &[Box<dyn KmerFilter>],
        on_partition: F,
    ) -> OKIResult<()> {
        let genomes    = &self.meta.genomes;
        let use_counts = self.meta.config.with_counts && !force_presence;
        let n_genomes  = genomes.len().max(1);
        let kmer_size  = self.kmer_size();

        // ── Header ────────────────────────────────────────────────────────────
        if debug {
            write!(out, "partition,layer,")?;
        }
        write!(out, "kmer")?;
        for g in genomes {
            write!(out, ",{}", g.label)?;
        }
        writeln!(out)?;

        // ── Rows — parallel over partitions ───────────────────────────────────
        let n = self.n_partitions();

        let write_row = |buf: &mut Vec<u8>, row: &[u32], prefix: &str| {
            let _ = buf.write_all(prefix.as_bytes());
            for &v in row { let _ = write!(buf, ",{v}"); }
            let _ = buf.write_all(b"\n");
        };

        let chunks: Vec<OKIResult<Vec<u8>>> = if let Some(limit) = head {
            // ── Bounded: atomic counter, early exit when limit reached ────────
            let remaining = AtomicUsize::new(limit);
            (0..n).into_par_iter().map(|i| {
                if remaining.load(Ordering::Relaxed) == 0 { return Ok(vec![]); }
                let mut buf = Vec::<u8>::new();
                let try_write = |buf: &mut Vec<u8>, row: &[u32], prefix: &str| -> bool {
                    match remaining.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |cur| {
                        if cur > 0 { Some(cur - 1) } else { None }
                    }) {
                        Err(_) => false,
                        Ok(_)  => { write_row(buf, row, prefix); true }
                    }
                };
                if debug {
                    self.partition
                        .iter_partition_kmers_located(i, use_counts, n_genomes, filters, |part, layer, kmer, row| {
                            let seq = String::from_utf8(kmer.to_ascii()).unwrap_or_else(|_| "?".repeat(kmer_size));
                            try_write(&mut buf, &row, &format!("{part},{layer},{seq}"))
                        })
                        .map_err(OKIError::Partition)?;
                } else {
                    self.partition
                        .iter_partition_kmers(i, use_counts, n_genomes, filters, |kmer, row| {
                            let seq = String::from_utf8(kmer.to_ascii()).unwrap_or_else(|_| "?".repeat(kmer_size));
                            try_write(&mut buf, &row, &seq)
                        })
                        .map_err(OKIError::Partition)?;
                }
                on_partition();
                Ok(buf)
            }).collect()
        } else {
            // ── Unbounded: no atomic, no contention ───────────────────────────
            (0..n).into_par_iter().map(|i| {
                let mut buf = Vec::<u8>::new();
                if debug {
                    self.partition
                        .iter_partition_kmers_located(i, use_counts, n_genomes, filters, |part, layer, kmer, row| {
                            let seq = String::from_utf8(kmer.to_ascii()).unwrap_or_else(|_| "?".repeat(kmer_size));
                            write_row(&mut buf, &row, &format!("{part},{layer},{seq}"));
                            true
                        })
                        .map_err(OKIError::Partition)?;
                } else {
                    self.partition
                        .iter_partition_kmers(i, use_counts, n_genomes, filters, |kmer, row| {
                            let seq = String::from_utf8(kmer.to_ascii()).unwrap_or_else(|_| "?".repeat(kmer_size));
                            write_row(&mut buf, &row, &seq);
                            true
                        })
                        .map_err(OKIError::Partition)?;
                }
                on_partition();
                Ok(buf)
            }).collect()
        };

        // ── Sequential write ──────────────────────────────────────────────────
        for chunk in chunks {
            out.write_all(&chunk?)?;
        }

        out.flush()?;
        Ok(())
    }
}
