//! Compact 2-bit kmer stored as a left-aligned [`RawKmer`].
//!
//! Nucleotide 0 occupies bits KMER_BITS−1 and KMER_BITS−2, nucleotide i occupies
//! bits KMER_BITS−1−2i and KMER_BITS−2−2i.
//! The low `KMER_BITS − 2·len` bits are always zero.
//!
//! The length is not stored in the struct — it is supplied by the [`KmerLength`]
//! type parameter.  Two public marker types cover the common cases:
//!
//! | Alias           | Marker   | Length source  |
//! |-----------------|----------|----------------|
//! | [`Kmer`]        | [`KLen`] | `params::k()`  |
//! | [`CanonicalKmer`]| [`KLen`]| `params::k()`  |
//! | [`Minimizer`]   | [`MLen`] | `params::m()`  |
//!
//! Tests that need a fixed length independent of the global singletons can use
//! [`ConstLen<N>`].

/// Underlying storage type for a raw (2-bit-per-base, left-aligned) kmer.
///
/// All kmer bit-width constants (`KMER_BITS`) and shift calculations are derived
/// from this type. To extend the maximum supported k beyond 32, change this alias
/// to `u128` and update `KMER_BITS` — callers that use `RawKmer` directly will
/// then pick up the wider type automatically.
pub type RawKmer = u64;

/// Bit width of [`RawKmer`]: the number of bits available for nucleotide encoding.
pub const KMER_BITS: usize = 64;

use serde::Serialize;
use std::io::{self, Write};
use std::marker::PhantomData;

use crate::Annotation;
use crate::Sequence;
use crate::encoding::{DEC4, encode_base};
use crate::params::{k, m};
use crate::sequence::mix64;

// ── KmerLength ────────────────────────────────────────────────────────────────

/// Marker trait that supplies a kmer length at runtime.
pub trait KmerLength: Copy + std::fmt::Debug + 'static {
    /// Returns the length this marker represents.
    fn len() -> usize;
}

/// Marker for the k-mer length (`params::k()`).
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct KLen;

/// Marker for the minimizer length (`params::m()`).
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct MLen;

/// Marker for a compile-time-constant length — useful for tests.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ConstLen<const N: usize>;

impl KmerLength for KLen {
    #[inline]
    fn len() -> usize {
        k()
    }
}

impl KmerLength for MLen {
    #[inline]
    fn len() -> usize {
        m()
    }
}

impl<const N: usize> KmerLength for ConstLen<N> {
    #[inline]
    fn len() -> usize {
        N
    }
}

// ── KmerError ─────────────────────────────────────────────────────────────────

/// Errors produced by kmer operations.
#[derive(Debug, Clone, PartialEq)]
pub enum KmerError {
    /// The requested kmer extends past the end of the sequence.
    OutOfBounds {
        /// Start position of the requested kmer.
        position: usize,
        /// Requested kmer length.
        k: usize,
        /// Actual sequence length.
        seql: usize,
    },
    /// k is zero or exceeds 32 (would not fit in a u64).
    InvalidK {
        /// The invalid k value.
        k: usize,
    },
}

impl std::fmt::Display for KmerError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            KmerError::OutOfBounds { position, k, seql } => write!(
                f,
                "kmer of length {k} at position {position} exceeds sequence length {seql}"
            ),
            KmerError::InvalidK { k } => write!(f, "k={k} is invalid: must be in 1..=32"),
        }
    }
}

impl std::error::Error for KmerError {}

#[derive(Serialize)]
struct KmerAnnotation {
    seq_length: usize,
}
impl Annotation for KmerAnnotation {}

// ── KmerOf ────────────────────────────────────────────────────────────────────

/// A DNA kmer of length `L::len()` encoded as a left-aligned [`RawKmer`] (2 bits/nucleotide, MSB-first).
/// The low `KMER_BITS − 2·L::len()` bits are always zero.
#[repr(transparent)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct KmerOf<L: KmerLength>(RawKmer, PhantomData<L>);

impl<L: KmerLength> KmerOf<L> {
    /// Wrap a raw left-aligned [`RawKmer`] value as a kmer.
    #[inline]
    pub const fn from_raw(raw: RawKmer) -> Self {
        KmerOf(raw, PhantomData)
    }

    /// Wrap a raw right-aligned [`RawKmer`] value, shifting it into left-aligned position.
    #[inline]
    pub fn from_raw_right(raw: RawKmer) -> Self {
        KmerOf(raw << (KMER_BITS - 2 * L::len()), PhantomData)
    }

    /// Return the raw left-aligned [`RawKmer`] value.
    #[inline]
    pub fn raw(&self) -> RawKmer {
        self.0
    }

    /// Return the raw right-aligned [`RawKmer`] value.
    #[inline]
    pub fn raw_right(&self) -> RawKmer {
        self.0 >> (KMER_BITS - 2 * L::len())
    }

