use obikentropy::EntropyTracker;
use obikseq::kmer::{Minimizer, hash_kmer};
use obikseq::params;

use crate::encoding::encode_nuc;

// ── Stack-allocated ring buffer ───────────────────────────────────────────────

/// Fixed-capacity ring buffer backed by a stack array.
/// N must be a power of two; operations are branchless via `% N`.
struct Ring<T: Copy + Default, const N: usize> {
    buf: [T; N],
    head: usize,
    len: usize,
}

impl<T: Copy + Default, const N: usize> Ring<T, N> {
    #[inline]
    fn new() -> Self {
        Self {
            buf: [T::default(); N],
            head: 0,
            len: 0,
        }
    }
    #[inline]
    fn clear(&mut self) {
        self.len = 0;
        self.head = 0;
    }

    #[inline]
    fn push_back(&mut self, val: T) {
        self.buf[(self.head + self.len) % N] = val;
        self.len += 1;
    }
    #[inline]
    fn pop_front(&mut self) -> T {
        let val = self.buf[self.head];
        self.head = (self.head + 1) % N;
        self.len -= 1;
        val
    }
    #[inline]
    fn pop_back(&mut self) -> T {
        self.len -= 1;
        self.buf[(self.head + self.len) % N]
    }
    #[inline]
    fn front(&self) -> Option<&T> {
        if self.len > 0 {
            Some(&self.buf[self.head])
        } else {
            None
        }
    }
    #[inline]
    fn back(&self) -> Option<&T> {
        if self.len > 0 {
            Some(&self.buf[(self.head + self.len - 1) % N])
        } else {
            None
        }
    }

}

// ── MmerItem ──────────────────────────────────────────────────────────────────

#[derive(Clone, Copy, Default)]
struct MmerItem {
    /// 0-based position of this m-mer's first base within the current segment.
    position: usize,
    /// Raw canonical m-mer value (right-aligned).
    canonical: u64,
    /// mix64 hash of the canonical m-mer, used as the random ordering key.
    hash: u64,
}

// ── RollingStat ───────────────────────────────────────────────────────────────

pub struct RollingStat {
    entropy_max_k: usize,
    k: usize,
    m: usize,
    rolling_k: u64,
    rolling_rck: u64,
    k_mask: u64,
    m_mask: u64,
    received: usize,

    // Minimizer selection state.
    minimier: Ring<MmerItem, 32>,

    // Entropy tracking, composed as a plain inline field so both concerns
    // update in the same streaming pass without being conflated in one
    // struct — see `obikentropy::EntropyTracker`.
    entropy: EntropyTracker,
}

impl RollingStat {
    pub fn new(entropy_max_k: usize) -> Self {
        let k = params::k();
        let m = params::m();
        Self {
            entropy_max_k,
            k,
            m,
            rolling_k: 0,
            rolling_rck: 0,
            k_mask: (!0u64) >> (64 - k * 2),
            m_mask: (!0u64) >> (64 - m * 2),
            received: 0,
            minimier: Ring::new(),
            entropy: EntropyTracker::new(k),
        }
    }

    pub fn reset(&mut self) {
        self.rolling_k = 0;
        self.rolling_rck = 0;
        self.received = 0;

        self.minimier.clear();
        self.entropy.reset();
    }

    pub fn push(&mut self, nuc: u8) {
        let k = self.k;
        let m = self.m;

        let bnuc = encode_nuc(nuc);
        let cnuc = bnuc ^ 3;

        self.rolling_k = ((self.rolling_k << 2) | (bnuc as u64)) & self.k_mask;
        self.rolling_rck =
            ((self.rolling_rck >> 2) | ((cnuc as u64) << ((k - 1) * 2))) & self.k_mask;

        self.received += 1;

        if self.received >= m {
            let possible_canonical_m =
                (self.rolling_k & self.m_mask).min(self.rolling_rck >> ((k - m) * 2));
            let possible_hash_m = hash_kmer(possible_canonical_m << 64 - m * 2);
            let possible_pos_m = self.received - m;

            while self
                .minimier
                .back()
                .map_or(false, |it| it.hash >= possible_hash_m)
            {
                self.minimier.pop_back();
            }
            self.minimier.push_back(MmerItem {
                position: possible_pos_m,
                canonical: possible_canonical_m,
                hash: possible_hash_m,
            });

            if self.received > k {
                while self
                    .minimier
                    .front()
                    .map_or(false, |it| it.position + k < self.received)
                {
                    self.minimier.pop_front();
                }
            }
        }

        self.entropy.push(self.received, self.rolling_k);
    }

    pub fn ready(&self) -> bool {
        self.received >= self.k
    }

    pub fn minimizer_position(&self) -> Option<usize> {
        if self.ready() {
            self.minimier.front().map(|it| it.position)
        } else {
            None
        }
    }

    pub fn canonical_minimizer_raw(&self) -> Option<u64> {
        if self.ready() {
            self.minimier.front().map(|it| it.canonical)
        } else {
            None
        }
    }

    pub fn canonical_minimizer(&self) -> Option<Minimizer> {
        self.canonical_minimizer_raw()
            .map(|raw| Minimizer::from_raw_unchecked(raw << (64 - self.m * 2)))
    }

    pub fn normalized_entropy(&self) -> Option<f64> {
        if !self.ready() {
            return None;
        }
        Some(self.entropy.normalized_entropy(self.entropy_max_k))
    }
}
