//! Canonical 2-bit DNA super-kmer with occurrence count.
//!
//! Delegates all sequence operations to [`PackedSeq`].
//!
//! On-disk header word (32 bits): `(count << 2) | tail` — 30-bit count, 2-bit tail.

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

use serde::Serialize;
use xxhash_rust::xxh3::xxh3_64;

use crate::Annotation;
use crate::Sequence;
use crate::kmer::{CanonicalKmer, Kmer, KmerError};
use crate::packed_seq::{MAX_KMERS_PER_CHUNK, PackedSeq, read_varint, write_varint};

// ── SKAnnotation ──────────────────────────────────────────────────────────────

#[derive(Serialize)]
struct SKAnnotation {
    seq_length: usize,
    count: u32,
}

impl Annotation for SKAnnotation {}

// ── SuperKmer ─────────────────────────────────────────────────────────────────

/// Canonical super-kmer: occurrence count + 2-bit packed DNA sequence.
///
/// Always stored in canonical form (lex min of forward and revcomp).
/// `PartialEq`/`Hash` compare only sequence content, ignoring count.
#[derive(Debug, Clone)]
pub struct SuperKmer {
    pub(crate) count: u32,
    pub(crate) inner: PackedSeq,
}

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

impl Eq for SuperKmer {}

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

impl Sequence for SuperKmer {
    type Canonical = SuperKmer;

    fn seql(&self) -> usize {
        self.inner.seql()
    }

    fn nucleotide(&self, i: usize) -> u8 {
        self.inner.nucleotide(i)
    }

    fn seq_hash(&self) -> u64 {
        xxh3_64(self.inner.seq_bytes())
    }

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

    fn annotation(&self) -> impl Annotation {
        SKAnnotation {
            seq_length: self.inner.seql(),
            count: self.count,
        }
    }
}

impl SuperKmer {
    /// Encode ASCII nucleotides (length ≥ 1) into a canonical SuperKmer.
    pub fn from_ascii(ascii: &[u8]) -> Self {
        let mut inner = PackedSeq::from_ascii(ascii);
        inner.canonicalize();
        Self { count: 1, inner }
    }

    /// Wrap a pre-built [`PackedSeq`], canonicalising in place.
    pub fn build(mut inner: PackedSeq) -> Self {
        inner.canonicalize();
        Self { count: 1, inner }
    }

    /// Serialise to compact binary: `[varint(count)][u8: n_kmers−1][packed bytes]` per chunk.
    ///
    /// Sequences with more than [`MAX_KMERS_PER_CHUNK`] kmers are transparently split into
    /// overlapping chunks (k−1 nucleotide overlap, same count per chunk). Each chunk is an
    /// independent, self-contained record — one [`read_from_binary`] call reads exactly one.
    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, "super-kmer shorter than k");
        let n_kmers = seql - k + 1;
        if n_kmers <= MAX_KMERS_PER_CHUNK {
            write_varint(w, self.count as u64)?;
            self.inner.write_to_binary(w)
        } else {
            let chunk_nucl = MAX_KMERS_PER_CHUNK + k - 1;
            let stride = MAX_KMERS_PER_CHUNK;
            let mut start = 0;
            loop {
                let end = (start + chunk_nucl).min(seql);
                let mut chunk = self.inner.sub(start, end);
                chunk.canonicalize();
                write_varint(w, self.count as u64)?;
                chunk.write_to_binary(w)?;
                if end == seql { break; }
                start += stride;
            }
            Ok(())
        }
    }

    /// 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 count = read_varint(r)? as u32;
        let inner = PackedSeq::read_from_binary(r)?;
        debug_assert!(inner.is_canonical(), "SuperKmer from disk is not canonical");
        Ok(Self { count, inner })
    }

    /// Sequence length in nucleotides.
    #[inline]
    pub fn seql(&self) -> usize {
        self.inner.seql()
    }

    /// Occurrence count.
    #[inline]
    pub fn count(&self) -> u32 {
        self.count
    }

    /// Increment occurrence count by 1.
    #[inline]
    pub fn increment(&mut self) {
        self.count += 1;
    }

    /// Add `n` to the occurrence count.
    #[inline]
    pub fn add(&mut self, n: u32) {
        self.count += n;
    }

    /// Set the occurrence count to `n`.
    #[inline]
    pub fn set_count(&mut self, n: u32) {
        self.count = n;
    }

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

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

    /// Extract the k-mer at position `i` using `params::k()`.
    #[inline]
    pub fn kmer(&self, i: usize) -> Result<Kmer, KmerError> {
        self.inner.kmer(i)
    }

    /// Extract the canonical k-mer at position `i`.
    #[inline]
    pub fn canonical_kmer(&self, i: usize) -> Result<CanonicalKmer, KmerError> {
        self.inner.canonical_kmer(i)
    }

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

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

    /// Borrowing iterator over all k-mers of length `params::k()` in order.
    ///
    /// Use when `self` is borrowed (e.g. from a `Vec<SuperKmer>`) and may be
    /// needed after iteration or iterated more than once.
    #[inline]
    pub fn iter_kmers(&self) -> impl Iterator<Item = Kmer> + '_ {
        self.inner.iter_kmers()
    }

    /// Borrowing iterator over all canonical k-mers in order.
    ///
    /// Use when `self` is borrowed.  Cannot be composed with `flat_map` over
    /// owned values — use [`into_canonical_kmers`](Self::into_canonical_kmers)
    /// in that case.
    #[inline]
    pub fn iter_canonical_kmers(&self) -> impl Iterator<Item = CanonicalKmer> + '_ {
        self.inner.iter_canonical_kmers()
    }

    /// Consuming iterator over all canonical k-mers in order.
    ///
    /// Moves `self` into the iterator; zero allocation.  Required when
    /// `SuperKmer` values are produced one-at-a-time (e.g. read from a file)
    /// and fed to `flat_map`:
    ///
    /// ```ignore
    /// file_reader.iter().flat_map(|sk| sk.into_canonical_kmers())
    /// ```
    #[inline]
    pub fn into_canonical_kmers(self) -> impl Iterator<Item = CanonicalKmer> {
        self.inner.into_canonical_kmers()
    }
}

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