    /// Encode the first `L::len()` nucleotides of an ASCII slice into a kmer.
    pub fn from_ascii(ascii: &[u8]) -> Result<Self, KmerError> {
        let k = L::len();
        if k == 0 || k > 32 {
            return Err(KmerError::InvalidK { k });
        }
        if ascii.len() < k {
            return Err(KmerError::OutOfBounds {
                position: 0,
                k,
                seql: ascii.len(),
            });
        }
        let mut val: RawKmer = 0;
        for i in 0..k {
            val = (val << 2) | encode_base(ascii[i]) as RawKmer;
        }
        Ok(KmerOf(val << (KMER_BITS - 2 * k), PhantomData))
    }

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

    /// Decode into ASCII nucleotides, writing into `writer`.
    #[inline]
    pub fn write_ascii<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        let k = L::len();
        let bytes = self.0.to_be_bytes();
        let full = k / 4;
        let rem = k % 4;
        for i in 0..full {
            writer.write_all(&DEC4[bytes[i] as usize].to_be_bytes())?;
        }
        if rem > 0 {
            let decoded = DEC4[bytes[full] as usize].to_be_bytes();
            writer.write_all(&decoded[..rem])?;
        }
        Ok(())
    }

    /// Compute the reverse complement.
    #[inline]
    pub fn revcomp(&self) -> Self {
        let k = L::len();
        let x = !self.0;
        let x = x.swap_bytes();
        let x = ((x >> 4) & 0x0F0F0F0F0F0F0F0F) | ((x & 0x0F0F0F0F0F0F0F0F) << 4);
        let x = ((x >> 2) & 0x3333333333333333) | ((x & 0x3333333333333333) << 2);
        KmerOf(x << (KMER_BITS - 2 * k), PhantomData)
    }

    /// Slide the window one base to the right: drop nucleotide 0, append `nuc` at position `L::len()-1`.
    pub fn push_right(self, nuc: u8) -> Self {
        let k = L::len();
        let shifted = self.0 << 2 & (RawKmer::MAX << (KMER_BITS - 2 * (k - 1)));
        KmerOf(shifted | ((nuc as RawKmer & 3) << (KMER_BITS - 2 * k)), PhantomData)
    }

    /// Slide the window one base to the left: drop nucleotide `L::len()-1`, prepend `nuc` at position 0.
    pub fn push_left(self, nuc: u8) -> Self {
        let k = L::len();
        let shifted = (self.0 >> 2) & (RawKmer::MAX << (KMER_BITS - 2 * k));
        KmerOf(shifted | ((nuc as RawKmer & 3) << (KMER_BITS - 2)), PhantomData)
    }

    /// Returns `true` if the last `L::len()-1` nucleotides of `self` equal the first `L::len()-1` of `other`.
    pub fn is_overlapping(self, other: Self) -> bool {
        let k = L::len();
        let mask = RawKmer::MAX << (KMER_BITS - 2 * (k - 1));
        (self.0 << 2 & mask) == (other.0 & mask)
    }
}

impl<L: KmerLength> Sequence for KmerOf<L> {
    type Canonical = CanonicalKmerOf<L>;

    fn seql(&self) -> usize {
        L::len()
    }

    fn seq_hash(&self) -> u64 {
        self.canonical().seq_hash()
    }

    #[inline]
    fn nucleotide(&self, i: usize) -> u8 {
        ((self.0 >> (KMER_BITS - 2 - 2 * i)) & 0b11) as u8
    }

    #[inline]
    fn canonical(&self) -> Self::Canonical {
        let rc = self.revcomp();
        CanonicalKmerOf(if self.0 <= rc.0 { self.0 } else { rc.0 }, PhantomData)
    }

    fn annotation(&self) -> impl Annotation {
        KmerAnnotation {
            seq_length: L::len(),
        }
    }
}
// ── CanonicalKmerOf ───────────────────────────────────────────────────────────

/// A [`KmerOf<L>`] guaranteed to be in canonical form (lexicographic minimum of
/// forward and reverse complement).
///
/// The only public constructors are [`KmerOf::canonical`] (verified) and
/// [`CanonicalKmerOf::from_raw_unchecked`] (trusted paths such as deserialisation).
#[repr(transparent)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct CanonicalKmerOf<L: KmerLength>(RawKmer, PhantomData<L>);

impl<L: KmerLength> CanonicalKmerOf<L> {
    /// Wrap a raw left-aligned [`RawKmer`] without verifying the canonical invariant.
    ///
    /// # Safety (logical)
    /// The caller must guarantee `raw == min(raw, revcomp(raw))`.
    #[inline]
    pub fn from_raw_unchecked(raw: RawKmer) -> Self {
        CanonicalKmerOf(raw, PhantomData)
    }

