use core::marker::PhantomData;
use core::mem;
use core::ops::Not;
use core::ptr::NonNull;

#[cfg(all(test, feature = "std"))]
extern crate std;

#[cfg(test)]
pub trait TestDebug: core::fmt::Debug {}
#[cfg(test)]
impl<T: core::fmt::Debug> TestDebug for T {}

#[cfg(not(test))]
pub trait TestDebug {}
#[cfg(not(test))]
impl<T> TestDebug for T {}

pub trait TestDebugInspect: Sized {
    fn test_debug_inspect_with<F: FnOnce(&Self)>(self, f: F) -> Self
    where
        Self: TestDebug,
    {
        f(&self);
        self
    }
}

impl<T: Sized> TestDebugInspect for T {}

#[allow(dead_code)]
pub trait DebugInspect: Sized {
    // fn debug_inspect(self) -> Self {
    //     dbg!(&self);
    //     self
    // }

    fn debug_inspect_with<F: FnOnce(&Self)>(self, f: F) -> Self
    where
        Self: core::fmt::Debug,
    {
        f(&self);
        self
    }
}

impl<T: Sized> DebugInspect for T {}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SearchResult<T> {
    FoundAt(T),
    NotFoundAt(T),
    Empty,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Color {
    Red,
    Black,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Side {
    Left,
    Right,
}

impl Not for Side {
    type Output = Self;

    fn not(self) -> Self::Output {
        match self {
            Side::Left => Side::Right,
            Side::Right => Side::Left,
        }
    }
}

pub unsafe trait UnsafeNode: TestDebug {
    type Key: Eq + Ord + TestDebug;
    fn left(&self) -> Option<NonNull<Self>>;
    fn right(&self) -> Option<NonNull<Self>>;
    fn parent(&self) -> Option<NonNull<Self>>;
    fn key(&self) -> &Self::Key;
    fn color(&self) -> Color;

    fn set_left(&self, left: Option<NonNull<Self>>);
    fn set_right(&self, right: Option<NonNull<Self>>);
    fn set_parent(&self, parent: Option<NonNull<Self>>);
    fn set_color(&self, color: Color);

    fn copy_meta_from(&self, other: &Self) {
        self.set_color(other.color());
        self.set_parent(other.parent());
        self.set_left(other.left());
        self.set_right(other.right());
    }
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Copy)]
pub enum LeftOrRight<T> {
    Left(T),
    Right(T),
}

impl<T> LeftOrRight<Option<T>> {
    pub fn transpose(self) -> Option<LeftOrRight<T>> {
        match self {
            LeftOrRight::Left(Some(t)) => Some(LeftOrRight::Left(t)),
            LeftOrRight::Right(Some(t)) => Some(LeftOrRight::Right(t)),
            _ => None,
        }
    }
}

impl<T> LeftOrRight<T> {
    pub fn map<U, F: FnOnce(T) -> U>(self, f: F) -> LeftOrRight<U> {
        match self {
            LeftOrRight::Left(t) => LeftOrRight::Left(f(t)),
            LeftOrRight::Right(t) => LeftOrRight::Right(f(t)),
        }
    }

    pub fn as_ref(&self) -> LeftOrRight<&T> {
        match self {
            LeftOrRight::Left(t) => LeftOrRight::Left(t),
            LeftOrRight::Right(t) => LeftOrRight::Right(t),
        }
    }

    pub fn as_inner(&self) -> &T {
        match self {
            LeftOrRight::Left(t) => t,
            LeftOrRight::Right(t) => t,
        }
    }

    pub fn as_inner_mut(&mut self) -> &mut T {
        match self {
            LeftOrRight::Left(t) => t,
            LeftOrRight::Right(t) => t,
        }
    }

    pub fn into_inner(self) -> T {
        match self {
            LeftOrRight::Left(t) => t,
            LeftOrRight::Right(t) => t,
        }
    }
}

#[derive(Debug)]
pub enum Handle<N> {
    EmptyRoot,
    Root(NonNull<N>),
    Child {
        parent: NonNull<N>,
        node: LeftOrRight<Option<NonNull<N>>>,
    },
}

impl<N: UnsafeNode> Eq for Handle<N> {}
impl<N: UnsafeNode> PartialEq for Handle<N> {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Self::Root(l0), Self::Root(r0)) => l0 == r0,
            (
                Self::Child {
                    parent: l_parent,
                    node: l_node,
                },
                Self::Child {
                    parent: r_parent,
                    node: r_node,
                },
            ) => l_parent == r_parent && l_node == r_node,
            _ => core::mem::discriminant(self) == core::mem::discriminant(other),
        }
    }
}

impl<N> Clone for Handle<N> {
    fn clone(&self) -> Self {
        match self {
            Handle::EmptyRoot => Handle::EmptyRoot,
            Handle::Root(node) => Handle::Root(*node),
            Handle::Child { parent, node } => Handle::Child {
                parent: *parent,
                node: *node,
            },
        }
    }
}

impl<N> !Sync for Handle<N> {}

