//use ahash::RandomState;
use hashbrown::HashMap;
use obikseq::k;
use obikseq::unitig::Unitig;
use obikseq::{CanonicalKmer, Kmer, Sequence};
use std::cell::Cell;
use std::fmt;
use xxhash_rust::xxh3::Xxh3Builder;

// ── Types ─────────────────────────────────────────────────────────────────────

type FastHashMap<K, V> = HashMap<K, V, Xxh3Builder>;

// ── Node ──────────────────────────────────────────────────────────────────────
//
//  bit layout (LSB first):
//    bit  0   : can_extend_right — exactly one right canonical neighbour exists
//    bit  1   : can_extend_left  — exactly one left  canonical neighbour exists
//    bit  2   : visited
//    bits 3–4 : right_nuc — index 0–3 (A/C/G/T) of that neighbour; valid iff bit 0 = 1
//    bits 5–6 : left_nuc  — index 0–3 (A/C/G/T) of that neighbour; valid iff bit 1 = 1
//    bit  7   : reserved (0)
//
//  "can_extend" = false covers both 0 neighbours and ≥2 neighbours; the only
//  information needed for traversal is "exactly one".

#[repr(transparent)]
#[derive(Debug, Clone, Copy, Default)]
pub struct Node(u8);

impl Node {
    /// Returns `true` if the node can be extended to the right.
    ///
    /// A single right neighbour exists.
    pub fn can_extend_right(self) -> bool {
        self.0 & 0b0000_0001 != 0
    }

    /// Returns `true` if the node can be extended to the left.
    ///
    /// A single left neighbour exists.
    pub fn can_extend_left(self) -> bool {
        self.0 & 0b0000_0010 != 0
    }

    /// Returns `true` if the node has been visited.
    pub fn is_visited(self) -> bool {
        self.0 & 0b0000_0100 != 0
    }

    /// Index of the unique right neighbour (0=A, 1=C, 2=G, 3=T).
    /// Only meaningful when `can_extend_right()` is true.
    pub fn right_nuc(self) -> u8 {
        (self.0 >> 3) & 0b11
    }

    /// Index of the unique left neighbour (0=A, 1=C, 2=G, 3=T).
    /// Only meaningful when `can_extend_left()` is true.
    pub fn left_nuc(self) -> u8 {
        (self.0 >> 5) & 0b11
    }

    /// Marks the node as visited.
    pub fn set_visited(&mut self) {
        if self.is_visited() {
            unreachable!("from: is_visited -> The node has already been visited")
        }
        self.0 |= 0b0000_0100;
    }

    /// Number of left neighbours.
    pub fn n_left_neighbours(self) -> u8 {
        if self.can_extend_left() {
            1
        } else {
            let v = (self.0 >> 5) & 0b11;
            v + (v != 0) as u8
        }
    }

    /// Number of right neighbours.
    pub fn n_right_neighbours(self) -> u8 {
        if self.can_extend_right() {
            1
        } else {
            let v = (self.0 >> 3) & 0b11;
            v + (v != 0) as u8
        }
    }

    /// `nuc` = Some(i) → exactly one neighbour (bit 0 set, bits 3–4 = nucleotide index).
    /// `nuc` = None    → 0 or ≥2 neighbours; `count` encoded in bits 3–4 as count.sat_sub(1).
    pub fn set_right(&mut self, count: u8, nuc: Option<u8>) {
        self.0 &= !(0b0000_0001 | 0b001_1000);
        if count == 1 {
            self.0 |= 0b0000_0001;
            if let Some(n) = nuc {
                self.0 |= (n & 0b11) << 3;
                return;
            }
            unreachable!("nuc must be Some when count is 1");
        }
        self.0 |= (count.saturating_sub(1).min(3)) << 3;
    }

    /// `nuc` = Some(i) → exactly one neighbour (bit 0 set, bits 3–4 = nucleotide index).
    /// `nuc` = None    → 0 or ≥2 neighbours; `count` encoded in bits 3–4 as count.sat_sub(1).
    pub fn set_left(&mut self, count: u8, nuc: Option<u8>) {
        self.0 &= !(0b0000_0010 | 0b0110_0000);
        if count == 1 {
            self.0 |= 0b0000_0010;
            if let Some(n) = nuc {
                self.0 |= (n & 0b11) << 5;
                return;
            }
            unreachable!("nuc must be Some when count is 1");
        }
        self.0 |= (count.saturating_sub(1).min(3)) << 5;
    }
}

impl fmt::Display for Node {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        const NUC: [char; 4] = ['A', 'C', 'G', 'T'];
        let r = if self.can_extend_right() {
            format!("→{}", NUC[self.right_nuc() as usize])
        } else {
            format!("→{}", self.n_right_neighbours())
        };
        let l = if self.can_extend_left() {
            format!("←{}", NUC[self.left_nuc() as usize])
        } else {
            format!("←{}", self.n_left_neighbours())
        };
        let v = if self.is_visited() { "V" } else { "." };
        write!(f, "Node({r} {l} {v})")
    }
}

