use super::*;
use crate::set_k;

fn make_seq(len: usize) -> Vec<u8> {
    (0..len).map(|i| b"ACGT"[i % 4]).collect()
}

fn ascii_revcomp(seq: &[u8]) -> Vec<u8> {
    seq.iter()
        .rev()
        .map(|&b| match b {
            b'A' => b'T',
            b'T' => b'A',
            b'C' => b'G',
            b'G' => b'C',
            _ => b'A',
        })
        .collect()
}

fn all_lengths() -> impl Iterator<Item = usize> {
    (1..=9).chain([255, 256, 257, 1000])
}

// ── from_ascii / canonical form ───────────────────────────────────────────────

#[test]
fn ascii_roundtrip_all_lengths() {
    for len in all_lengths() {
        let ascii = make_seq(len);
        let sk = SuperKmer::from_ascii(&ascii);
        // SuperKmer stores in canonical form; ACGT pattern is already canonical.
        assert_eq!(sk.to_ascii(), ascii, "roundtrip failed for len={len}");
    }
}

#[test]
fn from_ascii_canonical_all_bases() {
    // G×4 revcomp is C×4; T×4 revcomp is A×4.
    for (base, expected) in [(b'A', b'A'), (b'C', b'C'), (b'G', b'C'), (b'T', b'A')] {
        let ascii = vec![base; 4];
        let sk = SuperKmer::from_ascii(&ascii);
        assert_eq!(sk.to_ascii(), vec![expected; 4]);
    }
}

#[test]
fn from_ascii_is_canonical_all_lengths() {
    for len in all_lengths() {
        let ascii = make_seq(len);
        let sk = SuperKmer::from_ascii(&ascii);
        let fwd = sk.to_ascii();
        let rev = ascii_revcomp(&fwd);
        assert!(fwd <= rev, "not canonical for len={len}");
    }
}

// ── seql ──────────────────────────────────────────────────────────────────────

#[test]
fn seql_roundtrip() {
    for len in all_lengths() {
        let sk = SuperKmer::from_ascii(&make_seq(len));
        assert_eq!(sk.seql(), len, "seql() wrong for len={len}");
    }
}

// ── binary serialisation ──────────────────────────────────────────────────────

fn binary_test_lengths(k: usize) -> Vec<usize> {
    use crate::packed_seq::MAX_KMERS_PER_CHUNK;
    // Only single-chunk lengths: seql in [k, MAX_KMERS_PER_CHUNK+k-1].
    (k..=k + 5).chain([255, 256, 257, MAX_KMERS_PER_CHUNK + k - 1]).collect()
}

#[test]
fn binary_roundtrip() {
    set_k(4);
    let k = crate::params::k();
    for len in binary_test_lengths(k) {
        let mut sk = SuperKmer::from_ascii(&make_seq(len));
        sk.set_count(42);
        let mut buf = Vec::new();
        sk.write_to_binary(&mut buf).unwrap();
        let sk2 = SuperKmer::read_from_binary(&mut buf.as_slice()).unwrap();
        assert_eq!(sk.to_ascii(), sk2.to_ascii(), "sequence mismatch for len={len}");
        assert_eq!(sk2.count(), 42, "count mismatch for len={len}");
    }
}

#[test]
fn binary_split_roundtrip() {
    // A super-kmer > MAX_KMERS_PER_CHUNK kmers is split into multiple records on write.
    use crate::packed_seq::MAX_KMERS_PER_CHUNK;
    set_k(4);
    let k = crate::params::k();
    // seql = MAX_KMERS_PER_CHUNK + k = 260 → n_kmers = 257 > 256 → 2 chunks
    let seql = MAX_KMERS_PER_CHUNK + k;
    let mut sk = SuperKmer::from_ascii(&make_seq(seql));
    sk.set_count(7);
    let mut buf = Vec::new();
    sk.write_to_binary(&mut buf).unwrap();
    // Read all records back.
    let mut slice = buf.as_slice();
    let chunk0 = SuperKmer::read_from_binary(&mut slice).unwrap();
    let chunk1 = SuperKmer::read_from_binary(&mut slice).unwrap();
    assert!(slice.is_empty(), "unexpected trailing bytes");
    assert_eq!(chunk0.count(), 7);
    assert_eq!(chunk1.count(), 7);
    // Chunks cover the original sequence with k-1 overlap — no kmer lost.
    assert_eq!(chunk0.seql(), MAX_KMERS_PER_CHUNK + k - 1); // 259
    assert_eq!(chunk1.seql(), k);                            // 4 (1 kmer)
}