impl<N: UnsafeNode> Handle<N> {
    fn from_node(node: Option<NonNull<N>>) -> Self {
        let Some(node) = node else {
            return Handle::EmptyRoot;
        };

        let parent = unsafe { node.as_ref().parent() };

        match parent {
            Some(parent) => {
                let side = if unsafe { parent.as_ref().left() } == Some(node) {
                    LeftOrRight::Left(Some(node))
                } else {
                    LeftOrRight::Right(Some(node))
                };

                Handle::Child { parent, node: side }
            }
            None => Handle::Root(node),
        }
    }

    #[expect(dead_code)]
    fn refresh_from_node(&mut self) {
        *self = Self::from_node(self.node());
    }

    #[expect(dead_code)]
    fn is_root(&self) -> bool {
        matches!(self, Handle::Root(_) | Handle::EmptyRoot)
    }

    fn is_red(&self) -> bool {
        self.color() == Color::Red
    }

    fn is_black(&self) -> bool {
        self.color() == Color::Black
    }

    fn set_color(&mut self, color: Color) {
        if let Some(node) = self.node() {
            unsafe { node.as_ref().set_color(color) };
        }
    }

    fn side(&self) -> Option<Side> {
        match self {
            Handle::EmptyRoot => None,
            Handle::Root(_) => None,
            Handle::Child { node, .. } => match node {
                LeftOrRight::Left(_) => Some(Side::Left),
                LeftOrRight::Right(_) => Some(Side::Right),
            },
        }
    }

    fn child(&self, side: Side) -> Option<Self> {
        match side {
            Side::Left => self.left_child(),
            Side::Right => self.right_child(),
        }
    }

    fn left_child(&self) -> Option<Self> {
        let &parent = match self {
            Handle::EmptyRoot => return None,
            Handle::Root(parent) => parent,
            Handle::Child { node, .. } => node.as_inner().as_ref()?,
        };

        let child = unsafe { parent.as_ref().left() };

        Some(Handle::Child {
            parent,
            node: LeftOrRight::Left(child),
        })
    }

    fn right_child(&self) -> Option<Self> {
        let &parent = match self {
            Handle::EmptyRoot => return None,
            Handle::Root(parent) => parent,
            Handle::Child { node, .. } => node.as_inner().as_ref()?,
        };

        let child = unsafe { parent.as_ref().right() };

        Some(Handle::Child {
            parent,
            node: LeftOrRight::Right(child),
        })
    }

    fn right_child_extant(&self) -> Option<Self> {
        self.right_child().filter(|child| !child.is_nil())
    }

    fn left_child_extant(&self) -> Option<Self> {
        self.left_child().filter(|child| !child.is_nil())
    }

    fn into_has_node(self) -> Option<Self> {
        self.node().map(|_| self)
    }

    fn children(&self) -> (Option<Self>, Option<Self>) {
        (self.left_child(), self.right_child())
    }

    fn non_nil_children(&self) -> (Option<Self>, Option<Self>) {
        (
            self.left_child().and_then(Self::into_has_node),
            self.right_child().and_then(Self::into_has_node),
        )
    }

    fn sibling(&self) -> Option<Self> {
        let parent = self.parent()?;

        match self {
            Self::Child { node, .. } => match node {
                LeftOrRight::Left(_) => parent.right_child(),
                LeftOrRight::Right(_) => parent.left_child(),
            },
            _ => unreachable!(),
        }
    }

    fn set_left_child(&mut self, child: &mut Self) {
        if let Some(parent) = self.node() {
            let mut child_node = child.node();

            unsafe {
                parent.as_ref().set_left(child_node);
                if let Some(child) = child_node.as_mut().map(|c| c.as_mut()) {
                    child.set_parent(Some(parent));
                }
            }

            *child = Handle::Child {
                parent,
                node: LeftOrRight::Left(child_node),
            };
        }
    }

    fn set_right_child(&mut self, child: &mut Self) {
        if let Some(parent) = self.node() {
            let mut child_node = child.node();

            unsafe {
                parent.as_ref().set_right(child_node);
                if let Some(child) = child_node.as_mut().map(|c| c.as_mut()) {
                    child.set_parent(Some(parent));
                }
            }

            *child = Handle::Child {
                parent,
                node: LeftOrRight::Right(child_node),
            };
        }
    }

    fn node(&self) -> Option<NonNull<N>> {
        match self {
            Handle::EmptyRoot => None,
            Handle::Root(parent) => Some(*parent),
            Handle::Child { node, .. } => node.as_inner().as_ref().copied(),
        }
    }

    fn parent_and_side(&self) -> Option<LeftOrRight<Self>> {
        let (&parent, side) = match self {
            Handle::EmptyRoot => return None,
            Handle::Root(_) => return None,
            Handle::Child { parent, node } => (parent, node.map(|_| ())),
        };

        let grandparent = unsafe { parent.as_ref().parent() };

        let parent = if let Some(grandparent) = grandparent {
            let node = if unsafe { grandparent.as_ref().left() } == Some(parent) {
                LeftOrRight::Left(Some(parent))
            } else {
                LeftOrRight::Right(Some(parent))
            };

            Handle::Child {
                parent: grandparent,
                node,
            }
        } else {
            Handle::Root(parent)
        };

        Some(side.map(|_| parent))
    }

