use std::collections::{HashMap, VecDeque};
use std::io::{self, BufWriter, Write};
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::time::Instant;

use clap::Args;
use obikindex::KmerIndex;
use obikpartitionner::{KmerDesc, QueryHit, QueryStats};
use obikrope::Rope;
use obikseq::CanonicalKmer;
use obilayeredmap::IndexMode;
use obipipeline::{Throttled, ThrottleGuard, throttle};
use obiread::chunk::read_sequence_chunks_sized;
use obiread::record::{SeqRecord, parse_chunk};
use obiskbuilder::SuperKmerIter;
use obisys::{Reporter, Stage, available_memory_bytes, spinner};
use tracing::{debug, info};

// ── Pipeline data ─────────────────────────────────────────────────────────────

enum QueryData {
    Path(Throttled<PathBuf>),
    Chunk(Rope),
    Output(Vec<u8>),
}

// SAFETY: Rope contains Cell<u8> which is !Sync, but pipeline items are owned
// exclusively through channels — no item is ever shared across threads.
unsafe impl Send for QueryData {}
unsafe impl Sync for QueryData {}

// ── CLI ───────────────────────────────────────────────────────────────────────

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

    /// Input sequences (FASTA/FASTQ, optionally gzip-compressed)
    #[arg(num_args = 1..)]
    pub inputs: Vec<String>,

    /// Report per-position coverage vectors per genome (adds "coverage" to JSON)
    #[arg(long)]
    pub detail: bool,

    /// Enable 1-mismatch approximate matching
    #[arg(long)]
    pub mismatch: bool,

    /// Count k-mers absent from the index (adds kmer_missing annotation)
    #[arg(long)]
    pub count_missing: bool,

    /// Report per-genome presence (0/1) instead of raw counts
    #[arg(long)]
    pub force_presence: bool,

    /// Minimum accumulated match count to declare a genome present (implies --force-presence)
    #[arg(long, default_value_t = 1)]
    pub presence_threshold: u32,

    /// Override the Findere z parameter from index metadata
    #[arg(short = 'z', long)]
    pub findere_z: Option<usize>,

    /// Number of worker threads
    #[arg(
        short = 'T',
        long,
        default_value_t = obisys::effective_parallelism()
    )]
    pub threads: usize,

    /// I/O chunk size in MiB (default: auto-sized from available RAM and thread count)
    #[arg(long)]
    pub chunk_size: Option<usize>,

    /// Maximum number of input files open simultaneously.
    /// Defaults to threads/4 (minimum 1). Keep below the number of workers
    /// to ensure CPU workers are always available for the transform stage.
    #[arg(long)]
    pub max_open_files: Option<usize>,
}

impl QueryArgs {
    pub fn effective_max_open(&self) -> usize {
        self.max_open_files
            .unwrap_or_else(|| (self.threads / 4).max(1))
            .max(1)
    }
}

// ── QueryBatch ────────────────────────────────────────────────────────────────

/// A batch of query sequences, with k-mers deduplicated directly (not just at
/// the superkmer level) and pre-split by partition.
///
/// Superkmer *construction* (`SuperKmerIter`) is still required — it's the
/// mechanism that computes minimizers and partition routing — but the dedup
/// key is the canonical k-mer, not the superkmer: two different superkmers
/// that happen to share a k-mer (read overlaps, repeats, a SNP splitting an
/// otherwise-identical run) are deduplicated too, not just identical whole
/// superkmers. This also means each unique k-mer triggers at most one MPHF
/// lookup, not one per occurrence.
pub struct QueryBatch {
    /// Sequence ids in batch order.
    pub ids: Vec<String>,
    /// Raw sequence bytes (for output), in batch order.
    pub seqs: Vec<Vec<u8>>,
    /// Total kmer count per sequence (used for `--detail` coverage allocation).
    pub n_kmers: Vec<u32>,
    /// Deduplicated k-mer occurrences, one map per partition.
    pub by_partition: Vec<HashMap<CanonicalKmer, Vec<KmerDesc>>>,
}

