//! Compact 2-bit DNA sequence — shared substrate for [`SuperKmer`] and [`Unitig`].
//!
//! Encoding: A=00, C=01, G=10, T=11. Nucleotide 0 occupies bits 7–6 of `seq[0]`,
//! nucleotide i occupies bits `7−2*(i%4)` and `6−2*(i%4)` of `seq[i/4]`.
//! Padding bits in the last byte are always 0.
//!
//! The exact nucleotide count is recovered without storing it explicitly:
//!
//! ```text
//! seql = (seq.len() - 1) * 4 + tail_count(tail)
//! ```
//!
//! where `tail` encodes the number of valid nucleotides in the last byte (0 → 4).

use std::io::{self, Read, Write};

use bitvec::prelude::*;

use crate::Sequence;
use crate::encoding::{DEC4, encode_base};
use crate::kmer::{CanonicalKmer, Kmer, KmerError, KLen, KmerLength, KmerOf, MLen, Minimizer, RawKmer, KMER_BITS};
use crate::params::k;
use crate::revcomp_lookup::REVCOMP4;

/// Maximum kmers per stored chunk. Enforces the u8 max-kmer-index field in the binary format.
pub const MAX_KMERS_PER_CHUNK: usize = 256;

// ── PackedSeq ─────────────────────────────────────────────────────────────────

/// 2-bit packed DNA sequence of arbitrary length ≥ 1.
///
/// `tail` encodes the number of valid nucleotides in the last byte: 0 stands for 4,
/// so the range 0–3 covers all four cases. Padding bits are always 0.
#[derive(Debug, Clone)]
pub struct PackedSeq {
    pub(crate) tail: u8,
    pub(crate) seq: Box<[u8]>,
}

impl PartialEq for PackedSeq {
    fn eq(&self, other: &Self) -> bool {
        self.tail == other.tail && self.seq == other.seq
    }
}

impl Eq for PackedSeq {}

impl std::hash::Hash for PackedSeq {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.tail.hash(state);
        self.seq.hash(state);
    }
}

impl PackedSeq {
    /// Construct from pre-packed bytes and a `tail` value (0–3, where 0 means 4).
    /// Caller must guarantee padding bits in the last byte are zeroed.
    #[inline]
    pub fn new(tail: u8, seq: Box<[u8]>) -> Self {
        debug_assert!(tail <= 3, "tail must be 0–3");
        debug_assert!(!seq.is_empty(), "seq must be non-empty");
        Self { tail, seq }
    }

    /// Sequence length in nucleotides.
    #[inline]
    pub fn seql(&self) -> usize {
        (self.seq.len() - 1) * 4 + tail_count(self.tail)
    }

    /// Read-only view of the packed 2-bit bytes.
    #[inline]
    pub fn seq_bytes(&self) -> &[u8] {
        &self.seq
    }

    /// Extract nucleotide i (0-based from 5′ end) as a 2-bit value. Zero copy.
    #[inline]
    pub fn nucleotide(&self, i: usize) -> u8 {
        (self.seq[i / 4] >> (6 - 2 * (i % 4))) & 0b11
    }

    /// Encode an ASCII nucleotide slice (ACGT, length ≥ 1). Allocates once.
    pub fn from_ascii(ascii: &[u8]) -> Self {
        let seql = ascii.len();
        debug_assert!(seql >= 1);
        let n = byte_len(seql);
        let mut seq = vec![0u8; n];
        let full = seql / 4;
        for i in 0..full {
            seq[i] = encode_base(ascii[i * 4]) << 6
                | encode_base(ascii[i * 4 + 1]) << 4
                | encode_base(ascii[i * 4 + 2]) << 2
                | encode_base(ascii[i * 4 + 3]);
        }
        let rem = seql % 4;
        if rem > 0 {
            let mut last = 0u8;
            for j in 0..rem {
                last |= encode_base(ascii[full * 4 + j]) << (6 - 2 * j);
            }
            seq[full] = last;
        }
        Self::new(count_to_tail(seql), seq.into_boxed_slice())
    }

    /// Encode a slice of 2-bit nucleotide values (0=A…3=T, length ≥ 1). Allocates once.
    pub fn from_nucleotides(nucs: &[u8]) -> Self {
        let seql = nucs.len();
        debug_assert!(seql >= 1);
        let n = byte_len(seql);
        let mut seq = vec![0u8; n];
        for (i, &nuc) in nucs.iter().enumerate() {
            seq[i / 4] |= (nuc & 0b11) << (6 - 2 * (i % 4));
        }
        Self::new(count_to_tail(seql), seq.into_boxed_slice())
    }
}