    fn parent(&self) -> Option<Self> {
        self.parent_and_side().map(LeftOrRight::into_inner)
    }

    fn set_parent_handle(&mut self, parent: Option<LeftOrRight<Handle<N>>>) {
        match parent {
            Some(LeftOrRight::Left(mut parent)) => {
                parent.set_left_child(self);
            }
            Some(LeftOrRight::Right(mut parent)) => {
                parent.set_right_child(self);
            }
            None => match self.node() {
                Some(node) => {
                    unsafe {
                        node.as_ref().set_parent(None);
                    }
                    *self = Handle::Root(node);
                }
                None => {
                    *self = Handle::EmptyRoot;
                }
            },
        }
    }

    fn set_parent(&mut self, parent: Option<LeftOrRight<NonNull<N>>>) {
        let parent = parent.map(|p| p.map(|p| Handle::from_node(Some(p))));
        self.set_parent_handle(parent);
    }

    fn color(&self) -> Color {
        match self {
            Handle::Root(_) | Handle::EmptyRoot => Color::Black,
            Handle::Child { node, .. } => match node.as_inner() {
                Some(node) => unsafe { node.as_ref().color() },
                None => Color::Black,
            },
        }
    }

    /// rotate `self` left, returning the new root of the subtree, or
    /// `Err(self)` if `self` has no left child
    fn rotate_left<F: FnOnce(Self)>(mut self, on_root: F) -> Result<Self, Self> {
        let Some(mut y) = self.right_child() else {
            return Err(self);
        };

        let Some(mut b) = y.left_child() else {
            return Err(self);
        };

        let parent = self
            .parent_and_side()
            .and_then(|p| p.map(|p| p.node()).transpose());

        y.set_parent(parent);

        self.set_right_child(&mut b);

        y.set_left_child(&mut self);

        if matches!(y, Handle::Root(_)) {
            on_root(y.clone());
        }

        Ok(y)
    }

    /// rotate `self` right, returning the new root of the subtree, or
    /// `Err(self)` if `self` has no left child
    fn rotate_right<F: FnOnce(Self)>(mut self, on_root: F) -> Result<Self, Self> {
        let Some(mut y) = self.left_child() else {
            return Err(self);
        };

        let Some(mut b) = y.right_child() else {
            return Err(self);
        };

        let parent = self
            .parent_and_side()
            .and_then(|p| p.map(|p| p.node()).transpose());

        y.set_parent(parent);

        self.set_left_child(&mut b);

        y.set_right_child(&mut self);

        if matches!(y, Handle::Root(_)) {
            on_root(y.clone());
        }

        Ok(y)
    }

    /// rotate `self` into the position of its parent.
    /// Returns `Ok(self)` in the new position, or `Err(self)` if `self` is the
    /// root of the tree.
    fn rotate_up<F: FnOnce(Self)>(self, on_root: F) -> Result<Self, Self> {
        let Some(parent) = self.parent() else {
            on_root(self.clone());
            return Err(self);
        };

        match self.side() {
            Some(Side::Left) => parent.rotate_right(on_root).map_err(|_| self),
            Some(Side::Right) => parent.rotate_left(on_root).map_err(|_| self),
            None => unreachable!(),
        }
    }

    /// rotate `self` into the direction of `side`, returning the new root of the subtree, or
    /// `Err(self)` if `self` has no left child
    pub fn rotate<F: FnOnce(Self)>(self, side: Side, on_root: F) -> Result<Self, Self> {
        match side {
            Side::Left => self.rotate_left(on_root),
            Side::Right => self.rotate_right(on_root),
        }
    }

    /// Returns the least non-nil node in the subtree rooted at `self`, or
    /// `None` if the subtree is empty.
    pub fn minimum_of(&self) -> Option<Self> {
        let mut current = self.clone();
        while let Some(left) = current.left_child_extant() {
            current = left;
        }
        Some(current)
    }

    /// Returns the greatest non-nil node in the subtree rooted at `self`, or
    /// `None` if the subtree is empty.
    pub fn maximum_of(&self) -> Option<Self> {
        let mut current = self.clone();
        while let Some(right) = current.right_child_extant() {
            current = right;
        }
        Some(current)
    }

    /// Returns the next greater non-nil node in the tree, or `None` if `self`
    /// is the greatest node.
    pub fn next_of(&self) -> Option<Self> {
        match self.right_child_extant() {
            Some(right) => right.minimum_of(),
            _ => {
                let mut current = self.clone();
                while let Some(parent) = current.parent() {
                    if current.side() == Some(Side::Left) {
                        return Some(parent);
                    }

                    current = parent;
                }
                None
            }
        }
    }

    /// Returns the next smaller non-nil node in the tree, or `None` if `self`
    /// is the smallest node.
    pub fn next_back_of(&self) -> Option<Self> {
        match self.left_child_extant() {
            Some(left) => left.maximum_of(),
            None => {
                let mut current = self.clone();
                while let Some(parent) = current.parent() {
                    if parent.right_child().as_ref() == Some(&current) {
                        return Some(parent);
                    }
                    current = parent;
                }
                None
            }
        }
    }