impl QueryBatch {
    /// Build a batch from a vec of parsed sequence records, deduplicating
    /// k-mers and routing them to partitions in the same pass.
    pub fn from_records(
        records: Vec<SeqRecord>,
        k: usize,
        level_max: usize,
        theta: f64,
        n_partitions: usize,
    ) -> Self {
        let mut ids = Vec::with_capacity(records.len());
        let mut seqs = Vec::with_capacity(records.len());
        let mut n_kmers = Vec::with_capacity(records.len());
        let mask = (n_partitions as u64) - 1;
        let mut by_partition: Vec<HashMap<CanonicalKmer, Vec<KmerDesc>>> =
            (0..n_partitions).map(|_| HashMap::new()).collect();

        for (seq_idx, record) in records.into_iter().enumerate() {
            let mut kmer_offset = 0u32;

            for rsk in SuperKmerIter::new(&record.normalized, k, level_max, theta) {
                let part_idx = (rsk.minimizer().seq_hash() & mask) as usize;
                let map = &mut by_partition[part_idx];
                for (j, kmer) in rsk.superkmer().iter_canonical_kmers().enumerate() {
                    map.entry(kmer).or_default().push(KmerDesc {
                        seq_idx: seq_idx as u32,
                        pos: kmer_offset + j as u32,
                    });
                }
                let n = (rsk.seql() - k + 1) as u32;
                kmer_offset += n;
            }

            ids.push(record.id);
            seqs.push(record.sequence);
            n_kmers.push(kmer_offset);
        }

        Self {
            ids,
            seqs,
            n_kmers,
            by_partition,
        }
    }
}

// ── SmerIndex — sparse "was this k-mer found at all" bookkeeping ─────────────

/// Tracks, per (sequence, s-mer position), whether the k-mer was found in the
/// index at all — independent of *which* genome(s) matched. Sized
/// `total_smers` (one `bool` per s-mer occurrence in the chunk), **not**
/// multiplied by `n_genomes`: this is the O(1)-per-position bookkeeping that
/// `kmer_missing` needs (the leftmost-s-mer-of-window membership test), kept
/// dense because it's already cheap — the `n_genomes`-scaled data lives in
/// the sparse per-genome hit lists built alongside it (see `process_chunk`).
struct SmerIndex {
    in_index: Vec<bool>, // total_smers
    offsets: Vec<usize>, // offsets[i]..offsets[i+1] = s-mer range for sequence i
}

impl SmerIndex {
    fn new(n_kmers_per_seq: &[u32]) -> Self {
        let mut offsets = Vec::with_capacity(n_kmers_per_seq.len() + 1);
        let mut total = 0usize;
        offsets.push(0);
        for &n in n_kmers_per_seq {
            total += n as usize;
            offsets.push(total);
        }
        Self {
            in_index: vec![false; total],
            offsets,
        }
    }

    /// Mark the k-mer at (seq, kmer) as found in the index — independent of
    /// any particular genome's value. Called once per hit k-mer (stage 1 of
    /// `query_partition_with`), regardless of how the column-major fetch
    /// (stage 2) later reports per-genome values.
    fn mark_found(&mut self, seq: usize, kmer: usize) {
        let abs = self.offsets[seq] + kmer;
        self.in_index[abs] = true;
    }

    #[inline]
    fn is_in_index(&self, seq: usize, kmer: usize) -> bool {
        self.in_index[self.offsets[seq] + kmer]
    }
}

// ── Sparse Findere: per-genome run detection + sliding-window minimum ────────

/// One confirmed z-window: genome `g`'s window ending at k-mer `pos` (the
/// *leftmost* s-mer of the window, i.e. the k_user-mer's output position) is
/// fully present and nonzero, with window-minimum `value`.
type ConfirmedHit = (u32, u32, u32); // (seq_idx, pos_out, value)

