use std::io::{self, BufWriter, Write};
use std::path::PathBuf;
use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, Ordering};

use clap::Args;
use obidebruinj::GraphDeBruijn;
use obifastwrite::write_unitig;
use obikindex::KmerIndex;
use obisys::{Reporter, Stage, progress_bar, spinner};
use rayon::prelude::*;
use tracing::info;

use super::predicate::FilterArgs;

#[derive(Args)]
pub struct UnitigArgs {
    /// Index directory
    pub index: PathBuf,

    #[command(flatten)]
    pub filter: FilterArgs,
}

pub fn run(args: UnitigArgs) {
    let idx = KmerIndex::open(&args.index).unwrap_or_else(|e| {
        eprintln!("error opening index: {e}");
        std::process::exit(1);
    });

    let k = idx.kmer_size();
    let n = idx.n_partitions();
    let n_genomes = idx.meta().genomes.len().max(1);
    let use_counts = idx.meta().config.with_counts;

    info!("unitig: building de Bruijn graph from {n} partition(s) (k={k})");

    let filters = args.filter.build_filters(&idx.meta().genomes);
    let partition = idx.partition();
    let mut rep = Reporter::new();

    // ── Phase 1 : collect filtered kmers in parallel ──────────────────────────
    let pb = progress_bar("unitig", n as u64, "partitions");
    let stage = Stage::start("build graph");
    let g = (0..n)
        .into_par_iter()
        .fold(GraphDeBruijn::new, |mut local_g, i| {
            partition
                .iter_partition_kmers(i, use_counts, n_genomes, &filters, |kmer, _row| {
                    local_g.push(kmer);
                    true
                })
                .unwrap_or_else(|e| {
                    eprintln!("error reading partition {i}: {e}");
                    std::process::exit(1);
                });
            pb.inc(1);
            local_g
        })
        .reduce(GraphDeBruijn::new, |mut a, b| {
            a.merge(b);
            a
        });
    pb.finish_and_clear();
    rep.push(stage.stop());

    info!("unitig: {} distinct k-mers", g.len());

    // ── Phase 2 : compute degrees ─────────────────────────────────────────────
    let pb = spinner("degrees");
    let stage = Stage::start("compute degrees");
    g.compute_degrees_and_mark_starts();
    pb.finish_and_clear();
    rep.push(stage.stop());

    // ── Phase 3 : enumerate unitigs and write as FASTA ───────────────────────
    let pb = spinner("unitig");
    let out = Mutex::new(BufWriter::new(io::stdout()));
    let j = AtomicUsize::new(0);

    let stage = Stage::start("enumerate unitigs");
    g.for_each_unitig(|nuc_iter| {
        let unitig: obikseq::unitig::Unitig = nuc_iter.collect();
        let idx = j.fetch_add(1, Ordering::Relaxed);
        let mut w = out.lock().unwrap();
        write_unitig(&unitig, k, 0, idx, &mut *w).unwrap_or_else(|e| {
            eprintln!("write error: {e}");
            std::process::exit(1);
        });
        if idx % 10_000 == 0 {
            pb.set_message(format!("{idx} unitigs written"));
        }
    });
    pb.finish_and_clear();
    rep.push(stage.stop());

    out.into_inner().unwrap().flush().expect("flush error");
    rep.print();
}
