use super::*;
use obicompactvec::PersistentCompactIntMatrix;
use obikseq::{set_k, Kmer, Sequence as _, Unitig};
use obiskio::DEFAULT_BLOCK_BITS;
use crate::meta::IndexMode;
use tempfile::tempdir;

fn write_unitigs(dir: &Path, seqs: &[&[u8]]) {
    let mut w = UnitigFileWriter::create(&dir.join(UNITIGS_FILE)).unwrap();
    for s in seqs {
        w.write(&Unitig::from_ascii(s)).unwrap();
    }
    w.close().unwrap();
}

fn all_canonical_kmers(dir: &Path) -> Vec<CanonicalKmer> {
    UnitigFileReader::open_sequential(&dir.join(UNITIGS_FILE)).unwrap()
        .iter_indexed_canonical_kmers()
        .map(|(kmer, _, _)| kmer)
        .collect()
}

#[test]
fn build_and_query_all_kmers_found() {
    set_k(4);
    let dir = tempdir().unwrap();
    write_unitigs(dir.path(), &[b"AAAACGT"]);
    let kmers = all_canonical_kmers(dir.path());
    Layer::<()>::build(dir.path(), DEFAULT_BLOCK_BITS, &IndexMode::Exact).unwrap();
    let layer = Layer::<()>::open(dir.path(), &IndexMode::Exact).unwrap();
    for kmer in kmers {
        assert!(layer.query(kmer).is_some(), "kmer should be present");
    }
}

#[test]
fn counts_are_stored_and_retrieved() {
    set_k(4);
    let dir = tempdir().unwrap();
    write_unitigs(dir.path(), &[b"AAAACGT"]);
    let kmers = all_canonical_kmers(dir.path());
    let count_map: HashMap<CanonicalKmer, u32> =
        kmers.iter().enumerate().map(|(i, &k)| (k, i as u32 + 1)).collect();
    Layer::<PersistentCompactIntMatrix>::build(
        dir.path(),
        DEFAULT_BLOCK_BITS,
        &IndexMode::Exact,
        |kmer| count_map.get(&kmer).copied().unwrap_or(0),
    ).unwrap();
    let layer = Layer::<PersistentCompactIntMatrix>::open(dir.path(), &IndexMode::Exact).unwrap();
    for kmer in &kmers {
        let hit = layer.query(*kmer).expect("kmer must be present");
        assert_eq!(hit.data[0], count_map[kmer]);
    }
}

#[test]
fn query_absent_returns_none() {
    set_k(4);
    let dir = tempdir().unwrap();
    write_unitigs(dir.path(), &[b"AAAACGT"]);
    Layer::<()>::build(dir.path(), DEFAULT_BLOCK_BITS, &IndexMode::Exact).unwrap();
    let layer = Layer::<()>::open(dir.path(), &IndexMode::Exact).unwrap();
    let absent = Kmer::from_ascii(b"CCCC").unwrap().canonical();
    assert!(layer.query(absent).is_none());
}

#[test]
fn open_after_build_is_consistent() {
    set_k(4);
    let dir = tempdir().unwrap();
    write_unitigs(dir.path(), &[b"AAAACGT"]);
    let n = Layer::<PersistentCompactIntMatrix>::build(dir.path(), DEFAULT_BLOCK_BITS, &IndexMode::Exact, |_| 7).unwrap();
    assert_eq!(n, 4);
    let layer = Layer::<PersistentCompactIntMatrix>::open(dir.path(), &IndexMode::Exact).unwrap();
    let kmer = Kmer::from_ascii(b"AAAA").unwrap().canonical();
    let hit = layer.query(kmer).expect("AAAA must be present");
    assert_eq!(hit.data[0], 7);
}