    /// Return the raw left-aligned [`RawKmer`] value.
    #[inline]
    pub fn raw(&self) -> RawKmer {
        self.0
    }

    /// Decode into a freshly allocated ASCII `Vec<u8>`.
    #[inline]
    pub fn to_ascii(&self) -> Vec<u8> {
        self.into_kmer().to_ascii()
    }

    /// Decode into ASCII nucleotides, writing into `writer`.
    #[inline]
    pub fn write_ascii<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        self.into_kmer().write_ascii(writer)
    }

    /// Compute the reverse complement. The result is a raw [`KmerOf<L>`] — not
    /// necessarily canonical itself.
    #[inline]
    pub fn revcomp(&self) -> KmerOf<L> {
        self.into_kmer().revcomp()
    }

    /// Hash via `mix64`.
    #[inline]
    pub fn seq_hash(&self) -> u64 {
        hash_kmer(self.0)
    }

    /// Extract nucleotide i (0-based from 5′ end) as a 2-bit value.
    #[inline]
    pub fn nucleotide(&self, i: usize) -> u8 {
        self.into_kmer().nucleotide(i)
    }

    /// Return the four left canonical neighbours (each already canonical).
    pub fn left_canonical_neighbors(&self) -> [CanonicalKmerOf<L>; 4] {
        let k = L::len();
        let shifted = (self.0 >> 2) & (RawKmer::MAX << (KMER_BITS - 2 * k));
        [
            KmerOf::<L>(shifted, PhantomData).canonical(),
            KmerOf::<L>(shifted | ((1 as RawKmer) << (KMER_BITS - 2)), PhantomData).canonical(),
            KmerOf::<L>(shifted | ((2 as RawKmer) << (KMER_BITS - 2)), PhantomData).canonical(),
            KmerOf::<L>(shifted | ((3 as RawKmer) << (KMER_BITS - 2)), PhantomData).canonical(),
        ]
    }

    /// Return the four right canonical neighbours (each already canonical).
    pub fn right_canonical_neighbors(&self) -> [CanonicalKmerOf<L>; 4] {
        let k = L::len();
        let shifted = self.0 << 2 & (RawKmer::MAX << (KMER_BITS - 2 * (k - 1)));
        let shift = KMER_BITS - 2 * k;
        [
            KmerOf::<L>(shifted, PhantomData).canonical(),
            KmerOf::<L>(shifted | ((1 as RawKmer) << shift), PhantomData).canonical(),
            KmerOf::<L>(shifted | ((2 as RawKmer) << shift), PhantomData).canonical(),
            KmerOf::<L>(shifted | ((3 as RawKmer) << shift), PhantomData).canonical(),
        ]
    }

    /// Return the inner value as a raw [`KmerOf<L>`].
    #[inline]
    pub fn into_kmer(self) -> KmerOf<L> {
        KmerOf(self.0, PhantomData)
    }
}

impl<L: KmerLength> Sequence for CanonicalKmerOf<L> {
    type Canonical = CanonicalKmerOf<L>;

    fn seql(&self) -> usize {
        L::len()
    }

    fn seq_hash(&self) -> u64 {
        hash_kmer(self.0)
    }

    #[inline]
    fn nucleotide(&self, i: usize) -> u8 {
        ((self.0 >> (KMER_BITS - 2 - 2 * i)) & 0b11) as u8
    }

    fn canonical(&self) -> Self::Canonical {
        *self
    }

    fn annotation(&self) -> impl Annotation {
        KmerAnnotation {
            seq_length: L::len(),
        }
    }
}
impl<L: KmerLength> From<CanonicalKmerOf<L>> for KmerOf<L> {
    #[inline]
    fn from(ck: CanonicalKmerOf<L>) -> Self {
        ck.into_kmer()
    }
}

// ── Public type aliases ───────────────────────────────────────────────────────

/// A DNA k-mer using the global `params::k()` length.
pub type Kmer = KmerOf<KLen>;

/// A canonical k-mer using the global `params::k()` length.
pub type CanonicalKmer = CanonicalKmerOf<KLen>;

/// A minimizer: a canonical k-mer using the global `params::m()` length.
pub type Minimizer = CanonicalKmerOf<MLen>;

/// Compute a hash for a raw (left-aligned) kmer value.
///
/// This is a convenience wrapper around [`mix64`] that accepts the raw
/// 64-bit representation directly, which is useful when the canonical
/// invariant is not required or has already been handled.
#[inline]
pub fn hash_kmer(raw: RawKmer) -> u64 {
    mix64(raw ^ 0x9e3779b97f4a7c15)
}

// ── tests ─────────────────────────────────────────────────────────────────────
#[cfg(test)]
#[path = "tests/kmer.rs"]
mod tests;