/// Reduce one genome's raw sparse s-mer hits — `(seq_idx, pos_smer, raw_value)`,
/// unsorted, exactly as delivered by `QueryHit::Value` — into confirmed
/// z-windows, without ever visiting a position that had no hit at all.
///
/// A z-window is confirmed only when all z s-mers in it are present *and*
/// nonzero for this genome (matching the dense sliding-window's semantics,
/// where "not in index" or a zero value both contribute 0 to the window
/// minimum) — which can only happen inside a maximal run of consecutive
/// `pos_smer` values for the same sequence. `hits` is sorted in place by
/// `(seq_idx, pos_smer)` to expose those runs; the monotone-deque
/// window-minimum then runs per run, on run-relative indices, identical in
/// spirit to the dense version's whole-sequence scan.
///
/// Returns the confirmed hits plus `(n_runs, total_run_len)` for logging —
/// a low average run length relative to `z` means most hits fail to form a
/// complete window.
fn sparse_findere_for_genome(
    hits: &mut [(u32, u32, u32)],
    z: usize,
    presence: bool,
    threshold: u32,
) -> (Vec<ConfirmedHit>, usize, usize) {
    hits.sort_unstable_by_key(|&(seq, pos, _)| (seq, pos));

    let mut confirmed = Vec::new();
    let mut n_runs = 0usize;
    let mut total_run_len = 0usize;
    let mut dq: VecDeque<(usize, u32)> = VecDeque::new(); // (run-relative index, value)

    let mut i = 0;
    while i < hits.len() {
        let seq = hits[i].0;
        let mut j = i + 1;
        while j < hits.len() && hits[j].0 == seq && hits[j].1 == hits[j - 1].1 + 1 {
            j += 1;
        }
        let run = &hits[i..j];
        n_runs += 1;
        total_run_len += run.len();

        dq.clear();
        for (k, &(_, pos, val)) in run.iter().enumerate() {
            while dq.back().map_or(false, |&(_, v)| v >= val) {
                dq.pop_back();
            }
            dq.push_back((k, val));
            while dq.front().map_or(false, |&(fk, _)| fk + z <= k) {
                dq.pop_front();
            }
            if k + 1 >= z {
                let win_min = dq.front().unwrap().1;
                if win_min > 0 {
                    let pos_out = pos + 1 - z as u32;
                    let c = if presence { u32::from(win_min >= threshold) } else { win_min };
                    confirmed.push((seq, pos_out, c));
                }
            }
        }

        i = j;
    }

    (confirmed, n_runs, total_run_len)
}

// ── Per-sequence accumulator ──────────────────────────────────────────────────

struct SeqAcc {
    kmer_count: u32,
    kmer_missing: u32,
    genome_totals: Vec<u32>,
}

impl SeqAcc {
    fn new(n_genomes: usize) -> Self {
        Self {
            kmer_count: 0,
            kmer_missing: 0,
            genome_totals: vec![0u32; n_genomes],
        }
    }
}

// ── process_chunk ─────────────────────────────────────────────────────────────

