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

use clap::Args;
use obifastwrite::write_unitig;
use obikindex::KmerIndex;
use obikseq::set_k;
use obiskio::UnitigFileReader;
use rayon::prelude::*;
use tracing::info;

#[derive(Args)]
pub struct UnitigArgs {
    /// Index directory produced by the `index` command
    pub index: PathBuf,
}

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();
    set_k(k);
    let n = idx.n_partitions();
    info!("dumping unitigs from {n} partitions (k={k})");

    let stdout = Mutex::new(BufWriter::new(io::stdout()));

    (0..n).into_par_iter().for_each(|i| {
        let path = idx.layer_unitigs_path(i, 0);
        if !path.exists() {
            return;
        }

        // open_sequential: works with and without .idx (approx or exact index)
        let reader = UnitigFileReader::open_sequential(&path).unwrap_or_else(|e| {
            eprintln!("error opening unitigs (partition {i}): {e}");
            std::process::exit(1)
        });

        for (j, unitig) in reader.iter_unitigs() {
            let mut out = stdout.lock().unwrap();
            write_unitig(&unitig, k, i, j, &mut *out).unwrap_or_else(|e| {
                eprintln!("write error: {e}");
                std::process::exit(1)
            });
        }
    });

    stdout.into_inner().unwrap().flush().expect("flush error");
}