impl FromIterator<u8> for PackedSeq {
    /// Build from a stream of 2-bit nucleotide values (0=A…3=T).
    /// Packs on the fly — one allocation, no intermediate Vec.
    fn from_iter<I: IntoIterator<Item = u8>>(iter: I) -> Self {
        let iter = iter.into_iter();
        let (lower, _) = iter.size_hint();
        let mut seq: Vec<u8> = Vec::with_capacity(byte_len(lower.max(1)));
        let mut seql = 0usize;
        for nuc in iter {
            if seql % 4 == 0 {
                seq.push(0);
            }
            *seq.last_mut().unwrap() |= (nuc & 0b11) << (6 - 2 * (seql % 4));
            seql += 1;
        }
        debug_assert!(seql >= 1, "PackedSeq requires at least one nucleotide");
        Self::new(count_to_tail(seql), seq.into_boxed_slice())
    }
}

impl PackedSeq {

    /// Write ASCII nucleotides into `writer`. Zero allocation.
    pub fn write_ascii<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        let seql = self.seql();
        let full = seql / 4;
        for i in 0..full {
            writer.write_all(&DEC4[self.seq[i] as usize].to_be_bytes())?;
        }
        let rem = seql % 4;
        if rem > 0 {
            writer.write_all(&DEC4[self.seq[full] as usize].to_be_bytes()[..rem])?;
        }
        Ok(())
    }

    /// Decode into a fresh ASCII `Vec<u8>`. Allocates.
    #[inline]
    pub fn to_ascii(&self) -> Vec<u8> {
        let mut buf = Vec::with_capacity(self.seql());
        self.write_ascii(&mut buf).unwrap();
        buf
    }

    /// Reverse-complement in place. Zero allocation.
    pub fn revcomp_inplace(&mut self) {
        let seql = self.seql();
        let n = self.seq.len();
        {
            let bytes = &mut self.seq[..n];
            let (mut lo, mut hi) = (0, n - 1);
            while lo < hi {
                (bytes[lo], bytes[hi]) =
                    (REVCOMP4[bytes[hi] as usize], REVCOMP4[bytes[lo] as usize]);
                lo += 1;
                hi -= 1;
            }
            if lo == hi {
                bytes[lo] = REVCOMP4[bytes[lo] as usize];
            }
        }
        let shift = n * 8 - seql * 2;
        if shift > 0 {
            let bits = self.seq[..n].view_bits_mut::<Msb0>();
            bits.rotate_left(shift);
            let len = bits.len();
            bits[len - shift..].fill(false);
        }
        // tail is invariant: seql is unchanged by revcomp
    }

    /// Returns `true` if in canonical form (lexicographic minimum of forward and revcomp).
    pub fn is_canonical(&self) -> bool {
        let seql = self.seql();
        for i in 0..seql {
            let fwd = self.nucleotide(i);
            let rev = complement(self.nucleotide(seql - 1 - i));
            match fwd.cmp(&rev) {
                std::cmp::Ordering::Less => return true,
                std::cmp::Ordering::Greater => return false,
                std::cmp::Ordering::Equal => {}
            }
        }
        true
    }

    /// Put in canonical form in place. Returns `true` if already canonical. Zero allocation.
    #[inline]
    pub fn canonicalize(&mut self) -> bool {
        if self.is_canonical() {
            return true;
        }
        self.revcomp_inplace();
        false
    }

    /// Extract a kmer of length `L::len()` at nucleotide position `i`. Zero allocation.
    fn extract<L: KmerLength>(&self, i: usize) -> Result<KmerOf<L>, KmerError> {
        let len = L::len();
        let seql = self.seql();
        if i + len > seql {
            return Err(KmerError::OutOfBounds { position: i, k: len, seql });
        }
        let bits = self.seq.view_bits::<Msb0>();
        let raw: RawKmer = bits[i * 2..(i + len) * 2].load_be();
        Ok(KmerOf::from_raw(raw << (KMER_BITS - 2 * len)))
    }

    /// Extract the kmer of length `params::k()` at nucleotide position `i`. Zero allocation.
    #[inline]
    pub fn kmer(&self, i: usize) -> Result<Kmer, KmerError> {
        self.extract::<KLen>(i)
    }

    /// Extract the canonical m-mer (minimizer) of length `params::m()` at position `i`. Zero allocation.
    #[inline]
    pub fn mmer(&self, i: usize) -> Result<Minimizer, KmerError> {
        Ok(self.extract::<MLen>(i)?.canonical())
    }

    /// Extract the canonical kmer of length `params::k()` at position `i`. Zero allocation.
    #[inline]
    pub fn canonical_kmer(&self, i: usize) -> Result<CanonicalKmer, KmerError> {
        Ok(self.kmer(i)?.canonical())
    }

    /// Borrowing iterator over all k-mers of length `params::k()` in order.
    ///
    /// Borrows `self`; use when the sequence is held in a collection and may be
    /// iterated more than once, or when the sequence is needed after iteration.
    /// Cannot be used inside a `flat_map` closure that receives owned values —
    /// use [`into_kmers`](Self::into_kmers) in that case.
    #[inline]
    pub fn iter_kmers(&self) -> PackedSeqKmerIter<&PackedSeq> {
        PackedSeqKmerIter::new(self)
    }

    /// Borrowing iterator over all canonical k-mers of length `params::k()` in order.
    ///
    /// Same borrowing semantics as [`iter_kmers`](Self::iter_kmers).
    #[inline]
    pub fn iter_canonical_kmers(&self) -> impl Iterator<Item = CanonicalKmer> + '_ {
        self.iter_kmers().map(|km| km.canonical())
    }

    /// Consuming iterator over all k-mers of length `params::k()` in order.
    ///
    /// Moves `self` into the iterator ([`OwnedPackedSeqKmerIter`]); use when
    /// the sequence is produced one-at-a-time and ownership must transfer into
    /// the iterator — in particular inside a `flat_map` closure that receives
    /// owned `PackedSeq` or `SuperKmer` values.
    #[inline]
    pub fn into_kmers(self) -> OwnedPackedSeqKmerIter {
        OwnedPackedSeqKmerIter::new(self)
    }

    /// Consuming iterator over all canonical k-mers of length `params::k()` in order.
    ///
    /// Same ownership semantics as [`into_kmers`](Self::into_kmers).
    #[inline]
    pub fn into_canonical_kmers(self) -> impl Iterator<Item = CanonicalKmer> {
        self.into_kmers().map(|km| km.canonical())
    }

    /// Extract nucleotides `[start, end)` as a new [`PackedSeq`]. Allocates.
    pub fn sub(&self, start: usize, end: usize) -> Self {
        debug_assert!(end > start && end <= self.seql());
        let seql = end - start;
        let n = (seql + 3) / 4;
        let mut out = vec![0u8; n];
        let src = self.seq.view_bits::<Msb0>();
        let dst = out.view_bits_mut::<Msb0>();
        dst[..seql * 2].copy_from_bitslice(&src[start * 2..end * 2]);
        Self::new(count_to_tail(seql), out.into_boxed_slice())
    }

    /// Serialise one chunk to binary.
    ///
    /// Format: `[u8: n_kmers−1][packed bytes]`.
    /// The caller must ensure `seql ≥ k` and `seql − k + 1 ≤ MAX_KMERS_PER_CHUNK`.
    /// Use [`SuperKmer::write_to_binary`] for sequences that may exceed one chunk.
    pub fn write_to_binary<W: Write>(&self, w: &mut W) -> io::Result<()> {
        let k = crate::params::k();
        let seql = self.seql();
        debug_assert!(seql >= k, "sequence shorter than k");
        debug_assert!(
            seql - k + 1 <= MAX_KMERS_PER_CHUNK,
            "chunk exceeds MAX_KMERS_PER_CHUNK; split before calling write_to_binary"
        );
        w.write_all(&[(seql - k) as u8])?;
        w.write_all(&self.seq)
    }

    /// Deserialise one chunk from the binary format produced by [`write_to_binary`].
    /// Allocates exactly one `Box<[u8]>` for the packed bytes.
    pub fn read_from_binary<R: Read>(r: &mut R) -> io::Result<Self> {
        let mut buf = [0u8; 1];
        r.read_exact(&mut buf)?;
        let seql = buf[0] as usize + crate::params::k();
        let byte_len = (seql + 3) / 4;
        let tail = (seql % 4) as u8;
        let mut seq = vec![0u8; byte_len];
        r.read_exact(&mut seq)?;
        Ok(Self::new(tail, seq.into_boxed_slice()))
    }
}