#[test]
fn binary_packed_seq_roundtrip() {
    use crate::packed_seq::PackedSeq;
    set_k(4);
    let k = crate::params::k();
    for len in binary_test_lengths(k) {
        let ps = PackedSeq::from_ascii(&make_seq(len));
        let mut buf = Vec::new();
        ps.write_to_binary(&mut buf).unwrap();
        let ps2 = PackedSeq::read_from_binary(&mut buf.as_slice()).unwrap();
        assert_eq!(ps, ps2, "PackedSeq mismatch for len={len}");
    }
}

#[test]
fn binary_size_is_compact() {
    // ACGT with k=4: varint(count=1, 1 byte) + u8(n_kmers-1=0, 1 byte) + 1 packed byte = 3 bytes
    set_k(4);
    let sk = SuperKmer::from_ascii(b"ACGT");
    let mut buf = Vec::new();
    sk.write_to_binary(&mut buf).unwrap();
    assert_eq!(buf.len(), 3, "expected 3 bytes for 4-nt superkmer");
}

// ── count ─────────────────────────────────────────────────────────────────────

#[test]
fn count_starts_at_one() {
    let sk = SuperKmer::from_ascii(b"ACGT");
    assert_eq!(sk.count(), 1);
}

#[test]
fn increment_adds_one() {
    let mut sk = SuperKmer::from_ascii(b"ACGT");
    sk.increment();
    assert_eq!(sk.count(), 2);
    sk.increment();
    assert_eq!(sk.count(), 3);
}

#[test]
fn add_increases_count() {
    let mut sk = SuperKmer::from_ascii(b"ACGT");
    sk.add(42);
    assert_eq!(sk.count(), 43);
    sk.add(8);
    assert_eq!(sk.count(), 51);
}

#[test]
fn set_count_overwrites() {
    let mut sk = SuperKmer::from_ascii(b"ACGT");
    sk.add(100);
    sk.set_count(7);
    assert_eq!(sk.count(), 7);
}

#[test]
fn count_operations_preserve_seql() {
    for len in all_lengths() {
        let mut sk = SuperKmer::from_ascii(&make_seq(len));
        sk.increment();
        assert_eq!(sk.seql(), len, "increment altered seql for len={len}");
        sk.add(1000);
        assert_eq!(sk.seql(), len, "add altered seql for len={len}");
        sk.set_count(999);
        assert_eq!(sk.seql(), len, "set_count altered seql for len={len}");
    }
}

#[test]
fn count_does_not_affect_sequence() {
    let ascii = b"ACGTACGT".to_vec();
    let mut sk = SuperKmer::from_ascii(&ascii);
    sk.set_count(16_000_000);
    assert_eq!(sk.to_ascii(), ascii);
}

// ── kmer extraction ───────────────────────────────────────────────────────────

#[test]
fn kmer_first_matches_from_ascii() {
    set_k(4);
    let k = crate::params::k();
    let ascii = b"ACGTACGT";
    let sk = SuperKmer::from_ascii(ascii);
    let kmer = sk.kmer(0).unwrap();
    let expected = crate::kmer::Kmer::from_ascii(&ascii[..k]).unwrap();
    assert_eq!(kmer, expected);
}

#[test]
fn kmer_all_positions() {
    set_k(4);
    let k = crate::params::k();
    let ascii = b"ACGTACGTACGT";
    let sk = SuperKmer::from_ascii(ascii);
    for i in 0..=ascii.len() - k {
        let kmer = sk.kmer(i).unwrap();
        let expected = crate::kmer::Kmer::from_ascii(&ascii[i..i + k]).unwrap();
        assert_eq!(kmer, expected, "mismatch at position {i}");
    }
}

