use std::collections::{HashMap, VecDeque};
use std::io::{self, BufWriter, Write};
use std::path::PathBuf;
use std::sync::Arc;

use clap::Args;
use obikindex::KmerIndex;
use obikrope::Rope;
use obikseq::RoutableSuperKmer;
use obilayeredmap::IndexMode;
use obiread::chunk::read_sequence_chunks_sized;
use obiread::record::{SeqRecord, parse_chunk};
use obiskbuilder::SuperKmerIter;
use obisys::available_memory_bytes;
use tracing::info;

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

enum QueryData {
    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 = std::thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(1)
    )]
    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>,
}

// ── SKDesc — one occurrence of a superkmer in the batch ───────────────────────

/// Describes one occurrence of a superkmer in the query batch.
pub struct SKDesc {
    /// Index of the source sequence within the batch.
    pub seq_idx: u32,
    /// Kmer offset of the first kmer of this superkmer within its sequence.
    pub kmer_offset: u32,
}

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

/// A batch of query sequences with their superkmers deduplicated.
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 superkmer map.
    pub map: HashMap<RoutableSuperKmer, Vec<SKDesc>>,
}

impl QueryBatch {
    /// Build a batch from a vec of parsed sequence records.
    pub fn from_records(records: Vec<SeqRecord>, k: usize, level_max: usize, theta: f64) -> 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());
        // Upper-bound estimate: at most one superkmer per k bases.
        // Avoids repeated rehash on large chunks.
        let cap = records.iter().map(|r| r.normalized.len()).sum::<usize>() / k.max(1);
        let mut map: HashMap<RoutableSuperKmer, Vec<SKDesc>> = HashMap::with_capacity(cap);

        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 n = (rsk.seql() - k + 1) as u32;
                map.entry(rsk).or_default().push(SKDesc {
                    seq_idx: seq_idx as u32,
                    kmer_offset,
                });
                kmer_offset += n;
            }

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

        Self {
            ids,
            seqs,
            n_kmers,
            map,
        }
    }

    /// Split the superkmer map by partition index.
    pub fn split_by_partition(&self, n_partitions: usize) -> Vec<Vec<&RoutableSuperKmer>> {
        let mask = (n_partitions as u64) - 1;
        let mut by_part: Vec<Vec<&RoutableSuperKmer>> = vec![Vec::new(); n_partitions];
        for rsk in self.map.keys() {
            let part = (rsk.minimizer().seq_hash() & mask) as usize;
            by_part[part].push(rsk);
        }
        by_part
    }
}

// ── KmerResults — allocation-free ragged result matrix ───────────────────────

/// Flat storage for per-kmer query results across all sequences in a chunk.
///
/// Replaces `Vec<Vec<Option<Box<[u32]>>>>` — a single allocation for the whole
/// chunk instead of one `Box<[u32]>` per found k-mer.
struct KmerResults {
    data: Vec<u32>,      // total_kmers × n_genomes, row-major
    in_index: Vec<bool>, // total_kmers — true if the kmer was found in the index
    offsets: Vec<usize>, // offsets[i]..offsets[i+1] = kmer range for sequence i
    n_genomes: usize,
}

impl KmerResults {
    fn new(n_kmers_per_seq: &[u32], n_genomes: usize) -> 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 {
            data: vec![0u32; total * n_genomes],
            in_index: vec![false; total],
            offsets,
            n_genomes,
        }
    }

    fn n_kmers_for(&self, seq: usize) -> usize {
        self.offsets[seq + 1] - self.offsets[seq]
    }

    fn set(&mut self, seq: usize, kmer: usize, row: &[u32]) {
        let abs = self.offsets[seq] + kmer;
        self.in_index[abs] = true;
        let base = abs * self.n_genomes;
        self.data[base..base + self.n_genomes].copy_from_slice(row);
    }

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

    /// Value for genome `g` at (seq, kmer); meaningful only when `is_in_index`.
    #[inline]
    fn val(&self, seq: usize, kmer: usize, g: usize) -> u32 {
        self.data[(self.offsets[seq] + kmer) * self.n_genomes + g]
    }
}

