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

use clap::Args;
use obikindex::KmerIndex;
use obilayeredmap::IndexMode;
use obiread::chunk::read_sequence_chunks;
use obiread::record::{SeqRecord, parse_chunk};
use obikrope::Rope;
use obikseq::{RoutableSuperKmer, set_k, set_m};
use obiskbuilder::SuperKmerIter;
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,
}

// ── 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());
        let mut map: HashMap<RoutableSuperKmer, Vec<SKDesc>> = HashMap::new();

        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
    }
}

// ── 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],
        }
    }
}

// ── Findere z-window filter ───────────────────────────────────────────────────

/// Apply the Findere z-window filter to per-kmer query results for one superkmer.
/// Aggregate s-mer query results into k-mer answers using a Findere z-window.
///
/// Input:  N s-mer results  (indexed kmer size s = k − z + 1).
/// Output: N − z + 1 k-mer results (user kmer size k).
///
/// For each genome g independently: k-mer at position i is confirmed iff all z values
/// results[i..i+z][g] are nonzero (None counts as zero for all genomes).
/// Output values are taken from results[i]; genomes not confirmed are zeroed.
fn apply_findere(
    results: &[Option<Box<[u32]>>],
    z: usize,
    n_genomes: usize,
) -> Vec<Option<Box<[u32]>>> {
    let n = results.len();
    if z <= 1 {
        return results.iter().map(|r| r.as_ref().map(|row| row.clone())).collect();
    }
    if n < z {
        return Vec::new();
    }

    let out_n = n - z + 1;
    let mut confirmed = vec![vec![false; n_genomes]; out_n];

    for g in 0..n_genomes {
        let hit = |i: usize| results[i].as_ref().map_or(false, |r| r[g] > 0);

        let mut count: u32 = (0..z).filter(|&j| hit(j)).count() as u32;
        if count == z as u32 { confirmed[0][g] = true; }

        for i in 1..out_n {
            if hit(i - 1)     { count -= 1; }
            if hit(i + z - 1) { count += 1; }
            if count == z as u32 { confirmed[i][g] = true; }
        }
    }

    (0..out_n).map(|i| {
        let first = results[i].as_ref()?;
        let mut row: Box<[u32]> = first.clone();
        for g in 0..n_genomes {
            if !confirmed[i][g] { row[g] = 0; }
        }
        if row.iter().any(|&v| v > 0) { Some(row) } else { None }
    }).collect()
}

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

    // Per-sequence s-mer result vectors in global sequence position order.
    // All partitions fill into this structure before Findere is applied.
    let mut seq_results: Vec<Vec<Option<Box<[u32]>>>> = batch.n_kmers.iter()
        .map(|&n| vec![None; n as usize])
        .collect();

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

        let kmer_results = idx
            .partition()
            .query_partition(part_idx, part_sks, k, n_genomes, with_counts)
            .unwrap_or_else(|e| {
                eprintln!("query error on partition {part_idx}: {e}");
                std::process::exit(1);
            });

        for (rsk, sk_kmer_results) in part_sks.iter().zip(kmer_results.iter()) {
            let descs = batch.map.get(*rsk).expect("rsk must be in map");
            for desc in descs {
                let offset = desc.kmer_offset as usize;
                let dst = &mut seq_results[desc.seq_idx as usize];
                for (j, hit) in sk_kmer_results.iter().enumerate() {
                    dst[offset + j] = hit.as_ref().map(|r| r.clone());
                }
            }
        }
    }

    // Apply Findere on each complete sequence vector, then accumulate.
    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 accs: Vec<SeqAcc> =
        (0..n_seqs).map(|_| SeqAcc::new(n_genomes)).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;

    for seq_idx in 0..n_seqs {
        let filtered = apply_findere(&seq_results[seq_idx], effective_z, n_genomes);
        let acc = &mut accs[seq_idx];

        for (pos, hit) in filtered.iter().enumerate() {
            match hit {
                None => {
                    if seq_results[seq_idx][pos].is_none() {
                        acc.kmer_missing += 1;
                    }
                }
                Some(row) => {
                    acc.kmer_count += 1;
                    for (g, &v) in row.iter().enumerate() {
                        if v == 0 { continue; }
                        let contribution = if presence {
                            u32::from(v >= threshold)
                        } else {
                            v
                        };
                        acc.genome_totals[g] += contribution;
                        if detail {
                            cov[seq_idx][g][pos] += contribution;
                        }
                    }
                }
            }
        }
    }

    let mut buf = Vec::new();
    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);
    }));

    set_k(idx.kmer_size());
    set_m(idx.minimizer_size());

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

    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.
    let paths: Vec<PathBuf> = args.inputs.iter().map(PathBuf::from).collect();
    let all_chunks = paths.into_iter().flat_map(|path| {
        let path_str = path.to_str().unwrap_or("").to_owned();
        match read_sequence_chunks(&path_str) {
            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());
        }

        let ann_str = serde_json::to_string(&ann).unwrap_or_else(|_| "{}".to_string());

        // OBITools4 FASTA format: >id {"key":value,...}
        let _ = writeln!(out, ">{id} {ann_str}");
        let _ = out.write_all(seq);
        let _ = out.write_all(b"\n");
    }
}
