use std::path::PathBuf;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Instant;

use obisys::spinner;
use obiread::NucPage;
use obikpartitionner::KmerPartition;
use obipipeline::{ThrottleGuard, Throttled, throttle};
use obisys::{Reporter, Stage};
use tracing::info;

use crate::cli::PipelineData;

// ── Iterator that keeps the slot guard alive until the file is exhausted ──────

struct GuardedIter {
    inner:       Box<dyn Iterator<Item = NucPage> + Send>,
    _guard:      ThrottleGuard,
    flat_active: Arc<AtomicU32>,
}

impl Iterator for GuardedIter {
    type Item = NucPage;
    fn next(&mut self) -> Option<NucPage> {
        self.inner.next()
    }
}

impl Drop for GuardedIter {
    fn drop(&mut self) {
        self.flat_active.fetch_sub(1, Ordering::Relaxed);
    }
}

// ── scatter ───────────────────────────────────────────────────────────────────

/// Run scatter: normalise → build superkmers → route to partition → close.
/// Reports the "scatter" stage to `rep`.
pub fn scatter(
    kp: &mut KmerPartition,
    path_source: impl Iterator<Item = PathBuf> + Send + 'static,
    k: usize,
    level_max: usize,
    theta: f64,
    n_workers: usize,
    max_open: usize,
    rep: &mut Reporter,
) {
    use obikseq::RoutableSuperKmer;

    // Throttle in the source thread — never in a worker — to prevent deadlock.
    let throttled = throttle(path_source, max_open);

    let file_count      = Arc::new(AtomicU64::new(0));
    let flat_active     = Arc::new(AtomicU32::new(0));
    let transform_active = Arc::new(AtomicU32::new(0));

    let t = Stage::start("scatter");
    let pipe = obipipeline::make_pipe! {
        PipelineData : Throttled<PathBuf> => Vec<RoutableSuperKmer>,
        ||? {
            let file_count   = Arc::clone(&file_count);
            let flat_active  = Arc::clone(&flat_active);
            let k = k;
            move |pw: Throttled<PathBuf>| {
                let path  = pw.item;
                let guard = pw.guard;
                let n = file_count.fetch_add(1, Ordering::Relaxed) + 1;
                info!("indexing [{}]: {}", n, path.display());
                let path_str = path.to_str().unwrap_or("").to_owned();
                flat_active.fetch_add(1, Ordering::Relaxed);
                obiread::open_nuc_stream(&path_str, k)
                    .map(|iter| GuardedIter { inner: iter, _guard: guard, flat_active: Arc::clone(&flat_active) })
            }
        } : Path => NucPage,
        | {
            let transform_active = Arc::clone(&transform_active);
            move |page| {
                transform_active.fetch_add(1, Ordering::Relaxed);
                let result = obiskbuilder::build_superkmers_page(page, k, level_max, theta);
                transform_active.fetch_sub(1, Ordering::Relaxed);
                result
            }
        } : NucPage => Batch,
    };

    let pb = spinner("scatter");

    let mut total_bases: u64 = 0;
    let mut ema_rate: f64 = 0.0;
    let mut last_t = Instant::now();
    let mut last_bases: u64 = 0;
    let kmer_overlap = (k - 1) as u64;
    const ALPHA: f64 = 0.15;

    for batch in pipe.apply(throttled, n_workers, 1) {
        total_bases += batch.iter()
            .map(|sk| (sk.seql() as u64).saturating_sub(kmer_overlap))
            .sum::<u64>();
        let now = Instant::now();
        let dt = now.duration_since(last_t).as_secs_f64();
        if dt > 0.1 {
            let instant = (total_bases - last_bases) as f64 / dt;
            ema_rate = ALPHA * instant + (1.0 - ALPHA) * ema_rate;
            last_t = now;
            last_bases = total_bases;
            let bp = total_bases as f64;
            let (count_str, rate_str) = if bp >= 1e9 {
                (format!("{:.2} Gbp", bp / 1e9), format!("{:.0} Mbp/s", ema_rate / 1e6))
            } else {
                (format!("{:.0} Mbp", bp / 1e6), format!("{:.0} Mbp/s", ema_rate / 1e6))
            };
            let n_files = file_count.load(Ordering::Relaxed);
            let r = flat_active.load(Ordering::Relaxed);
            let c = transform_active.load(Ordering::Relaxed);
            pb.set_message(format!("{count_str}  {rate_str}  {n_files} files  [R:{r} C:{c}]"));
        }
        kp.write_batch(batch).unwrap_or_else(|e| {
            eprintln!("error: {e}");
            std::process::exit(1);
        });
    }
    pb.finish_and_clear();
    kp.close().expect("close error");
    rep.push(t.stop());
}