    pub fn is_nil(&self) -> bool {
        match self {
            Handle::EmptyRoot => true,
            Handle::Root(_) => false,
            Handle::Child { node, .. } => node.as_inner().is_none(),
        }
    }

    pub fn make_nil(&mut self) -> Option<NonNull<N>> {
        match self {
            Handle::EmptyRoot => None,
            &mut Handle::Root(node) => {
                *self = Handle::EmptyRoot;
                Some(node)
            }
            &mut Handle::Child { mut node, .. } => unsafe {
                let (left, right) = self.non_nil_children();

                if let Some(mut left) = left {
                    left.set_parent(None);
                }

                if let Some(mut right) = right {
                    right.set_parent(None);
                }

                if let Some(parent) = self
                    .parent_and_side()
                    .and_then(|p| p.map(|p| p.node()).transpose())
                {
                    match parent {
                        LeftOrRight::Left(parent) => parent.as_ref().set_left(None),
                        LeftOrRight::Right(parent) => parent.as_ref().set_right(None),
                    }
                }

                node.as_inner_mut().take()
            },
        }
    }

    #[must_use = "inserting a node may replace an existing node, which must be deallocated"]
    pub fn insert(&mut self, new_node: NonNull<N>) -> Option<NonNull<N>> {
        match self {
            Self::EmptyRoot => {
                *self = Handle::Root(new_node);
                None
            }
            Self::Root(old) => {
                let old = *old;
                *self = Handle::Root(new_node);
                Some(old)
            }
            Self::Child { node, parent } => {
                let old = match node {
                    LeftOrRight::Left(old) => unsafe {
                        parent.as_ref().set_left(Some(new_node));
                        old.replace(new_node)
                    },
                    LeftOrRight::Right(old) => unsafe {
                        parent.as_ref().set_right(Some(new_node));
                        old.replace(new_node)
                    },
                };

                unsafe {
                    if let Some(old) = old {
                        new_node.as_ref().copy_meta_from(old.as_ref());
                    } else {
                        new_node.as_ref().set_parent(Some(*parent));
                    }
                }

                old
            }
        }
    }
}

pub struct RBTree<N: UnsafeNode> {
    root: Option<NonNull<N>>,
}

impl<N: UnsafeNode> RBTree<N> {
    pub fn new() -> Self {
        Self { root: None }
    }

    fn root_handle(&self) -> Handle<N> {
        match self.root {
            Some(root) => Handle::Root(root),
            None => Handle::EmptyRoot,
        }
    }

    fn set_root_handle(&mut self, mut handle: Handle<N>) {
        handle.set_parent(None);
        handle.set_color(Color::Black);
        self.root = handle.node();
    }

    pub fn find_by_key<Q>(&self, key: &Q) -> SearchResult<Handle<N>>
    where
        N::Key: core::borrow::Borrow<Q>,
        Q: Ord + ?Sized,
    {
        use core::borrow::Borrow;
        use core::cmp::Ordering::*;

        let mut current = self.root_handle();

        loop {
            let node = match &current {
                Handle::Root(node) => *node,
                Handle::Child { node, .. } => match node.into_inner() {
                    Some(node) => node,
                    None => return SearchResult::NotFoundAt(current),
                },
                _ => return SearchResult::Empty,
            };

            match unsafe { node.as_ref().key().borrow().cmp(key) } {
                Less => {
                    current = current.right_child().expect("current is an occupied node");
                }
                Greater => {
                    current = current.left_child().expect("current is an occupied node");
                }
                Equal => {
                    return SearchResult::FoundAt(current);
                }
            }
        }
    }