impl AsRef<PackedSeq> for PackedSeq {
    fn as_ref(&self) -> &PackedSeq { self }
}

// ── PackedSeqKmerIter ─────────────────────────────────────────────────────────

/// Sliding-window k-mer iterator over a [`PackedSeq`]. Zero allocation.
///
/// The type parameter `S` controls ownership of the underlying sequence:
///
/// | `S`           | Alias                   | Obtained via          | Use when                                              |
/// |---------------|-------------------------|-----------------------|-------------------------------------------------------|
/// | `&PackedSeq`  | *(none, inferred)*      | [`PackedSeq::iter_kmers`]   | The `PackedSeq` is borrowed and will be reused or iterated multiple times. |
/// | `PackedSeq`   | [`OwnedPackedSeqKmerIter`] | [`PackedSeq::into_kmers`] | The `PackedSeq` is consumed and will not be needed after iteration — required when the sequence is produced one-at-a-time (e.g. read from a file) and passed to `flat_map`. |
///
/// The borrowing form cannot be used inside a `flat_map` closure that receives
/// owned values: the sequence would be dropped at the end of the closure while
/// the iterator still holds a reference to it.  The owning form moves the
/// sequence into the iterator, so no reference escapes the closure.
pub struct PackedSeqKmerIter<S> {
    seq: S,
    mask: RawKmer,
    lshift: usize,
    current: RawKmer,
    pos: usize,
    max_pos: usize,
}

