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

use clap::Args;
use obifastwrite::write_scatter;
use obikseq::{RoutableSuperKmer, set_k, set_m};

use obipipeline::{Throttled, throttle};

use crate::cli::{CommonArgs, PipelineData, partitions_to_bits};

#[derive(Args)]
pub struct SuperkmerArgs {
    #[command(flatten)]
    pub common: CommonArgs,
}

// ── Stage functions ───────────────────────────────────────────────────────────

fn write_batch(
    batch: Vec<RoutableSuperKmer>,
    out: &mut BufWriter<io::Stdout>,
    partition_bits: usize,
    k: usize,
    m: usize,
) -> io::Result<()> {
    let partition_mask = (1u64 << partition_bits) - 1;
    for rsk in batch {
        let minimizer = *rsk.minimizer();
        let partition = (minimizer.seq_hash() & partition_mask) as usize;
        write_scatter(rsk.superkmer(), out, k, m, partition, minimizer)?;
    }
    Ok(())
}

// ── Entry point ───────────────────────────────────────────────────────────────

pub fn run(args: SuperkmerArgs) {
    args.common.validate();

    let k = args.common.kmer_size;
    let m = args.common.minimizer_size;
    let theta = args.common.theta;
    let level_max = args.common.level_max;
    let partition_bits = partitions_to_bits(args.common.partitions);
    let n_workers = args.common.threads.max(1);
    let max_open = args.common.effective_max_open();

    set_k(k);
    set_m(m);

    let path_source = throttle(args.common.seqfile_paths(), max_open);

    let pipe = obipipeline::make_pipe! {
        PipelineData : Throttled<PathBuf> => Vec<RoutableSuperKmer>,
        ||? {
            let k = k;
            move |pw: Throttled<PathBuf>| {
                let path_str = pw.item.to_str().unwrap_or("").to_owned();
                let _guard = pw.guard;
                obiread::open_nuc_stream(&path_str, k)
            }
        } : Path => NucPage,
        | { move |page| obiskbuilder::build_superkmers_page(page, k, level_max, theta) } : NucPage => Batch,
    };

    let mut out = BufWriter::new(io::stdout());
    for batch in pipe.apply(path_source, n_workers, 1) {
        write_batch(batch, &mut out, partition_bits, k, m).expect("write error");
    }
    out.flush().expect("flush error");
}