fn process_chunk(
    idx: &KmerIndex,
    rope: Rope,
    k: usize,
    n_genomes: usize,
    n_partitions: usize,
    with_counts: bool,
    effective_z: usize,
    detail: bool,
    count_missing: bool,
    force_presence: bool,
    presence_threshold: u32,
) -> Vec<u8> {
    let chunk_start = Instant::now();
    let chunk_bytes = rope.len();

    let records = parse_chunk(&rope, k);
    if records.is_empty() {
        return Vec::new();
    }

    let batch = QueryBatch::from_records(records, k, 6, 0.7, n_partitions);
    let n_seqs = batch.ids.len();

    // Estimate QueryBatch::by_partition's actual memory footprint: the
    // k-mer-level dedup map (roadmap point 5) — one HashMap<CanonicalKmer,
    // Vec<KmerDesc>> per partition, sized by *unique* k-mers, not shrunk by
    // dedup. On real workloads with a low intra-chunk duplication rate this
    // can dwarf every other per-chunk structure, including the sparse
    // Findere ones logged further down — unlike those, chunk_bytes's formula
    // (run()) does not account for this at all today. Measured by allocated
    // capacity, not logical length, to reflect real memory pressure
    // (HashMap/Vec growth slack) — `by_partition` is alive for the entire
    // process_chunk call (never drained, only iterated by reference), so
    // this is its footprint for the whole chunk lifetime, not a transient.
    let hashmap_slot_bytes = (std::mem::size_of::<CanonicalKmer>()
        + std::mem::size_of::<Vec<KmerDesc>>()
        + 1) as u64; // +1 ≈ hashbrown control byte per slot
    let by_partition_map_bytes: u64 = batch
        .by_partition
        .iter()
        .map(|m| m.capacity() as u64 * hashmap_slot_bytes)
        .sum();
    let by_partition_desc_bytes: u64 = batch
        .by_partition
        .iter()
        .flat_map(|m| m.values())
        .map(|v| v.capacity() as u64 * std::mem::size_of::<KmerDesc>() as u64)
        .sum();
    let by_partition_bytes = by_partition_map_bytes + by_partition_desc_bytes;

    debug!(
        n_unique_kmers_total = batch.by_partition.iter().map(|m| m.len() as u64).sum::<u64>(),
        by_partition_map_bytes,
        by_partition_desc_bytes,
        by_partition_bytes,
        chunk_bytes,
        "by_partition memory retained"
    );

    // Sparse bookkeeping for the whole chunk:
    // - smer_index: O(total_smers) — is this s-mer in the index at all.
    // - by_genome[g]: raw (seq_idx, pos_smer, value) hits for genome g, only
    //   ever containing nonzero entries (query_partition_with never emits a
    //   QueryHit::Value for a zero value) — empty for every genome this chunk
    //   never matched, which is the common case for unrelated queries.
    let mut smer_index = SmerIndex::new(&batch.n_kmers);
    let mut by_genome: Vec<Vec<(u32, u32, u32)>> = (0..n_genomes).map(|_| Vec::new()).collect();

    // Dedup-ratio bookkeeping: occurrences (from batch.n_kmers, computed
    // before dedup) vs. unique k-mers actually queried (query_stats) — the
    // entire justification for k-mer-level dereplication (see query.md,
    // Future work point 5). If this ratio stays close to 1.0 on real data,
    // dereplication isn't paying for itself and that should show up here.
    let n_occurrences: u64 = batch.n_kmers.iter().map(|&n| n as u64).sum();
    let mut query_stats = QueryStats::default();

    for (part_idx, kmers) in batch.by_partition.iter().enumerate() {
        if kmers.is_empty() {
            continue;
        }

        let stats = idx.partition()
            .query_partition_with(
                part_idx,
                kmers,
                n_genomes,
                with_counts,
                |event| match event {
                    QueryHit::Found(descs) => {
                        for desc in descs {
                            smer_index.mark_found(desc.seq_idx as usize, desc.pos as usize);
                        }
                    }
                    QueryHit::Value(descs, g, v) => {
                        for desc in descs {
                            by_genome[g].push((desc.seq_idx, desc.pos, v));
                        }
                    }
                },
            )
            .unwrap_or_else(|e| {
                eprintln!("query error on partition {part_idx}: {e}");
                std::process::exit(1);
            });
        query_stats += stats;
    }

    debug!(
        n_occurrences,
        n_unique_kmers = query_stats.n_unique_kmers,
        n_mphf_calls = query_stats.n_mphf_calls,
        n_hits = query_stats.n_hits,
        n_columns_scanned = query_stats.n_columns_scanned,
        n_col_get_calls = query_stats.n_col_get_calls,
        "k-mer dedup + column-major fetch"
    );

    // ── Sparse Findere: per-genome run detection + sliding-window minimum ────
    //
    // Confirmed z-windows, per genome, replace the dense win_min matrix:
    // total retained memory is O(actual hits), not O(total_smers × n_genomes)
    // — the whole point of this pass. See sparse_findere_for_genome's doc for
    // why run detection is equivalent to the dense scan's semantics.
    let presence  = force_presence || !with_counts;
    let threshold = presence_threshold;
    let z         = effective_z;

    let n_kmers_out: Vec<usize> = batch
        .n_kmers
        .iter()
        .map(|&n| {
            let n = n as usize;
            if n >= z { n - z + 1 } else { 0 }
        })
        .collect();
    let mut out_offsets = Vec::with_capacity(n_seqs + 1);
    {
        let mut total = 0usize;
        out_offsets.push(0);
        for &n in &n_kmers_out {
            total += n;
            out_offsets.push(total);
        }
    }
    let total_out = *out_offsets.last().unwrap_or(&0);

    let n_dense_would_be = n_occurrences as u64 * n_genomes as u64;
    let mut n_sparse_entries = 0u64;
    let mut n_runs_total = 0usize;
    let mut run_len_total = 0usize;

    let mut confirmed_by_genome: Vec<Vec<ConfirmedHit>> = Vec::with_capacity(n_genomes);
    for hits in &mut by_genome {
        n_sparse_entries += hits.len() as u64;
        let (confirmed, n_runs, run_len) = sparse_findere_for_genome(hits, z, presence, threshold);
        n_runs_total += n_runs;
        run_len_total += run_len;
        confirmed_by_genome.push(confirmed);
    }

    debug!(
        n_dense_would_be,
        n_sparse_entries,
        n_runs = n_runs_total,
        avg_run_len = if n_runs_total > 0 { run_len_total as f64 / n_runs_total as f64 } else { 0.0 },
        z,
        "sparse Findere"
    );

    // Actual bytes retained by the sparse hit structures (by_genome +
    // confirmed_by_genome, both alive simultaneously at this point — see the
    // chunk-size formula's comment in `run()`), by allocated capacity rather
    // than logical length so this reflects real memory pressure including
    // Vec growth slack. `empirical_multiplier` is directly comparable to
    // BYTES_PER_KMER_PER_GENOME (`run()`) — the ratio a cluster run's logs
    // need to judge whether that constant is over- or under-conservative for
    // real data, instead of guessing.
    const HIT_ENTRY_BYTES: u64 = std::mem::size_of::<(u32, u32, u32)>() as u64;
    let by_genome_bytes: u64 = by_genome.iter().map(|v| v.capacity() as u64 * HIT_ENTRY_BYTES).sum();
    let confirmed_bytes: u64 = confirmed_by_genome.iter().map(|v| v.capacity() as u64 * HIT_ENTRY_BYTES).sum();
    let retained_bytes = by_genome_bytes + confirmed_bytes;

    debug!(
        by_genome_bytes,
        confirmed_bytes,
        retained_bytes,
        chunk_bytes,
        empirical_multiplier = retained_bytes as f64 / chunk_bytes.max(1) as f64,
        "sparse memory retained"
    );

    // ── Accumulate: genome totals (per genome, from confirmed hits) ──────────
    let mut accs: Vec<SeqAcc> = (0..n_seqs).map(|_| SeqAcc::new(n_genomes)).collect();
    let mut confirmed_any = vec![false; total_out];

    for (g, hits) in confirmed_by_genome.iter().enumerate() {
        for &(seq_idx, pos_out, c) in hits {
            let abs_out = out_offsets[seq_idx as usize] + pos_out as usize;
            confirmed_any[abs_out] = true;
            accs[seq_idx as usize].genome_totals[g] += c;
        }
    }

    // ── Accumulate: kmer_count / kmer_missing (per position, genome-independent) ─
    for seq_idx in 0..n_seqs {
        let out_n = n_kmers_out[seq_idx];
        let acc = &mut accs[seq_idx];
        for pos in 0..out_n {
            let abs_out = out_offsets[seq_idx] + pos;
            if confirmed_any[abs_out] {
                acc.kmer_count += 1;
            } else if !smer_index.is_in_index(seq_idx, pos) {
                acc.kmer_missing += 1;
            }
        }
    }

    // ── Coverage (--detail): densify only when actually requested ────────────
    let mut cov: Vec<Vec<Vec<u32>>> = if detail {
        n_kmers_out.iter().map(|&n| vec![vec![0u32; n]; n_genomes]).collect()
    } else {
        Vec::new()
    };
    if detail {
        for (g, hits) in confirmed_by_genome.iter().enumerate() {
            for &(seq_idx, pos_out, c) in hits {
                cov[seq_idx as usize][g][pos_out as usize] += c;
            }
        }
    }

    // Capacity estimate: actual sequence + ID bytes, plus JSON overhead per record.
    // JSON per record ≈ 50 fixed chars + ~20 per genome (label + count value) + 100 (overhead).
    let seq_bytes: usize = batch.seqs.iter().map(|s| s.len()).sum();
    let id_bytes: usize = batch.ids.iter().map(|s| s.len()).sum();
    let cap = seq_bytes + id_bytes + n_seqs * (4 + 50 + n_genomes * 20) + 100;
    let mut buf = Vec::with_capacity(cap);
    emit_batch(
        &batch,
        &accs,
        idx.meta(),
        count_missing,
        detail,
        &cov,
        &mut buf,
    );

    debug!(
        chunk_bytes,
        n_seqs,
        n_smers = batch.n_kmers.iter().map(|&n| n as u64).sum::<u64>(),
        wall_ms = chunk_start.elapsed().as_millis() as u64,
        "process_chunk"
    );

    buf
}