    #[must_use]
    pub fn insert_node(&mut self, new_node: NonNull<N>) -> Option<NonNull<N>> {
        let node_ref = unsafe { new_node.as_ref() };

        let mut entry = match self
            .find_by_key(node_ref.key())
            .test_debug_inspect_with(|_res| {
                #[cfg(all(test, feature = "std"))]
                eprintln!("find_by_key({:?}) = {:?}", node_ref.key(), _res);
            }) {
            SearchResult::FoundAt(mut entry) => {
                return entry.insert(new_node);
            }
            SearchResult::Empty => {
                node_ref.set_color(Color::Black);
                node_ref.set_parent(None);
                self.set_root_handle(Handle::Root(new_node));

                return None;
            }
            SearchResult::NotFoundAt(mut entry) => {
                unsafe { new_node.as_ref().set_color(Color::Red) };
                _ = entry.insert(new_node);

                entry
            }
        };

        // Fixing

        // we've introduced a new red node, but one of the invariants of the red-black tree is that red nodes cannot have red children.
        // In the case that our new node's parent is red, we need to fix the tree:
        while let Some(mut parent) = entry.parent()
            && parent.is_red()
        {
            // since the parent is red, it must have a grandparent (and an
            // uncle), since the root of the tree is always black.
            let mut gp = parent
                .parent()
                .expect("parent is red, so it must have a grandparent");

            let mut uncle = parent
                .sibling()
                .expect("parent is red, so it must have a grandparent, and thus an uncle");

            if uncle.is_red() {
                // Case 1:
                // If the uncle is red, we recolour both the parent and uncle black
                // and the grandparent red, preserving the black-height of the tree.
                //
                // In the next loop, we look at the grandparent, which might have
                // had a red parent.

                parent.set_color(Color::Black);
                uncle.set_color(Color::Black);
                gp.set_color(Color::Red);

                entry = gp;
            } else {
                // Case 2:
                // If the uncle is black, we perform one (two) rotations.
                // our current subtree (starting from the grandparent) looks like this:
                //
                //        (1)           (2)
                //        GP(B) <or>    GP(B)
                //       /  \          /  \
                //     P(R)  U(B)    P(R)  U(B)
                //    /   \         /   \
                //  ...   E(R)    E(R)  ...

                if uncle.side() == entry.side() {
                    // If the uncle has the same sidedness as the entry (1),
                    // we rotate the entry up into the parent's position and end
                    // up with a tree with the same shape.

                    entry
                        .rotate_up(|_| panic!("entry cannot become root"))
                        .unwrap_or_else(|_| panic!("entry is not root, so it must have a parent"));
                }

                // We rotate the grandparent towards the uncle to end up with
                // the following shape:
                //
                //       P(B)
                //      /   \
                //    E(R)   GP(R)
                //          /  \
                //        ...   U(B)
                // the grandparent is painted red in order to preserve the black-height of the tree.
                gp.set_color(Color::Red);

                let mut parent = gp
                    .rotate(
                        uncle
                            .side()
                            .expect("uncle is not root, so it must have a side"),
                        |root| self.set_root_handle(root),
                    )
                    .unwrap_or_else(|_| panic!("grandparent has children"));

                // Whatever node ends up in the grandparent's position (either
                // `entry` or `parent`) is painted black, resolving the red-red
                // violation and replacing the grandparents black-level within
                // the tree.
                parent.set_color(Color::Black);

                // We are done fixing the tree, so we can break out of the loop.
                break;
            }
        }

        self.root_handle().set_color(Color::Black);

        None
    }

    #[must_use]
    pub fn remove<Q>(&mut self, key: &Q) -> Option<NonNull<N>>
    where
        N::Key: core::borrow::Borrow<Q>,
        Q: Ord + ?Sized,
    {
        let SearchResult::FoundAt(z) = self.find_by_key(key) else {
            return None;
        };

        // Y is either Z, the removed node in the case that Z has at most
        // one child, or Y is Z's successor which is guaranteed to have at most one
        // child (the right child).
        let y = match z.non_nil_children() {
            (Some(_), Some(_)) => z
                .next_of()
                .expect("z has a right child, so it must have a successor"),
            _ => z.clone(),
        };

        // In either case, X is the only child of Y, if it exists.
        let mut x = match y.children() {
            (Some(left), Some(right)) => {
                if !right.is_nil() {
                    right
                } else {
                    left
                }
            }
            _ => unreachable!("y is present, so it has children"),
        };

        let parent = y.parent_and_side();
        let color = y.color();

        // If Y is Z's successor, move Y's data into Z (or move Z's meta into Y).
        if y != z {
            unsafe {
                y.node()
                    .expect("z is occupied")
                    .as_ref()
                    .copy_meta_from(z.node().expect("z is occupied").as_ref());
            }
        }

        let Some(parent) = parent else {
            // If Y was the root, and the tree is empty.
            self.set_root_handle(x);
            return z.node();
        };

        // X is promoted to Y's position, and Y is unlinked from the tree.
        x.set_parent_handle(Some(parent));

        if !x.is_nil() {
            // If X is not nil, it must be red, since its parent, Y, must be
            // black (otherwise a red node would have a red child), and since Y
            // has only one child and nil leafs are implicitly black, Y would
            // have been imbalanced if X were black.
            // Since X is replacing a black node (Y), the black-height is
            // preserved by painting X black.
            x.set_color(Color::Black);

            return z.node();
        }

        // Fixing

        if color == Color::Black {
            // If Y was black, then we have a black-height violation, since X
            // was nil (and thus black) and replaced a black node (Y).

            // x is a NIL leaf and doubly black.
            let mut x = x;

            while x != self.root_handle() && x.is_black() {
                let mut parent = x.parent().expect("x is not root, so it must have a parent");

                // W is X's sibling, and must exist because the subtree at X has
                // a black-height of 2, and the two subtrees of the parent must
                // have the same black-height.
                let mut w = x.sibling().expect("w exists because x is black-deficient");

                // Case 1: W is red -> parent and W's children are black
                if w.is_red() {
                    w.set_color(Color::Black);
                    parent.set_color(Color::Red);

                    parent = w
                        .rotate_up(|root| self.set_root_handle(root))
                        .unwrap_or_else(|_| panic!("w is the child of parent"));

                    // X's sibling has changed
                    w = x.sibling().expect("w exists because x is black-deficient");
                }

                // Case 2: W is black
                assert!(
                    w.is_black(),
                    "w is black because it was red in the previous case"
                );

                // W's children exist because W exists
                match (
                    w.left_child().map(|c| c.color()).unwrap_or(Color::Black),
                    w.right_child().map(|c| c.color()).unwrap_or(Color::Black),
                ) {
                    (Color::Black, Color::Black) => {
                        // Case 2a: W's children are both black

                        // X carries a phantom black, and its sibling W is black
                        // we can remove one black from X and W by colouring W
                        // red and giving X's phantom black to the
                        // parent.
                        w.set_color(Color::Red);
                        x = parent;
                        continue;
                    }
                    (Color::Red, Color::Black) | (Color::Black, Color::Red)
                        if w.child(x.side().unwrap()).unwrap().is_red() =>
                    {
                        // Case 2b: X's near-cousin is red and its far-cousin is black

                        // colour the near-cousin black and w red, rotate so
                        // that the near cousin becomes the sibling of x
                        let mut near_cousin = w
                            .child(x.side().unwrap())
                            .expect("near cousin exists because it is red");

                        near_cousin.set_color(Color::Black);
                        w.set_color(Color::Red);
                        w = near_cousin
                            .rotate_up(|_| panic!("this shouldn't be root"))
                            .unwrap_or_else(|_| panic!("near cousin is the child of w"));

                        // Fall through to case 2c
                    }
                    _ => {}
                }

                // Case 2c: X's far-cousin is red

                // we can now rotate around the parent to balance the
                // subtree at parent without increasing the
                // black-height.
                // However, in the case that the far-cousin is red and
                // the parent is black, this will result, after
                // rotating, in the far-cousin's path having fewer black
                // nodes than the parent's path: we have taken a black
                // node from above the far-cousin and moved it into its
                // sibling branch.
                // This is fixed by colouring w the colour of the
                // parent, and colouring the parent and X's far-cousin
                // black.

                w.set_color(parent.color());
                parent.set_color(Color::Black);
                let mut far_cousin = w
                    .child(x.side().unwrap().not())
                    .expect("far cousin exists because it is red");
                far_cousin.set_color(Color::Black);

                w.rotate_up(|root| self.set_root_handle(root))
                    .unwrap_or_else(|_| panic!("w is the child of parent"));
            }
        }

        z.node()
    }