#[test]
fn kmer_out_of_bounds() {
    set_k(4);
    let sk = SuperKmer::from_ascii(b"ACGT"); // seql=4, k=4
    assert!(sk.kmer(1).is_err()); // 1 + 4 > 4
}

// ── canonical_kmer ────────────────────────────────────────────────────────────

#[test]
fn canonical_kmer_is_min_of_kmer_and_revcomp() {
    set_k(4);
    let k = crate::params::k();
    let sk = SuperKmer::from_ascii(b"ACGTACGT");
    for i in 0..=(sk.seql() - k) {
        let ck = sk.canonical_kmer(i).unwrap();
        let fwd = sk.kmer(i).unwrap();
        assert_eq!(ck, fwd.canonical());
    }
}

#[test]
fn canonical_kmer_palindrome_unchanged() {
    set_k(4);
    let sk = SuperKmer::from_ascii(b"ACGT"); // ACGT is its own revcomp
    let ck = sk.canonical_kmer(0).unwrap();
    let fwd = sk.kmer(0).unwrap();
    assert_eq!(ck.into_kmer(), fwd);
}

#[test]
fn canonical_kmer_errors_propagate() {
    set_k(4);
    let sk = SuperKmer::from_ascii(b"ACGT");
    assert!(sk.canonical_kmer(1).is_err()); // out of bounds: 1 + 4 > 4
}

// ── iter_kmers ────────────────────────────────────────────────────────────────

#[test]
fn iter_kmers_count() {
    set_k(4);
    let k = crate::params::k();
    let ascii = b"ACGTACGTACGT";
    let sk = SuperKmer::from_ascii(ascii);
    assert_eq!(sk.iter_kmers().count(), ascii.len() - k + 1);
}

#[test]
fn iter_kmers_first_is_kmer_0() {
    set_k(4);
    let ascii = b"ACGTACGT";
    let sk = SuperKmer::from_ascii(ascii);
    let first = sk.iter_kmers().next().unwrap();
    assert_eq!(first, sk.kmer(0).unwrap());
}

#[test]
fn iter_kmers_matches_kmer_at_each_position() {
    set_k(4);
    let ascii = b"ACGTACGTACGT";
    let sk = SuperKmer::from_ascii(ascii);
    let kmers: Vec<crate::kmer::Kmer> = sk.iter_kmers().collect();
    for (i, &km) in kmers.iter().enumerate() {
        assert_eq!(km, sk.kmer(i).unwrap(), "mismatch at pos {i}");
    }
}

#[test]
fn iter_kmers_single_when_seql_eq_k() {
    set_k(4);
    let k = crate::params::k();
    let ascii = make_seq(k);
    let sk = SuperKmer::from_ascii(&ascii);
    assert_eq!(sk.iter_kmers().count(), 1);
    assert_eq!(sk.iter_kmers().next().unwrap(), sk.kmer(0).unwrap());
}

#[test]
fn iter_kmers_crosses_byte_boundary() {
    set_k(4);
    let ascii = b"ACGTACGTACGT";
    let sk = SuperKmer::from_ascii(ascii);
    let kmers: Vec<crate::kmer::Kmer> = sk.iter_kmers().collect();
    for boundary in [3usize, 4, 7, 8] {
        if boundary + 1 < kmers.len() {
            assert_eq!(
                kmers[boundary],
                sk.kmer(boundary).unwrap(),
                "pos={boundary}"
            );
        }
    }
}

#[test]
fn iter_kmers_long_sequence() {
    set_k(4);
    let k = crate::params::k();
    let ascii = make_seq(200);
    let sk = SuperKmer::from_ascii(&ascii);
    let kmers: Vec<crate::kmer::Kmer> = sk.iter_kmers().collect();
    assert_eq!(kmers.len(), 200 - k + 1);
    for (i, &km) in kmers.iter().enumerate() {
        assert_eq!(km, sk.kmer(i).unwrap(), "pos={i}");
    }
}