// ── GraphDeBruijn ─────────────────────────────────────────────────────────────

pub struct GraphDeBruijn {
    nodes: FastHashMap<CanonicalKmer, Cell<Node>>,
}

impl GraphDeBruijn {
    pub fn new() -> Self {
        Self {
            nodes: FastHashMap::with_hasher(Xxh3Builder::new()),
        }
    }

    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            nodes: FastHashMap::with_capacity_and_hasher(capacity, Xxh3Builder::new()),
        }
    }

    /// Insert a canonical kmer into the graph. No-op if already present.
    pub fn push(&mut self, kmer: CanonicalKmer) {
        self.nodes
            .entry(kmer)
            .or_insert_with(|| Cell::new(Node::default()));
    }

    /// For every node, find its unique right/left canonical neighbour (if any)
    /// and store the nucleotide index in the Node flags.
    ///
    /// Single pass thanks to Cell interior mutability.
    pub fn compute_degrees(&self) {
        for (&kmer, cell) in &self.nodes {
            let (rc, rn) = count_neighbors(kmer.right_canonical_neighbors(), &self.nodes);
            let (lc, ln) = count_neighbors(kmer.left_canonical_neighbors(), &self.nodes);

            let mut node = cell.get();
            node.set_right(rc, rn);
            node.set_left(lc, ln);
            cell.set(node);
        }
    }

    /// Iterates over the right neighbors of `kmer`.
    pub fn iter_right_neighbors(
        &self,
        kmer: CanonicalKmer,
    ) -> impl Iterator<Item = CanonicalKmer> + '_ {
        kmer.right_canonical_neighbors()
            .into_iter()
            .filter_map(|kmer| {
                self.nodes.get(&kmer)?;
                Some(kmer)
            })
    }

    /// Iterates over the left neighbors of `kmer`.
    pub fn iter_left_neighbors(
        &self,
        kmer: CanonicalKmer,
    ) -> impl Iterator<Item = CanonicalKmer> + '_ {
        kmer.left_canonical_neighbors()
            .into_iter()
            .filter_map(|kmer| {
                self.nodes.get(&kmer)?;
                Some(kmer)
            })
    }

    pub fn is_visited(&self, kmer: &CanonicalKmer) -> Option<bool> {
        self.nodes.get(kmer).map(|cell| cell.get().is_visited())
    }

    pub fn set_visited(&self, kmer: CanonicalKmer) {
        if let Some(cell) = self.nodes.get(&kmer) {
            let mut node = cell.get();
            node.set_visited();
            cell.set(node);
        }
    }

    /// Returns the single right neighbor of `kmer`, if it exists.
    pub fn the_single_right_neighbor(&self, kmer: CanonicalKmer) -> Option<CanonicalKmer> {
        let node = self.nodes.get(&kmer)?.get();
        if !node.can_extend_right() {
            return None;
        }
        let next = kmer.into_kmer().push_right(node.right_nuc()).canonical();
        self.nodes.contains_key(&next).then_some(next)
    }

    /// Returns the single left neighbor of `kmer`, if it exists.
    pub fn the_single_left_neighbor(&self, kmer: CanonicalKmer) -> Option<CanonicalKmer> {
        let node = self.nodes.get(&kmer)?.get();
        if !node.can_extend_left() {
            return None;
        }
        let next = kmer.into_kmer().push_left(node.left_nuc()).canonical();
        self.nodes.contains_key(&next).then_some(next)
    }

    /// Internal iterator over unitig-start nodes; drives `iter_unitig`.
    ///
    /// MUST NOT be consumed standalone: the second pass finds cycle nodes only
    /// because `iter_unitig` lazily interleaves chain traversal between the two passes.
    ///
    /// Two passes:
    ///   1. Chain ends / isolated nodes (at most one extension missing):
    ///        - `!can_extend_left`  → yield canonical form
    ///        - `!can_extend_right` → yield reverse complement
    ///   2. Nodes still unvisited → part of a cycle; yield canonical form.
    fn start_iter(&self) -> impl Iterator<Item = (CanonicalKmer, Option<Kmer>)> + '_ {
        StartIter::new(self)
    }

    fn next_unitig_kmer(&self, kmer: Kmer) -> Option<Kmer> {
        let canonical = kmer.canonical();
        let node = self.nodes.get(&canonical)?.get();

        let direct = kmer.raw() == canonical.raw();

        if (direct && !node.can_extend_right()) || (!direct && !node.can_extend_left()) {
            return None;
        }

        let next_c: CanonicalKmer = if direct {
            canonical
                .into_kmer()
                .push_right(node.right_nuc())
                .canonical()
        } else {
            canonical.into_kmer().push_left(node.left_nuc()).canonical()
        };

        let cell = self.nodes.get(&next_c)?;
        let next_node = cell.get();
        if next_node.is_visited() {
            return None;
        }

        let oriented = oriented_next(kmer, next_c);
        let ndirect = oriented.raw() == next_c.raw();

        if (ndirect && next_node.n_right_neighbours() > 1)
            || (!ndirect && next_node.n_left_neighbours() > 1)
        {
            return None;
        }

        let mut updated = next_node;
        updated.set_visited();
        cell.set(updated);
        Some(oriented)
    }

    fn iter_unitig_kmers(&self, start: Kmer) -> UnitigIter<'_> {
        UnitigIter {
            graph: self,
            current: Some(start),
        }
    }

    pub fn iter_unitig(&self) -> impl Iterator<Item = Unitig> + '_ {
        let k = k();
        self.start_iter().map(move |(start, first_next)| {
            let mut nucs: Vec<u8> = (0..k).map(|i| start.nucleotide(i)).collect();
            if let Some(next_c) = first_next {
                for kmer in self.iter_unitig_kmers(next_c) {
                    nucs.push(kmer.nucleotide(k - 1));
                }
            }
            Unitig::from_nucleotides(&nucs)
        })
    }

    pub fn len(&self) -> usize {
        self.nodes.len()
    }

    pub fn is_empty(&self) -> bool {
        self.nodes.is_empty()
    }
}