    pub fn iter_nodes(&self) -> TreeNodeIter<'_, N> {
        TreeNodeIter {
            range: TreeRange::full_range(self.root_handle()),
        }
    }

    pub fn iter(&self) -> TreeIter<'_, N> {
        TreeIter {
            range: TreeRange::full_range(self.root_handle()),
        }
    }
}

impl<N: UnsafeNode> Default for RBTree<N> {
    fn default() -> Self {
        Self::new()
    }
}

#[derive(Debug, Clone)]
enum RangeHandle<N> {
    Root(Handle<N>),
    Node(Handle<N>),
}

impl<N> RangeHandle<N> {
    fn into_inner(self) -> Handle<N> {
        match self {
            RangeHandle::Root(handle) => handle,
            RangeHandle::Node(handle) => handle,
        }
    }
}

impl<N: UnsafeNode + Eq> Eq for RangeHandle<N> {}

impl<N: UnsafeNode> PartialEq for RangeHandle<N> {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Self::Root(l0), Self::Root(r0)) => l0 == r0,
            (Self::Node(l0), Self::Node(r0)) => l0 == r0,
            _ => false,
        }
    }
}

struct TreeRange<'a, N: UnsafeNode> {
    start: Option<RangeHandle<N>>,
    end: Option<RangeHandle<N>>,
    _pd: PhantomData<&'a ()>,
}

impl<'a, N: UnsafeNode + 'a> TreeRange<'a, N> {
    #[expect(dead_code)]
    fn new(start: Handle<N>, end: Handle<N>) -> Self {
        Self {
            start: Some(RangeHandle::Node(start)),
            end: Some(RangeHandle::Node(end)),
            _pd: PhantomData,
        }
    }
    fn full_range(root: Handle<N>) -> Self {
        Self {
            start: Some(RangeHandle::Root(root.clone())),
            end: Some(RangeHandle::Root(root)),
            _pd: PhantomData,
        }
    }

    fn init_front(&mut self) -> Option<&mut Handle<N>> {
        if let Some(RangeHandle::Root(root)) = &self.start {
            self.start = Some(RangeHandle::Node(
                root.minimum_of().unwrap_or_else(|| root.clone()),
            ));
        }

        match &mut self.start {
            None => None,
            Some(RangeHandle::Node(handle)) => Some(handle),
            _ => unreachable!(),
        }
    }

    fn init_back(&mut self) -> Option<&mut Handle<N>> {
        if let Some(RangeHandle::Root(root)) = &self.end {
            self.end = Some(RangeHandle::Node(
                root.maximum_of().unwrap_or_else(|| root.clone()),
            ));
        }

        match &mut self.end {
            None => None,
            Some(RangeHandle::Node(handle)) => Some(handle),
            _ => unreachable!(),
        }
    }

    fn next(&mut self) -> Option<Handle<N>> {
        let next = self.init_front()?.next_of().map(RangeHandle::Node);
        let current = mem::replace(&mut self.start, next);

        match self.end {
            None => None,
            _ => {
                if self.start == self.end {
                    self.end = None;
                }

                current.map(RangeHandle::into_inner)
            }
        }
    }

    fn next_back(&mut self) -> Option<Handle<N>> {
        let next = self.init_back()?.next_back_of().map(RangeHandle::Node);
        let current = mem::replace(&mut self.end, next);

        match self.start {
            None => None,
            _ => {
                if self.start == self.end {
                    self.start = None;
                }

                current.map(RangeHandle::into_inner)
            }
        }
    }
}