/// Owning variant of [`PackedSeqKmerIter`]: the [`PackedSeq`] is moved into
/// the iterator rather than borrowed.  Use this when the sequence is consumed
/// (e.g. via [`PackedSeq::into_kmers`] or [`SuperKmer::into_canonical_kmers`])
/// so that the iterator can be stored or passed to `flat_map` without lifetime
/// constraints.
pub type OwnedPackedSeqKmerIter = PackedSeqKmerIter<PackedSeq>;

impl<S: AsRef<PackedSeq>> PackedSeqKmerIter<S> {
    fn new(seq: S) -> Self {
        let ps = seq.as_ref();
        let seql = ps.seql();
        let klen = k();
        let lshift = KMER_BITS - klen * 2;
        let mask = ((!0u128) << (lshift + 2)) as RawKmer;
        let current = if seql >= klen {
            ps.extract::<KLen>(0).map(|km| km.raw()).unwrap_or(0)
        } else {
            0
        };
        Self { seq, mask, lshift, current, pos: klen, max_pos: seql }
    }
}

impl<S: AsRef<PackedSeq>> Iterator for PackedSeqKmerIter<S> {
    type Item = Kmer;

    fn next(&mut self) -> Option<Kmer> {
        if self.pos > self.max_pos {
            return None;
        }
        let result = Kmer::from_raw(self.current);
        if self.pos < self.max_pos {
            let inner_shift = 6 - 2 * (self.pos & 3);
            let nuc = ((self.seq.as_ref().seq[self.pos / 4] >> inner_shift) & 3) as RawKmer;
            self.current = ((self.current << 2) & self.mask) | (nuc << self.lshift);
        }
        self.pos += 1;
        Some(result)
    }
}

// ── varint (LEB128) ───────────────────────────────────────────────────────────

pub(crate) fn write_varint<W: Write>(w: &mut W, mut val: u64) -> io::Result<()> {
    loop {
        let mut byte = (val & 0x7F) as u8;
        val >>= 7;
        if val != 0 {
            byte |= 0x80;
        }
        w.write_all(&[byte])?;
        if val == 0 {
            break;
        }
    }
    Ok(())
}

pub(crate) fn read_varint<R: Read>(r: &mut R) -> io::Result<u64> {
    let mut val = 0u64;
    let mut shift = 0u32;
    let mut buf = [0u8; 1];
    loop {
        r.read_exact(&mut buf)?;
        let byte = buf[0];
        val |= ((byte & 0x7F) as u64) << shift;
        if byte & 0x80 == 0 {
            break;
        }
        shift += 7;
        if shift >= 64 {
            return Err(io::Error::new(io::ErrorKind::InvalidData, "varint overflow"));
        }
    }
    Ok(val)
}

// ── helpers ───────────────────────────────────────────────────────────────────

#[inline]
fn complement(base: u8) -> u8 {
    !base & 0b11
}

#[inline]
fn byte_len(seql: usize) -> usize {
    (seql + 3) / 4
}

/// Nucleotide count → `tail` value: 0 encodes 4, 1–3 are identity.
#[inline]
pub(crate) fn count_to_tail(seql: usize) -> u8 {
    (seql % 4) as u8
}

/// `tail` value → nucleotide count in last byte: 0 means 4.
#[inline]
pub(crate) fn tail_count(tail: u8) -> usize {
    if tail == 0 { 4 } else { tail as usize }
}