// ── GuardedChunkIter — keeps the throttle slot guard alive until the file is exhausted ──

/// Wraps a per-file `Rope` chunk iterator together with its `ThrottleGuard`,
/// so the guard (and the throttle slot it holds) is only released once the
/// file has been fully read — never earlier, never held past that point.
struct GuardedChunkIter {
    inner:      Box<dyn Iterator<Item = Rope> + Send>,
    _guard:     ThrottleGuard,
    files_open: Arc<AtomicU32>,
}

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

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

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

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

    let k = idx.kmer_size();
    let n_genomes = idx.meta().genomes.len();
    let n_partitions = idx.n_partitions();
    let with_counts = idx.meta().config.with_counts;
    let n_workers = args.threads.max(1);

    // Chunk size: each chunk stays in memory for its entire processing lifetime.
    //
    // Per-chunk memory is no longer a dense n_genomes-wide buffer (removed in
    // the sparse Findere rework, see process_chunk) — it now scales with
    // *actual hit count*, not with total_kmers_in_chunk × n_genomes
    // unconditionally. BYTES_PER_KMER_PER_GENOME below is therefore a
    // pathological-case bound, not a typical-case estimate: it protects
    // against a fully-dense hit pattern (every k-mer of the query matching
    // every genome — a degenerate case, e.g. low-complexity input theta-
    // filtering should mostly reject, or an index of near-duplicate genomes),
    // where by_genome and confirmed_by_genome (process_chunk) both end up
    // holding one (seq_idx, pos, value) entry — 3 × u32 = 12 bytes, vs. 4
    // bytes for the old dense encoding, where position was implicit in the
    // array index — per (k-mer, genome) pair, and *coexist simultaneously*
    // (by_genome isn't freed before confirmed_by_genome is built), for a
    // worst case of ~24 bytes/pair before Vec growth slack. `cov` remains
    // fully dense when --detail is set (unaffected by the sparse rework),
    // still roughly doubling the n_genomes-scaled cost.
    //
    // For realistic, sparse hit patterns actual memory is far below this
    // bound — see the "sparse memory retained" debug log in process_chunk,
    // which reports the empirical bytes-per-raw-byte multiplier actually
    // observed per chunk, directly comparable to BYTES_PER_KMER_PER_GENOME
    // below. Tightening this constant for typical-case throughput (at the
    // cost of pathological-case safety margin) is a deliberate tuning
    // decision to make from that data, not something to guess at here.
    //
    // BASE_OVERHEAD approximates what scales with chunk_bytes alone,
    // independent of n_genomes: the Rope itself, parsed SeqRecord sequence +
    // normalised bytes, the superkmer dedup map, and the JSON output buffer.
    // Like the n_genomes-scaled term, this is an estimate — validate against
    // actual peak RSS (Stage::stop's `rss` in the summary table) on real
    // workloads rather than trusting it blindly.
    //
    // We target ≤ 50 % of available RAM across all concurrent workers
    // (SAFETY_FACTOR).
    const BASE_OVERHEAD: u64 = 4;
    const BYTES_PER_KMER_PER_GENOME: u64 = 8; // pathological-case bound — see comment above
    const SAFETY_FACTOR: u64 = 2;

    let detail_factor: u64 = if args.detail { 2 } else { 1 };
    let overhead_multiplier =
        BASE_OVERHEAD + n_genomes as u64 * BYTES_PER_KMER_PER_GENOME * detail_factor;

    let chunk_bytes = args
        .chunk_size
        .map(|mb| mb * 1024 * 1024)
        .unwrap_or_else(|| {
            let avail = available_memory_bytes();
            let computed = avail / (n_workers as u64 * overhead_multiplier * SAFETY_FACTOR);
            computed.clamp(4 * 1024 * 1024, 256 * 1024 * 1024) as usize
        });

    debug!(
        chunk_bytes,
        n_genomes,
        detail = args.detail,
        overhead_multiplier,
        estimated_peak_chunk_bytes = chunk_bytes as u64 * overhead_multiplier,
        "chunk-size formula resolved"
    );

    let effective_z: usize = args
        .findere_z
        .unwrap_or_else(|| match idx.meta().config.evidence {
            IndexMode::Approx { z, .. } | IndexMode::Hybrid { z, .. } => z as usize,
            IndexMode::Exact => 1,
        });

    info!(
        "query: k={k}, {} genome(s), with_counts={with_counts}, z={effective_z}, \
         mismatch={}, detail={}",
        n_genomes, args.mismatch, args.detail
    );

    if args.mismatch {
        eprintln!("warning: --mismatch not yet implemented, ignored");
    }

    let detail = args.detail;
    let count_missing = args.count_missing;
    let force_presence = args.force_presence;
    let presence_threshold = args.presence_threshold;

    // Throttled iterator over input file paths: at most `effective_max_open()`
    // files are open at once. Opening + decompressing + chunking each file is
    // now a Flat pipeline stage, executed across the `n_workers` pool — not
    // serialised in the pipe's dedicated source thread (see steps::scatter /
    // cmd::superkmer for the same pattern applied to indexing).
    info!("query: chunk_size={}MiB, max_open_files={}", chunk_bytes / (1024 * 1024), args.effective_max_open());

    let paths: Vec<PathBuf> = args.inputs.iter().map(PathBuf::from).collect();
    let path_source = throttle(paths.into_iter(), args.effective_max_open());

    // Instrumentation: total bytes processed (for the EMA throughput readout),
    // number of files currently open/being chunked, and number of chunks
    // currently being processed by a worker — all read from the spinner loop
    // below, updated from inside the pipe closures.
    let total_bytes = Arc::new(AtomicU64::new(0));
    let files_open = Arc::new(AtomicU32::new(0));
    let chunks_active = Arc::new(AtomicU32::new(0));

    let pipe = obipipeline::make_pipe! {
        QueryData : Throttled<PathBuf> => Vec<u8>,
        || {
            let files_open = Arc::clone(&files_open);
            move |pw: Throttled<PathBuf>| -> GuardedChunkIter {
                let path = pw.item;
                let guard = pw.guard;
                let path_str = path.to_str().unwrap_or("").to_owned();
                files_open.fetch_add(1, Ordering::Relaxed);
                let open_start = Instant::now();
                // Hard-exit on file-open failure (mirrors the previous behaviour):
                // propagating this as a pipeline Err would hit a known scheduler
                // hang on early stage errors (obipipeline::scheduler::WorkerPool::run
                // breaks its main loop without unblocking the still-running source
                // thread, so the final `h.join()` never returns) — worth fixing in
                // obipipeline itself, but out of scope here; sidestepping it like the
                // original code already did is the safe choice for this change.
                let iter = read_sequence_chunks_sized(&path_str, chunk_bytes).unwrap_or_else(|e| {
                    eprintln!("error opening {path_str}: {e}");
                    std::process::exit(1);
                });
                debug!(
                    path = %path_str,
                    open_ms = open_start.elapsed().as_millis() as u64,
                    "opened query input file"
                );
                let err_path = path_str.clone();
                GuardedChunkIter {
                    inner: Box::new(iter.filter_map(move |r| match r {
                        Ok(rope) => Some(rope),
                        Err(e) => {
                            eprintln!("read error: {err_path}: {e}");
                            None
                        }
                    })),
                    _guard: guard,
                    files_open: Arc::clone(&files_open),
                }
            }
        } : Path => Chunk,
        | {
            let idx = Arc::clone(&idx);
            let total_bytes = Arc::clone(&total_bytes);
            let chunks_active = Arc::clone(&chunks_active);
            move |rope: Rope| {
                chunks_active.fetch_add(1, Ordering::Relaxed);
                let bytes = rope.len() as u64;
                let out = process_chunk(
                    &idx, rope, k, n_genomes, n_partitions, with_counts,
                    effective_z, detail, count_missing, force_presence, presence_threshold,
                );
                total_bytes.fetch_add(bytes, Ordering::Relaxed);
                chunks_active.fetch_sub(1, Ordering::Relaxed);
                out
            }
        } : Chunk => Output,
    };

    let t = Stage::start("query");
    let pb = spinner("query");

    let mut ema_rate: f64 = 0.0;
    let mut last_t = Instant::now();
    let mut last_bytes: u64 = 0;
    const ALPHA: f64 = 0.15;

    let mut out = BufWriter::new(io::stdout());
    for block in pipe.apply(path_source, n_workers, 2) {
        if !block.is_empty() {
            out.write_all(&block).expect("write error");
        }

        let now = Instant::now();
        let dt = now.duration_since(last_t).as_secs_f64();
        if dt > 0.1 {
            let total = total_bytes.load(Ordering::Relaxed);
            let instant = (total - last_bytes) as f64 / dt;
            ema_rate = ALPHA * instant + (1.0 - ALPHA) * ema_rate;
            last_t = now;
            last_bytes = total;
            let bp = total as f64;
            let (count_str, rate_str) = if bp >= 1e9 {
                (format!("{:.2} GB", bp / 1e9), format!("{:.0} MB/s", ema_rate / 1e6))
            } else {
                (format!("{:.0} MB", bp / 1e6), format!("{:.0} MB/s", ema_rate / 1e6))
            };
            let active = chunks_active.load(Ordering::Relaxed);
            let open = files_open.load(Ordering::Relaxed);
            pb.set_message(format!("{count_str}  {rate_str}  [files open: {open}, chunks in flight: {active}]"));
        }
    }
    out.flush().expect("flush error");

    pb.finish_and_clear();

    let mut rep = Reporter::new();
    rep.push(t.stop());
    rep.print();
}