pub struct TreeNodeIter<'a, N: UnsafeNode> {
    range: TreeRange<'a, N>,
}

impl<'a, N: UnsafeNode + 'a> Iterator for TreeNodeIter<'a, N> {
    type Item = Handle<N>;

    fn next(&mut self) -> Option<Self::Item> {
        self.range.next()
    }
}

impl<'a, N: UnsafeNode + 'a> DoubleEndedIterator for TreeNodeIter<'a, N> {
    fn next_back(&mut self) -> Option<Self::Item> {
        self.range.next_back()
    }
}

pub struct TreeIter<'a, N: UnsafeNode> {
    range: TreeRange<'a, N>,
}

impl<'a, N: UnsafeNode + 'a> Iterator for TreeIter<'a, N> {
    type Item = &'a N::Key;

    fn next(&mut self) -> Option<Self::Item> {
        self.range
            .next()
            .map(|n| unsafe { n.node().unwrap().as_ref().key() })
    }
}

impl<'a, N: UnsafeNode + 'a> DoubleEndedIterator for TreeIter<'a, N> {
    fn next_back(&mut self) -> Option<Self::Item> {
        self.range
            .next_back()
            .map(|n| unsafe { n.node().unwrap().as_ref().key() })
    }
}

#[cfg(test)]
mod tests {
    use std::cell::Cell;

    use super::*;

    #[derive(Debug)]
    struct TestNode {
        key: i32,
        left: Cell<Option<NonNull<TestNode>>>,
        right: Cell<Option<NonNull<TestNode>>>,
        parent: Cell<Option<NonNull<TestNode>>>,
        color: Cell<Color>,
    }

    unsafe impl UnsafeNode for TestNode {
        type Key = i32;

        fn left(&self) -> Option<NonNull<Self>> {
            self.left.get()
        }

        fn right(&self) -> Option<NonNull<Self>> {
            self.right.get()
        }

        fn parent(&self) -> Option<NonNull<Self>> {
            self.parent.get()
        }

        fn key(&self) -> &Self::Key {
            &self.key
        }

        fn color(&self) -> Color {
            self.color.get()
        }

        fn set_left(&self, left: Option<NonNull<Self>>) {
            self.left.set(left);
        }

        fn set_right(&self, right: Option<NonNull<Self>>) {
            self.right.set(right);
        }

        fn set_parent(&self, parent: Option<NonNull<Self>>) {
            self.parent.set(parent);
        }

        fn set_color(&self, color: Color) {
            self.color.set(color);
        }
    }

    impl TestNode {
        fn new(key: i32) -> Self {
            Self {
                key,
                left: Cell::new(None),
                right: Cell::new(None),
                parent: Cell::new(None),
                color: Cell::new(Color::Red),
            }
        }
    }

    #[test]
    fn next_of() {
        let mut tree = RBTree::<TestNode>::new();
        _ = tree.insert_node(Box::into_non_null(Box::new(TestNode::new(1))));

        assert_eq!(tree.root_handle().next_of(), None);
        assert_eq!(tree.root_handle().next_back_of(), None);
    }