// --- StartIter -----------------------------------------------------------------
struct StartIter<'a> {
    graph: &'a GraphDeBruijn,
    nodes: hashbrown::hash_map::Iter<'a, CanonicalKmer, Cell<Node>>,
    suspended: Vec<CanonicalKmer>,
    in_cycle_pass: bool,
}

impl<'a> StartIter<'a> {
    fn new(graph: &'a GraphDeBruijn) -> Self {
        Self {
            graph,
            nodes: graph.nodes.iter(),
            suspended: Vec::new(),
            in_cycle_pass: false,
        }
    }
}

impl<'a> Iterator for StartIter<'a> {
    type Item = (CanonicalKmer, Option<Kmer>);

    fn next(&mut self) -> Option<(CanonicalKmer, Option<Kmer>)> {
        loop {
            let current = if let Some(k) = self.suspended.pop() {
                k
            } else {
                match self.nodes.next() {
                    Some((&k, _)) => k,
                    None => {
                        if self.in_cycle_pass {
                            return None;
                        }
                        self.in_cycle_pass = true;
                        self.nodes = self.graph.nodes.iter();
                        match self.nodes.next() {
                            Some((&k, _)) => k,
                            None => return None,
                        }
                    }
                }
            };

            let node = match self.graph.nodes.get(&current) {
                Some(c) => c.get(),
                None => continue,
            };
            if node.is_visited() {
                continue;
            }
            if !self.in_cycle_pass && node.can_extend_left() {
                continue;
            }

            self.graph.set_visited(current);

            if let Some(next) = self.graph.the_single_right_neighbor(current) {
                if self.graph.is_visited(&next).unwrap_or(true) {
                    return Some((current, None));
                }
                self.graph.set_visited(next);
                let oriented = oriented_next(current.into_kmer(), next);
                return Some((current, Some(oriented)));
            }

            let mut first_neighbor: Option<CanonicalKmer> = None;
            for neighbor in self.graph.iter_right_neighbors(current) {
                if self.graph.is_visited(&neighbor).unwrap_or(true) {
                    continue;
                }
                if first_neighbor.is_none() {
                    self.graph.set_visited(neighbor);
                    first_neighbor = Some(neighbor);
                } else {
                    self.suspended.push(neighbor);
                }
            }

            let oriented = match first_neighbor {
                Some(neighbor) => Some(oriented_next(current.into_kmer(), neighbor)),
                None => None,
            };
            return Some((current, oriented));
        }
    }
}

// ── UnitigIter ────────────────────────────────────────────────────────────────

struct UnitigIter<'a> {
    graph: &'a GraphDeBruijn,
    current: Option<Kmer>,
}

impl Iterator for UnitigIter<'_> {
    type Item = Kmer;

    fn next(&mut self) -> Option<Kmer> {
        let current = self.current?;
        self.current = self.graph.next_unitig_kmer(current);
        Some(current)
    }
}

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

fn oriented_next(from: Kmer, to: CanonicalKmer) -> Kmer {
    if from.is_overlapping(to.into_kmer()) {
        to.into_kmer()
    } else {
        to.revcomp()
    }
}

/// Returns `Some(i)` if exactly one of the four canonical neighbours exists in
/// the graph, where `i` is its index (0=A, 1=C, 2=G, 3=T). Returns `None` for
/// zero or ≥2 existing neighbours.
fn count_neighbors(
    neighbors: [CanonicalKmer; 4],
    nodes: &FastHashMap<CanonicalKmer, Cell<Node>>,
) -> (u8, Option<u8>) {
    let mut count = 0u8;
    let mut first = None;
    for (i, neighbour) in neighbors.iter().enumerate() {
        if nodes.contains_key(neighbour) {
            count += 1;
            if first.is_none() {
                first = Some(i as u8);
            }
        }
    }
    if count == 1 {
        (1, first)
    } else {
        (count, None)
    }
}

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