// ── Output ────────────────────────────────────────────────────────────────────

fn emit_batch(
    batch: &QueryBatch,
    accs: &[SeqAcc],
    meta: &obikindex::meta::IndexMeta,
    count_missing: bool,
    detail: bool,
    cov: &[Vec<Vec<u32>>],
    out: &mut impl Write,
) {
    for (seq_idx, (id, seq)) in batch.ids.iter().zip(batch.seqs.iter()).enumerate() {
        let acc = &accs[seq_idx];
        let mut ann = serde_json::Map::new();

        ann.insert("kmer_count".into(), acc.kmer_count.into());
        if count_missing {
            ann.insert("kmer_missing".into(), acc.kmer_missing.into());
        }

        let mut match_map = serde_json::Map::new();
        for (g, genome) in meta.genomes.iter().enumerate() {
            if acc.genome_totals[g] != 0 {
                match_map.insert(genome.label.clone(), acc.genome_totals[g].into());
            }
        }
        ann.insert("kmer_strict_matches".into(), match_map.into());

        if detail && !cov.is_empty() {
            let mut cov_map = serde_json::Map::new();
            for (g, genome) in meta.genomes.iter().enumerate() {
                let v: Vec<serde_json::Value> = cov[seq_idx][g].iter().map(|&x| x.into()).collect();
                cov_map.insert(genome.label.clone(), v.into());
            }
            ann.insert("coverage".into(), cov_map.into());
        }

        // OBITools4 FASTA format: >id {"key":value,...}
        let _ = out.write_all(b">");
        let _ = out.write_all(id.as_bytes());
        let _ = out.write_all(b" ");
        let _ = serde_json::to_writer(&mut *out, &ann);
        let _ = out.write_all(b"\n");
        let _ = out.write_all(seq);
        let _ = out.write_all(b"\n");
    }
}

#[cfg(test)]
#[path = "tests/query.rs"]
mod tests;