    #[test]
    fn rotate() {
        struct DummyTree {
            a: NonNull<TestNode>,
            x: NonNull<TestNode>,
            b: NonNull<TestNode>,
            y: NonNull<TestNode>,
            c: NonNull<TestNode>,
        }

        impl DummyTree {
            fn new() -> Self {
                let a = Box::into_non_null(Box::new(TestNode::new(1)));
                let x = Box::into_non_null(Box::new(TestNode::new(2)));
                let b = Box::into_non_null(Box::new(TestNode::new(3)));
                let y = Box::into_non_null(Box::new(TestNode::new(4)));
                let c = Box::into_non_null(Box::new(TestNode::new(5)));

                unsafe {
                    x.as_ref().set_left(Some(a));
                    x.as_ref().set_right(Some(y));
                    y.as_ref().set_left(Some(b));
                    y.as_ref().set_right(Some(c));

                    a.as_ref().set_parent(Some(x));
                    y.as_ref().set_parent(Some(x));

                    b.as_ref().set_parent(Some(y));
                    c.as_ref().set_parent(Some(y));
                }

                Self { a, x, b, y, c }
            }
        }

        impl Drop for DummyTree {
            fn drop(&mut self) {
                unsafe {
                    _ = Box::from_raw(self.a.as_ptr());
                    _ = Box::from_raw(self.x.as_ptr());
                    _ = Box::from_raw(self.b.as_ptr());
                    _ = Box::from_raw(self.y.as_ptr());
                    _ = Box::from_raw(self.c.as_ptr());
                }
            }
        }

        struct PanicOnDrop;
        impl Drop for PanicOnDrop {
            fn drop(&mut self) {
                panic!("PanicOnDrop dropped");
            }
        }

        let mut panic = Some(PanicOnDrop);

        let tree = DummyTree::new();
        let new_root = Handle::Root(tree.x).rotate_left(|_| mem::forget(panic.take()));
        assert_eq!(new_root, Ok(Handle::Root(tree.y)));

        assert_eq!(unsafe { tree.y.as_ref().left() }, Some(tree.x));
        assert_eq!(unsafe { tree.y.as_ref().right() }, Some(tree.c));
        assert_eq!(unsafe { tree.x.as_ref().left() }, Some(tree.a));
        assert_eq!(unsafe { tree.x.as_ref().right() }, Some(tree.b));

        assert_eq!(unsafe { tree.a.as_ref().parent() }, Some(tree.x));
        assert_eq!(unsafe { tree.b.as_ref().parent() }, Some(tree.x));
        assert_eq!(unsafe { tree.c.as_ref().parent() }, Some(tree.y));
        assert_eq!(unsafe { tree.x.as_ref().parent() }, Some(tree.y));
        assert_eq!(unsafe { tree.y.as_ref().parent() }, None);

        let mut panic = Some(PanicOnDrop);
        let new_root = Handle::Root(tree.y).rotate_right(|_| mem::forget(panic.take()));
        assert_eq!(new_root, Ok(Handle::Root(tree.x)));

        assert_eq!(unsafe { tree.x.as_ref().left() }, Some(tree.a));
        assert_eq!(unsafe { tree.x.as_ref().right() }, Some(tree.y));
        assert_eq!(unsafe { tree.y.as_ref().left() }, Some(tree.b));
        assert_eq!(unsafe { tree.y.as_ref().right() }, Some(tree.c));

        assert_eq!(unsafe { tree.a.as_ref().parent() }, Some(tree.x));
        assert_eq!(unsafe { tree.y.as_ref().parent() }, Some(tree.x));
        assert_eq!(unsafe { tree.b.as_ref().parent() }, Some(tree.y));
        assert_eq!(unsafe { tree.c.as_ref().parent() }, Some(tree.y));
        assert_eq!(unsafe { tree.x.as_ref().parent() }, None);

        let mut panic = Some(PanicOnDrop);
        let root = Handle::from_node(Some(tree.y)).rotate_up(|_| mem::forget(panic.take()));
        assert_eq!(root, Ok(Handle::Root(tree.y)));

        assert_eq!(unsafe { tree.y.as_ref().left() }, Some(tree.x));
        assert_eq!(unsafe { tree.y.as_ref().right() }, Some(tree.c));
        assert_eq!(unsafe { tree.x.as_ref().left() }, Some(tree.a));
        assert_eq!(unsafe { tree.x.as_ref().right() }, Some(tree.b));

        assert_eq!(unsafe { tree.a.as_ref().parent() }, Some(tree.x));
        assert_eq!(unsafe { tree.b.as_ref().parent() }, Some(tree.x));
        assert_eq!(unsafe { tree.c.as_ref().parent() }, Some(tree.y));
        assert_eq!(unsafe { tree.x.as_ref().parent() }, Some(tree.y));
        assert_eq!(unsafe { tree.y.as_ref().parent() }, None);
    }

    #[test]
    fn insert() {
        let mut tree = RBTree::<TestNode>::new();

        let nodes: Vec<_> = (0..10)
            .map(|i| Box::into_raw(Box::new(TestNode::new(i))))
            .collect();

        for &node in &nodes {
            eprintln!("Inserting node with key: {}", unsafe { (*node).key });
            _ = tree.insert_node(unsafe { NonNull::new_unchecked(node) });

            eprintln!("Tree after insertion:");
            for n in tree.iter_nodes() {
                eprintln!(
                    "\tNode: {:?} => {:?}",
                    unsafe { n.node().map(|n| n.as_ref().key) },
                    n
                );
            }
        }

        for i in 0..10 {
            let result = tree.find_by_key(&i);
            eprintln!("{i}: {result:?}");
            assert!(matches!(result, SearchResult::FoundAt(_)));
        }
    }

    #[test]
    fn remove() {
        let mut tree = RBTree::<TestNode>::new();

        let nodes: Vec<_> = (0..10)
            .map(|i| Box::into_raw(Box::new(TestNode::new(i))))
            .collect();

        for &node in &nodes {
            _ = tree.insert_node(unsafe { NonNull::new_unchecked(node) });
        }

        for i in 0..10 {
            assert!(matches!(tree.find_by_key(&i), SearchResult::FoundAt(_)));

            eprintln!("Removing node with key: {}", i);
            let removed_node = tree.remove(&i);
            assert!(removed_node.is_some());

            assert!(matches!(
                tree.find_by_key(&i),
                SearchResult::NotFoundAt(_) | SearchResult::Empty
            ));

            eprintln!("Tree after removal:");
            for n in tree.iter_nodes() {
                eprintln!(
                    "\tNode: {:?} => {:?}",
                    unsafe { n.node().map(|n| n.as_ref().key) },
                    n
                );
            }
        }

        for i in 0..10 {
            let result = tree.find_by_key(&i);
            eprintln!("{i}: {result:?}");
            assert!(matches!(result, SearchResult::Empty));
        }
    }
}