// ── 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 records = parse_chunk(&rope, k);
    if records.is_empty() {
        return Vec::new();
    }

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

    // Flat result matrix — one allocation for the whole chunk.
    let mut results = KmerResults::new(&batch.n_kmers, n_genomes);

    let by_part = batch.split_by_partition(n_partitions);

    for (part_idx, part_sks) in by_part.iter().enumerate() {
        if part_sks.is_empty() {
            continue;
        }

        idx.partition()
            .query_partition_with(
                part_idx,
                part_sks,
                k,
                n_genomes,
                with_counts,
                |sk_idx, kmer_idx, row| {
                    let rsk = part_sks[sk_idx];
                    let descs = batch.map.get(rsk).expect("rsk must be in map");
                    for desc in descs {
                        results.set(
                            desc.seq_idx as usize,
                            desc.kmer_offset as usize + kmer_idx,
                            row,
                        );
                    }
                },
            )
            .unwrap_or_else(|e| {
                eprintln!("query error on partition {part_idx}: {e}");
                std::process::exit(1);
            });
    }

    // Sliding window minimum — one reusable buffer and one deque per batch.
    //
    // win_min[pos * n_genomes + g] = min count across the z-window [pos, pos+z)
    // for genome g, where "not in index" counts as 0.
    //
    // win_min > 0  ↔  all z consecutive kmers are in the index with count > 0
    //               ↔  Findere confirmation (for z=1 this degenerates to the
    //                  simple case with no overhead).
    //
    // Works uniformly for count matrices and presence/absence (0/1) matrices.
    let max_n_kmers = batch.n_kmers.iter().map(|&n| n as usize).max().unwrap_or(0);
    let mut win_min = vec![0u32; max_n_kmers * n_genomes];

    let mut accs: Vec<SeqAcc> = (0..n_seqs).map(|_| SeqAcc::new(n_genomes)).collect();

    let n_kmers_out: Vec<usize> = batch
        .n_kmers
        .iter()
        .map(|&n| {
            let n = n as usize;
            if n >= effective_z { n - effective_z + 1 } else { 0 }
        })
        .collect();

    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()
    };

    let presence  = force_presence || !with_counts;
    let threshold = presence_threshold;
    let z         = effective_z;

    // Deque reused across all (seq, genome) pairs.
    let mut dq: VecDeque<(usize, u32)> = VecDeque::with_capacity(z + 1);

    for seq_idx in 0..n_seqs {
        let n     = results.n_kmers_for(seq_idx);
        let out_n = n_kmers_out[seq_idx];
        if out_n == 0 { continue; }

        let mins = &mut win_min[..out_n * n_genomes];
        mins.fill(0);

        // ── Per-genome sliding window minimum ─────────────────────────────────
        for g in 0..n_genomes {
            dq.clear();
            for i in 0..n {
                let v_i = if results.is_in_index(seq_idx, i) {
                    results.val(seq_idx, i, g)
                } else {
                    0
                };
                // Evict elements that have left the window.
                while dq.front().map_or(false, |&(f, _)| f + z <= i) {
                    dq.pop_front();
                }
                // Maintain monotone non-decreasing back→front for minimum at front.
                while dq.back().map_or(false, |&(_, v)| v >= v_i) {
                    dq.pop_back();
                }
                dq.push_back((i, v_i));
                // Window [pos, pos+z) is complete when i = pos + z - 1.
                if i + 1 >= z {
                    let pos = i + 1 - z;
                    mins[pos * n_genomes + g] = dq.front().unwrap().1;
                }
            }
        }

        // ── Accumulate ────────────────────────────────────────────────────────
        let acc = &mut accs[seq_idx];
        for pos in 0..out_n {
            let any = (0..n_genomes).any(|g| mins[pos * n_genomes + g] > 0);
            if !any {
                if !results.is_in_index(seq_idx, pos) {
                    acc.kmer_missing += 1;
                }
                continue;
            }
            acc.kmer_count += 1;
            for g in 0..n_genomes {
                let v = mins[pos * n_genomes + g];
                if v == 0 { continue; }
                let c = if presence { u32::from(v >= threshold) } else { v };
                acc.genome_totals[g] += c;
                if detail { cov[seq_idx][g][pos] += 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,
    );
    buf
}

// ── 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.
    // Overhead per raw byte is ~8× (Rope + parsed records + superkmers + results).
    // We target ≤ 50 % of available RAM across all concurrent workers.
    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 * 16);
            computed.clamp(4 * 1024 * 1024, 256 * 1024 * 1024) as usize
        });

    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;

    // Flat iterator over all Rope chunks from all input files.
    // I/O runs in the source thread; chunk processing is parallelised by the pipe.
    info!("query: chunk_size={}MiB", chunk_bytes / (1024 * 1024));

    let paths: Vec<PathBuf> = args.inputs.iter().map(PathBuf::from).collect();
    let all_chunks = paths.into_iter().flat_map(move |path| {
        let path_str = path.to_str().unwrap_or("").to_owned();
        match read_sequence_chunks_sized(&path_str, chunk_bytes) {
            Ok(iter) => Box::new(iter.filter_map(|r| match r {
                Ok(rope) => Some(rope),
                Err(e) => {
                    eprintln!("read error: {e}");
                    None
                }
            })) as Box<dyn Iterator<Item = Rope> + Send>,
            Err(e) => {
                eprintln!("error opening {path_str}: {e}");
                std::process::exit(1);
            }
        }
    });

    let pipe = obipipeline::make_pipe! {
        QueryData : Rope => Vec<u8>,
        | {
            let idx = Arc::clone(&idx);
            move |rope: Rope| {
                process_chunk(
                    &idx, rope, k, n_genomes, n_partitions, with_counts,
                    effective_z, detail, count_missing, force_presence, presence_threshold,
                )
            }
        } : Chunk => Output,
    };

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

// ── 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() {
            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");
    }
}
