const std = @import("std");
// const t = stdx.testing;
const t = std.testing;

const log = std.log.scoped(.rb_tree);

const Color = enum(u1) {
    Black,
    Red,
};

const CompactUnorderedList = @import("compact.zig").CompactUnorderedList;
const CompactNull = @import("compact.zig").CompactNull;

/// Based on Zig's rb node pointer based implementation:
/// https://github.com/ziglang/std-lib-orphanage/blob/master/std/rb.zig
/// Deletion logic was redone from https://www.youtube.com/watch?v=CTvfzU_uNKE as guidance.
/// Correction for case 5 in the video: it should not have a parent black condition.
/// Visualize: https://www.cs.usfca.edu/~galles/visualization/RedBlack.html
pub fn RbTree(comptime Id: type, comptime Value: type, comptime Context: type, comptime Compare: fn (Value, Value, Context) std.math.Order) type {
    return struct {
        root: OptId,
        buf: CompactUnorderedList(Id, Node),
        /// The current number of nodes. Does not count detached nodes.
        size: usize,
        ctx: Context,

        const Self = @This();
        const OptId = Id;
        const NullId = CompactNull(Id);

        const Node = struct {
            left: OptId,
            right: OptId,
            val: Value,

            parent: OptId,
            color: Color,

            fn getParentOpt(self: Node) ?Id {
                if (self.parent == NullId) {
                    return null;
                } else {
                    return self.parent;
                }
            }

            fn isRoot(self: Node) bool {
                return self.parent == NullId;
            }

            fn setChild(self: *Node, child: Id, is_left: bool) void {
                if (is_left) {
                    self.left = child;
                } else {
                    self.right = child;
                }
            }
        };

        pub fn init(alloc: std.mem.Allocator, ctx: Context) Self {
            return .{
                .root = NullId,
                .buf = CompactUnorderedList(Id, Node).init(alloc),
                .size = 0,
                .ctx = ctx,
            };
        }

        pub fn deinit(self: *Self) void {
            self.buf.deinit();
        }

        pub fn clearRetainingCapacity(self: *Self) void {
            self.buf.clearRetainingCapacity();
            self.root = NullId;
            self.size = 0;
        }

        /// Re-sorts a tree with a compare function.
        pub fn sort(self: *Self, ctx: anytype, cmp: fn (Value, Value, @TypeOf(ctx)) std.math.Order) !void {
            self.root = NullId;
            var iter = self.buf.iterator();
            while (iter.next()) |node| {
                self.buf.remove(iter.cur_id);
                _ = try self.insertCustom(node.val, ctx, cmp);
            }
        }

        pub fn first(self: Self) ?Id {
            if (self.root == NullId) {
                return null;
            }
            var id = self.root;
            var node = self.buf.getNoCheck(id);
            while (node.left != NullId) {
                id = node.left;
                node = self.buf.getNoCheck(id);
            }
            return id;
        }

        pub fn firstFrom(self: Self, node_id: Id) Id {
            var id = node_id;
            var node = self.buf.getNoCheck(id);
            while (node.left != NullId) {
                id = node.left;
                node = self.buf.getNoCheck(id);
            }
            return id;
        }

        pub fn last(self: Self) ?Id {
            if (self.root == NullId) {
                return null;
            }
            var id = self.root;
            var node = self.buf.getNoCheck(id);
            while (node.right != NullId) {
                id = node.right;
                node = self.buf.getNoCheck(id);
            }
            return id;
        }

        pub fn allocValuesInOrder(self: Self, alloc: std.mem.Allocator) []const Value {
            var vals = std.ArrayList(Value).initCapacity(alloc, self.size) catch unreachable;
            var cur = self.first();
            while (cur) |id| {
                vals.appendAssumeCapacity(self.get(id).?);
                cur = self.getNext(id);
            }
            return vals.toOwnedSlice();
        }

        pub fn allocNodeIdsInOrder(self: Self, alloc: std.mem.Allocator) []const Id {
            var node_ids = std.ArrayList(Id).initCapacity(alloc, self.size) catch unreachable;
            var cur = self.first();
            while (cur) |id| {
                node_ids.appendAssumeCapacity(id);
                cur = self.getNext(id);
            }
            return node_ids.toOwnedSlice();
        }

        /// User code typically wouldn't need the node, but it is available for testing.
        pub fn getNode(self: Self, id: Id) ?Node {
            return self.buf.get(id);
        }

        pub fn get(self: Self, id: Id) ?Value {
            if (self.buf.get(id)) |node| {
                return node.val;
            } else return null;
        }

        pub fn getNoCheck(self: Self, id: Id) Value {
            return self.buf.getNoCheck(id).val;
        }

        pub fn getPtr(self: Self, id: Id) ?*Value {
            if (self.buf.getPtr(id)) |node| {
                return &node.val;
            } else return null;
        }

        pub fn getPtrNoCheck(self: Self, id: Id) *Value {
            return &self.buf.getPtrNoCheck(id).val;
        }

        fn getSibling(self: Self, parent: *Node, child_id: Id) OptId {
            _ = self;
            if (parent.left == child_id) {
                return parent.right;
            } else {
                return parent.left;
            }
        }

        // Replacement node (with it's children) are relinked in place of node . Node (with it's children) become detached from tree.
        fn transplant(self: *Self, node_id: Id, node: *Node, r_id: Id, mb_rnode: ?*Node) void {
            if (node_id == self.root) {
                self.root = r_id;
            } else {
                const parent = self.buf.getPtrNoCheck(node.parent);
                parent.setChild(r_id, parent.left == node_id);
            }
            if (mb_rnode) |rnode| {
                rnode.parent = node.parent;
            }
        }

        pub fn removeDetached(self: *Self, node_id: Id) void {
            self.buf.remove(node_id);
        }

        pub fn remove(self: *Self, node_id: Id) anyerror!void {
            try self.detach(node_id);
            self.buf.remove(node_id);
        }

        /// If node is not part of tree, error is returned.
        /// Detaches the node from the tree, rebalances tree, but does not remove it. The caller is responsible for calling removeDetached later on.
        pub fn detach(self: *Self, node_id: Id) anyerror!void {
            try self.detachInternal(node_id);
            self.size -= 1;
        }

        fn detachInternal(self: *Self, node_id: Id) anyerror!void {
            var node = self.buf.getPtr(node_id) orelse return error.DoesNotExist;

            var to_fix_nid: OptId = undefined;

            if (node.left == NullId) {
                const mb_rnode: ?*Node = if (node.right != NullId) self.buf.getPtrNoCheck(node.right) else null;
                self.transplant(node_id, node, node.right, mb_rnode);
                to_fix_nid = node.right;
            } else if (node.right == NullId) {
                const mb_rnode: ?*Node = if (node.left != NullId) self.buf.getPtrNoCheck(node.left) else null;
                self.transplant(node_id, node, node.left, mb_rnode);
                to_fix_nid = node.left;
            } else {
                const r_id = self.firstFrom(node.right);
                const r_node = self.buf.getPtrNoCheck(r_id);

                // Normally this would be a value copy but since ids are tied to their nodes, the replacement node is relinked in place of the target node.
                // Transplant only relinks parent of replacement node.
                const r_parent = r_node.parent;
                self.transplant(node_id, node, r_id, r_node);
                if (r_parent != node_id) {
                    const rp = self.buf.getPtrNoCheck(r_parent);
                    rp.setChild(node_id, rp.left == r_id);
                    node.parent = r_parent;
                } else {
                    node.parent = r_id;
                }

                // Swap colors.
                const tmp_color = r_node.color;
                r_node.color = node.color;
                node.color = tmp_color;

                // Copy r_node value to node; a recursive call to delete node will be called.
                var orig_val = node.val;
                node.val = r_node.val;
                defer {
                    // Revert back to node's original val in case the user wanted to only detach the node (and not remove it).
                    node.val = orig_val;
                }

                // Relink r_node left and node left.
                // Note: r_node shouldn't have a left child since firstFrom would have returned it the child instead.
                self.buf.getPtrNoCheck(node.left).parent = r_id;
                r_node.left = node.left;
                node.left = NullId;

                // Relink r_node right and node right.
                // Note: tmp_right can't be null since node should have two children.
                const tmp_right = node.right;
                if (r_node.right != NullId) {
                    self.buf.getPtrNoCheck(r_node.right).parent = node_id;
                }
                node.right = r_node.right;
                if (tmp_right != r_id) {
                    self.buf.getPtrNoCheck(tmp_right).parent = r_id;
                    r_node.right = tmp_right;
                }

                // Reduced to zero or one children case. Recurse once.
                try self.detachInternal(node_id);
                return;
            }

            // If red was removed, it is done.
            if (node.color == .Red) {
                return;
            } else {
                // If black was removed and replacement is red. Paint replacement black and we are done.
                if (to_fix_nid == NullId) {
                    // Double black. Perform fix.
                    self.detachFixUp(to_fix_nid, node.parent);
                } else {
                    const r_node = self.buf.getPtrNoCheck(to_fix_nid);
                    if (r_node.color == .Red) {
                        // Can mark black and be done since a black was removed.
                        r_node.color = .Black;
                    } else {
                        // Double black. Perform fix.
                        self.detachFixUp(to_fix_nid, r_node.parent);
                    }
                }
            }
        }

        /// Handle double black cases.
        /// Assumes node is a double black. Since the node could be null, the current parent is also required.
        fn detachFixUp(self: *Self, node_id: OptId, parent_id: OptId) void {
            // Case 1: Root case.
            if (parent_id == NullId) {
                return;
            }

            const parent = self.buf.getPtrNoCheck(parent_id);
            const s_id = self.getSibling(parent, node_id);
            const is_right_sibling = parent.left == node_id;
            // Sibling must exist since node is a double black.
            const sibling = self.buf.getPtrNoCheck(s_id);
            const s_left: ?*Node = if (sibling.left != NullId) self.buf.getPtrNoCheck(sibling.left) else null;
            const s_right: ?*Node = if (sibling.right != NullId) self.buf.getPtrNoCheck(sibling.right) else null;
            const s_left_black = s_left == null or s_left.?.color == .Black;
            const s_right_black = s_right == null or s_right.?.color == .Black;

            if (parent.color == .Black and sibling.color == .Red) {
                if (s_left_black and s_right_black) {
                    if (is_right_sibling) {
                        // Case 2: parent is black, right sibling is red and has two black children.
                        self.rotateLeft(parent_id, parent);
                    } else {
                        // Case 2: parent is black, left sibling is red and has two black children.
                        self.rotateRight(parent_id, parent);
                    }
                    parent.color = .Red;
                    sibling.color = .Black;
                    self.detachFixUp(node_id, parent_id);
                    return;
                }
            }

            if (parent.color == .Black and sibling.color == .Black) {
                // Case 3: left or right sibling with both black children.
                if (s_left_black and s_right_black) {
                    sibling.color = .Red;
                    // Recurse at parent.
                    self.detachFixUp(parent_id, parent.parent);
                    return;
                }
            }

            if (parent.color == .Red and sibling.color == .Black) {
                // Case 4: left or right sibling with both black children.
                if (s_left_black and s_right_black) {
                    parent.color = .Black;
                    sibling.color = .Red;
                    return;
                }
            }

            if (sibling.color == .Black) {
                // Case 5: parent is black, right sibling is black, sibling has red left child and black right child.
                if (is_right_sibling and s_left != null and s_left.?.color == .Red and s_right_black) {
                    self.rotateRight(s_id, sibling);
                    sibling.color = .Red;
                    s_left.?.color = .Black;
                    // Call again to check for case 6.
                    self.detachFixUp(node_id, parent_id);
                    return;
                }
                // Case 5: parent is black, left sibiling is black, sibling has red right child and black left child.
                if (!is_right_sibling and s_right != null and s_right.?.color == .Red and s_left_black) {
                    self.rotateLeft(s_id, sibling);
                    sibling.color = .Red;
                    s_right.?.color = .Black;
                    // Call again to check for case 6.
                    self.detachFixUp(node_id, parent_id);
                    return;
                }
                // Case 6: right sibling with red right child.
                if (is_right_sibling and s_right != null and s_right.?.color == .Red) {
                    self.rotateLeft(parent_id, parent);
                    sibling.color = parent.color;
                    parent.color = .Black;
                    s_right.?.color = .Black;
                    return;
                }
                // Case 6: left sibling with red left child.
                if (!is_right_sibling and s_left != null and s_left.?.color == .Red) {
                    self.rotateRight(parent_id, parent);
                    sibling.color = parent.color;
                    parent.color = .Black;
                    s_left.?.color = .Black;
                    return;
                }
            }
        }

        pub fn insert(self: *Self, val: Value) !Id {
            return try self.insertCustom(val, self.ctx, Compare);
        }

        /// Duplicate keys are not allowed.
        pub fn insertCustom(self: *Self, val: Value, ctx: anytype, cmp: fn (Value, Value, @TypeOf(ctx)) std.math.Order) !Id {
            var maybe_id: ?Id = undefined;
            var maybe_parent: ?Id = undefined;
            var is_left: bool = undefined;

            maybe_id = self.doLookup(val, &maybe_parent, &is_left, ctx, cmp);
            if (maybe_id) |_| {
                return error.DuplicateKey;
            }

            const new_id = try self.buf.add(.{
                .left = NullId,
                .right = NullId,
                .color = .Red,
                .parent = maybe_parent orelse NullId,
                .val = val,
            });
            self.size += 1;

            if (maybe_parent) |parent| {
                self.buf.getPtrNoCheck(parent).setChild(new_id, is_left);
            } else {
                self.root = new_id;
            }

            var node_id = new_id;
            while (true) {
                var node = self.buf.getPtrNoCheck(node_id);
                const parent_id = node.getParentOpt() orelse break;
                var parent = self.buf.getPtrNoCheck(parent_id);
                if (parent.color == .Black) {
                    // Current is red, parent is black. Nothing left to do.
                    break;
                }
                // If parent is red, there must be a grand parent that is black.
                var grandpa_id = parent.getParentOpt() orelse unreachable;
                var grandpa = self.buf.getPtrNoCheck(grandpa_id);

                if (parent_id == grandpa.left) {
                    var opt_psibling = grandpa.right;
                    const mb_psibling: ?*Node = if (opt_psibling != NullId) self.buf.getPtrNoCheck(opt_psibling) else null;
                    if (mb_psibling == null or mb_psibling.?.color == .Black) {
                        // Case #5, parent is red, parent sibling is black, node is inner grandchild. Rotate left first.
                        if (node_id == parent.right) {
                            self.rotateLeft(parent_id, parent);
                            // Make sure reattached right is black since parent is red.
                            if (parent.right != NullId) {
                                self.buf.getPtrNoCheck(parent.right).color = .Black;
                            }
                            parent = node;
                        }
                        parent.color = .Black;
                        grandpa.color = .Red;
                        self.rotateRight(grandpa_id, grandpa);
                        // Make sure reattached left is black since grandpa is red.
                        if (grandpa.left != NullId) {
                            self.buf.getPtrNoCheck(grandpa.left).color = .Black;
                        }
                    } else {
                        // parent and parent sibling are both red. Set both to black, and grand parent to red.
                        parent.color = .Black;
                        mb_psibling.?.color = .Black;
                        grandpa.color = .Red;
                        node_id = grandpa_id;
                    }
                } else {
                    var opt_psibling = grandpa.left;
                    const mb_psibling: ?*Node = if (opt_psibling != NullId) self.buf.getPtrNoCheck(opt_psibling) else null;
                    if (mb_psibling == null or mb_psibling.?.color == .Black) {
                        // Case #5, parent is red, parent sibling is black, node is inner grandchild. Rotate right first.
                        if (node_id == parent.left) {
                            self.rotateRight(parent_id, parent);
                            // Make sure reattached left is black since parent is red.
                            if (parent.left != NullId) {
                                self.buf.getPtrNoCheck(parent.left).color = .Black;
                            }
                            parent = node;
                        }
                        parent.color = .Black;
                        grandpa.color = .Red;
                        self.rotateLeft(grandpa_id, grandpa);
                        // Make sure reattached right is black since grandpa is red.
                        if (grandpa.right != NullId) {
                            self.buf.getPtrNoCheck(grandpa.right).color = .Black;
                        }
                    } else {
                        // parent and parent sibling are both red. Set both to black, and grand parent to red.
                        parent.color = .Black;
                        mb_psibling.?.color = .Black;
                        grandpa.color = .Red;
                        node_id = grandpa_id;
                    }
                }
            }
            // This was an insert, there is at least one node.
            self.buf.getPtrNoCheck(self.root).color = .Black;
            return new_id;
        }

        /// lookup searches for the value of key, using binary search. It will
        /// return a pointer to the node if it is there, otherwise it will return null.
        /// Complexity guaranteed O(log n), where n is the number of nodes book-kept
        /// by tree.
        pub fn lookup(self: Self, val: Value) ?Id {
            var parent: ?Id = undefined;
            var is_left: bool = undefined;
            return self.doLookup(val, &parent, &is_left, self.ctx, Compare);
        }

        /// Lookup with a different compare function.
        pub fn lookupCustom(self: Self, val: Value, ctx: anytype, cmp: fn (Value, Value, @TypeOf(ctx)) std.math.Order) ?Id {
            var parent: ?Id = undefined;
            var is_left: bool = undefined;
            return self.doLookup(val, &parent, &is_left, ctx, cmp);
        }

        /// Lookup that also sets the parent and is_left to indicate where it was found or where it would be inserted.
        pub fn lookupCustomLoc(self: Self, val: Value, ctx: anytype, cmp: fn (Value, Value, @TypeOf(ctx)) std.math.Order, parent: *?Id, is_left: *bool) ?Id {
            return self.doLookup(val, parent, is_left, ctx, cmp);
        }

        /// If there is a match, the Id is returned.
        /// Otherwise, the insert slot is provided by pparent and is_left.
        fn doLookup(self: Self, val: Value, pparent: *?Id, is_left: *bool, ctx: anytype, cmp: fn (Value, Value, @TypeOf(ctx)) std.math.Order) ?Id {
            var opt_id = self.root;

            pparent.* = null;
            is_left.* = false;

            while (opt_id != NullId) {
                const node = self.buf.getNoCheck(opt_id);
                const res = cmp(val, node.val, ctx);
                if (res == .eq) {
                    return opt_id;
                }
                pparent.* = opt_id;
                switch (res) {
                    .lt => {
                        is_left.* = true;
                        opt_id = node.left;
                    },
                    .gt => {
                        is_left.* = false;
                        opt_id = node.right;
                    },
                    .eq => unreachable, // handled above
                }
            }
            return null;
        }

        ///     e
        ///    /
        ///  (a)
        ///  / \
        /// b   c
        ///    /
        ///   d
        /// Given a, rotate a to b and c to a. a's parent becomes c and right becomes d. c's parent bcomes e and left becomes a.
        fn rotateLeft(self: *Self, node_id: Id, node: *Node) void {
            if (node.right == NullId) {
                unreachable;
            }
            var right = self.buf.getPtrNoCheck(node.right);
            if (!node.isRoot()) {
                var parent = self.buf.getPtrNoCheck(node.parent);
                if (parent.left == node_id) {
                    parent.left = node.right;
                } else {
                    parent.right = node.right;
                }
                right.parent = node.parent;
            } else {
                self.root = node.right;
                right.parent = NullId;
            }
            node.parent = node.right;
            node.right = right.left;
            if (node.right != NullId) {
                self.buf.getPtrNoCheck(node.right).parent = node_id;
            }
            right.left = node_id;
        }

        /// Works similarily to rotateLeft for the right direction.
        fn rotateRight(self: *Self, node_id: Id, node: *Node) void {
            if (node.left == NullId) {
                unreachable;
            }
            var left = self.buf.getPtrNoCheck(node.left);
            if (!node.isRoot()) {
                var parent = self.buf.getPtrNoCheck(node.parent);
                if (parent.left == node_id) {
                    parent.left = node.left;
                } else {
                    parent.right = node.left;
                }
                left.parent = node.parent;
            } else {
                self.root = node.left;
                left.parent = NullId;
            }
            node.parent = node.left;
            node.left = left.right;
            if (node.left != NullId) {
                self.buf.getPtrNoCheck(node.left).parent = node_id;
            }
            left.right = node_id;
        }

        pub fn getPrev(self: Self, id: Id) ?Id {
            var node = self.buf.getNoCheck(id);
            if (node.left != NullId) {
                var cur = node.left;
                node = self.buf.getNoCheck(cur);
                while (node.right != NullId) {
                    cur = node.right;
                    node = self.buf.getNoCheck(cur);
                }
                return cur;
            }
            var cur = id;
            while (true) {
                if (node.parent != NullId) {
                    var p = self.buf.getNoCheck(node.parent);
                    if (cur != p.left) {
                        return node.parent;
                    }
                    cur = node.parent;
                    node = p;
                } else {
                    return null;
                }
            }
        }

        pub fn getNext(self: Self, id: Id) ?Id {
            var node = self.buf.getNoCheck(id);
            if (node.right != NullId) {
                var cur = node.right;
                node = self.buf.getNoCheck(cur);
                while (node.left != NullId) {
                    cur = node.left;
                    node = self.buf.getNoCheck(cur);
                }
                return cur;
            }
            var cur = id;
            while (true) {
                if (node.parent != NullId) {
                    var p = self.buf.getNoCheck(node.parent);
                    if (cur != p.right) {
                        return node.parent;
                    }
                    cur = node.parent;
                    node = p;
                } else {
                    return null;
                }
            }
        }

        /// Checks whether the tree is a valid red black tree.
        pub fn isValid(self: Self) bool {
            if (self.root == NullId) {
                return true;
            }
            const root = self.buf.getNoCheck(self.root);
            if (root.color != .Black) {
                return false;
            }
            var black_cnt: u32 = undefined;
            return self.isValid2(self.root, &black_cnt);
        }

        fn isValid2(self: Self, id: OptId, out_black_cnt: *u32) bool {
            if (id == NullId) {
                out_black_cnt.* = 1;
                return true;
            }
            const node = self.buf.getNoCheck(id);
            if (node.color == .Red) {
                // Red node. Both children must be black.
                if (node.left != NullId) {
                    const left = self.buf.getNoCheck(node.left);
                    if (left.color != .Black) {
                        return false;
                    }
                }
                if (node.right != NullId) {
                    const right = self.buf.getNoCheck(node.right);
                    if (right.color != .Black) {
                        return false;
                    }
                }
            }
            var left_black_cnt: u32 = undefined;
            if (!self.isValid2(node.left, &left_black_cnt)) {
                return false;
            }
            var right_black_cnt: u32 = undefined;
            if (!self.isValid2(node.right, &right_black_cnt)) {
                return false;
            }
            if (left_black_cnt != right_black_cnt) {
                return false;
            }
            if (node.color == .Black) {
                out_black_cnt.* = left_black_cnt + 1;
            } else {
                out_black_cnt.* = left_black_cnt;
            }
            return true;
        }

        fn dump(self: Self) void {
            self.dump2(self.root);
        }

        fn dump2(self: Self, id: Id) void {
            if (id == NullId) {
                return;
            }
            const node = self.buf.getNoCheck(id);
            log.debug("{}-{} -> {} {}", .{ id, node.color, node.left, node.right });
            self.dump2(node.left);
            self.dump2(node.right);
        }

        fn insertDetached(self: *Self, val: Value, color: Color) !Id {
            return try self.buf.add(.{
                .left = NullId,
                .right = NullId,
                .color = color,
                .parent = NullId,
                .val = val,
            });
        }

        fn attachLeft(self: *Self, attached_parent: Id, detached: Id) !void {
            const parent = self.buf.getPtrNoCheck(attached_parent);
            if (parent.left != NullId) {
                return error.CantAttach;
            }
            parent.left = detached;
            const left = self.buf.getPtrNoCheck(detached);
            left.parent = attached_parent;
            self.size += 1;
        }

        fn attachRight(self: *Self, attached_parent: Id, detached: Id) !void {
            const parent = self.buf.getPtrNoCheck(attached_parent);
            if (parent.right != NullId) {
                return error.CantAttach;
            }
            parent.right = detached;
            const right = self.buf.getPtrNoCheck(detached);
            right.parent = attached_parent;
            self.size += 1;
        }
    };
}

fn testCompare(left: u32, right: u32, _: void) std.math.Order {
    if (left < right) {
        return .lt;
    } else if (left == right) {
        return .eq;
    } else if (left > right) {
        return .gt;
    }
    unreachable;
}

fn testCompareReverse(left: u32, right: u32, _: void) std.math.Order {
    return testCompare(right, left, {});
}

test "getNext, getPrev" {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(10);
    const b = try tree.insert(5);
    const c = try tree.insert(15);

    try t.expectEqual(tree.getNext(b).?, a);
    try t.expectEqual(tree.getNext(a).?, c);
    try t.expectEqual(tree.getNext(c), null);

    try t.expectEqual(tree.getPrev(b), null);
    try t.expectEqual(tree.getPrev(a).?, b);
    try t.expectEqual(tree.getPrev(c).?, a);
}

test "Insert where rotations need to set reattached nodes to black." {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(10);
    const b = try tree.insertDetached(9, .Black);
    const c = try tree.insertDetached(20, .Red);
    const d = try tree.insertDetached(8, .Red);
    const e = try tree.insertDetached(19, .Black);
    const f = try tree.insertDetached(21, .Black);
    const g = try tree.insertDetached(17, .Red);
    try tree.attachLeft(a, b);
    try tree.attachRight(a, c);
    try tree.attachLeft(b, d);
    try tree.attachLeft(c, e);
    try tree.attachRight(c, f);
    try tree.attachLeft(e, g);
    try t.expectEqual(tree.isValid(), true);

    const h = try tree.insert(18);
    try t.expectEqual(tree.isValid(), true);

    const vals = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(vals);
    try t.expectEqualSlices(TestId, vals, &.{ d, b, a, g, h, e, c, f });
}

test "Insert case #1: current node parent is black" {
    var tree = initTestTree();
    defer tree.deinit();

    const root = try tree.insert(10);
    const node = try tree.insert(9);

    try t.expectEqual(tree.getNode(root).?.color, .Black);
    try t.expectEqual(tree.getNode(root).?.left, node);
    try t.expectEqual(tree.getNode(node).?.color, .Red);
    try t.expectEqual(tree.getNode(node).?.parent, root);
}

test "Insert case #2/#3/#4: both parent and parent sibling are red" {
    var tree = initTestTree();
    defer tree.deinit();

    const root = try tree.insert(10);
    const parent = try tree.insert(9);
    const psibling = try tree.insert(11);
    const node = try tree.insert(8);

    try t.expectEqual(tree.getNode(root).?.color, .Black);
    try t.expectEqual(tree.getNode(root).?.left, parent);
    try t.expectEqual(tree.getNode(root).?.right, psibling);
    try t.expectEqual(tree.getNode(parent).?.color, .Black);
    try t.expectEqual(tree.getNode(parent).?.left, node);
    try t.expectEqual(tree.getNode(psibling).?.color, .Black);
    try t.expectEqual(tree.getNode(node).?.color, .Red);
    try t.expectEqual(tree.getNode(node).?.parent, parent);
}

test "Insert case #5: parent is red but parent sibling is black, node is inner grandchild" {
    var tree = initTestTree();
    defer tree.deinit();

    const root = try tree.insert(100);
    const parent = try tree.insert(50);
    _ = try tree.insert(150);
    _ = try tree.insert(25);
    const node = try tree.insert(75);
    _ = try tree.insert(80);
    _ = try tree.insert(70);
    _ = try tree.insert(60);

    try t.expectEqual(tree.getNode(node).?.parent, TestNullId);
    try t.expectEqual(tree.getNode(node).?.color, .Black);
    try t.expectEqual(tree.getNode(node).?.right, root);
    try t.expectEqual(tree.getNode(node).?.left, parent);

    try t.expectEqual(tree.getNode(root).?.color, .Red);
    try t.expectEqual(tree.getNode(root).?.parent, node);

    try t.expectEqual(tree.getNode(parent).?.color, .Red);
    try t.expectEqual(tree.getNode(parent).?.parent, node);
}

test "Insert case #6: parent is red but parent sibling is black, node is outer grandchild" {
    var tree = initTestTree();
    defer tree.deinit();

    const root = try tree.insert(100);
    const parent = try tree.insert(50);
    _ = try tree.insert(150);
    const node = try tree.insert(25);
    _ = try tree.insert(70);
    _ = try tree.insert(20);
    _ = try tree.insert(30);
    _ = try tree.insert(15);

    try t.expectEqual(tree.getNode(parent).?.parent, TestNullId);
    try t.expectEqual(tree.getNode(parent).?.color, .Black);
    try t.expectEqual(tree.getNode(parent).?.right, root);
    try t.expectEqual(tree.getNode(parent).?.left, node);

    try t.expectEqual(tree.getNode(root).?.color, .Red);
    try t.expectEqual(tree.getNode(root).?.parent, parent);

    try t.expectEqual(tree.getNode(node).?.color, .Red);
    try t.expectEqual(tree.getNode(node).?.parent, parent);
}

test "Insert in order." {
    var tree = initTestTree();
    defer tree.deinit();

    _ = try tree.insert(1);
    _ = try tree.insert(2);
    _ = try tree.insert(3);
    _ = try tree.insert(4);
    _ = try tree.insert(5);
    _ = try tree.insert(6);
    _ = try tree.insert(7);
    _ = try tree.insert(8);
    _ = try tree.insert(9);
    _ = try tree.insert(10);

    const vals = tree.allocValuesInOrder(t.allocator);
    defer t.allocator.free(vals);
    try t.expectEqualSlices(TestId, vals, &.{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 });
}

test "Insert in reverse order." {
    var tree = initTestTree();
    defer tree.deinit();

    _ = try tree.insert(10);
    _ = try tree.insert(9);
    _ = try tree.insert(8);
    _ = try tree.insert(7);
    _ = try tree.insert(6);
    _ = try tree.insert(5);
    _ = try tree.insert(4);
    _ = try tree.insert(3);
    _ = try tree.insert(2);
    _ = try tree.insert(1);

    const vals = tree.allocValuesInOrder(t.allocator);
    defer t.allocator.free(vals);
    try t.expectEqualSlices(TestId, vals, &.{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 });
}

test "inserting and looking up" {
    var tree = initTestTree();
    defer tree.deinit();
    const orig: u32 = 1000;
    const node_id = try tree.insert(orig);
    // Assert that identical value finds the same pointer
    try t.expectEqual(tree.lookup(1000), node_id);
    // Assert that insert duplicate returns error.
    try t.expectError(error.DuplicateKey, tree.insert(1000));
    try t.expectEqual(tree.lookup(1000), node_id);
    try t.expectEqual(tree.get(node_id).?, orig);
    // Assert that if looking for a non-existing value, return null.
    try t.expectEqual(tree.lookup(1234), null);
}

test "multiple inserts, followed by calling first and last" {
    // if (@import("builtin").arch == .aarch64) {
    //     // TODO https://github.com/ziglang/zig/issues/3288
    //     return error.SkipZigTest;
    // }
    var tree = initTestTree();
    defer tree.deinit();

    _ = try tree.insert(0);
    _ = try tree.insert(1);
    _ = try tree.insert(2);
    const third_id = try tree.insert(3);
    try t.expectEqual(tree.get(tree.first().?).?, 0);
    try t.expectEqual(tree.get(tree.last().?).?, 3);
    try t.expectEqual(tree.lookup(3), third_id);
    tree.sort({}, testCompareReverse) catch unreachable;
    try t.expectEqual(tree.get(tree.first().?).?, 3);
    try t.expectEqual(tree.get(tree.last().?).?, 0);
    try t.expectEqual(tree.lookupCustom(3, {}, testCompareReverse), third_id);
}

test "Remove root with no children." {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(10);
    try t.expectEqual(tree.root, a);

    try tree.remove(a);
    const node_ids = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(node_ids);
    try t.expectEqualSlices(TestId, node_ids, &.{});
}

test "Remove root with left red child." {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(10);
    const b = try tree.insert(5);
    try t.expectEqual(tree.root, a);
    try t.expectEqual(tree.getNode(a).?.left, b);
    try t.expectEqual(tree.getNode(b).?.color, .Red);

    try tree.remove(a);
    try t.expectEqual(tree.getNode(b).?.color, .Black);

    const node_ids = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(node_ids);
    try t.expectEqualSlices(TestId, node_ids, &.{b});
}

test "Remove root with right red child." {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(10);
    const b = try tree.insert(15);
    try t.expectEqual(tree.root, a);
    try t.expectEqual(tree.getNode(a).?.right, b);
    try t.expectEqual(tree.getNode(b).?.color, .Red);

    try tree.remove(a);
    try t.expectEqual(tree.getNode(b).?.color, .Black);

    const node_ids = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(node_ids);
    try t.expectEqualSlices(TestId, node_ids, &.{b});
}

test "Remove root with two red children." {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(10);
    const b = try tree.insert(5);
    const c = try tree.insert(15);
    try t.expectEqual(tree.root, a);
    try t.expectEqual(tree.getNode(a).?.left, b);
    try t.expectEqual(tree.getNode(a).?.right, c);
    try t.expectEqual(tree.getNode(b).?.color, .Red);
    try t.expectEqual(tree.getNode(c).?.color, .Red);

    try tree.remove(a);
    try t.expectEqual(tree.root, c);
    try t.expectEqual(tree.getNode(c).?.color, .Black);
    try t.expectEqual(tree.getNode(c).?.left, b);
    try t.expectEqual(tree.getNode(b).?.color, .Red);

    const node_ids = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(node_ids);
    try t.expectEqualSlices(TestId, node_ids, &.{ b, c });
}

test "Remove red non-root." {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(10);
    const b = try tree.insert(5);
    try t.expectEqual(tree.root, a);
    try t.expectEqual(tree.getNode(a).?.left, b);
    try t.expectEqual(tree.getNode(b).?.color, .Red);

    try tree.remove(b);

    const node_ids = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(node_ids);
    try t.expectEqualSlices(TestId, node_ids, &.{a});
}

test "Remove black non-root with left red child." {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(3);
    const b = try tree.insert(2);
    const c = try tree.insert(4);
    const d = try tree.insert(1);
    try t.expectEqual(tree.root, a);
    try t.expectEqual(tree.getNode(b).?.color, .Black);
    try t.expectEqual(tree.getNode(d).?.color, .Red);

    try tree.remove(b);

    try t.expectEqual(tree.getNode(d).?.parent, a);
    try t.expectEqual(tree.getNode(d).?.color, .Black);

    const node_ids = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(node_ids);
    try t.expectEqualSlices(TestId, node_ids, &.{ d, a, c });
}

test "Remove black non-root with right red child." {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(3);
    const b = try tree.insert(2);
    const c = try tree.insert(4);
    const d = try tree.insert(5);
    try t.expectEqual(tree.root, a);
    try t.expectEqual(tree.getNode(c).?.color, .Black);
    try t.expectEqual(tree.getNode(d).?.color, .Red);

    try tree.remove(c);

    try t.expectEqual(tree.getNode(d).?.parent, a);
    try t.expectEqual(tree.getNode(d).?.color, .Black);

    const node_ids = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(node_ids);
    try t.expectEqualSlices(TestId, node_ids, &.{ b, a, d });
}

test "Remove non-root with double black case: Parent is red, right sibling is black, and sibling's children are black." {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(15);
    const b = try tree.insert(10);
    const c = try tree.insert(20);
    const d = try tree.insert(17);
    const e = try tree.insert(23);
    const f = try tree.insert(25);
    try tree.remove(f);
    try t.expectEqual(tree.getNode(c).?.color, .Red);
    try t.expectEqual(tree.getNode(d).?.color, .Black);
    try t.expectEqual(tree.getNode(e).?.color, .Black);
    try t.expectEqual(tree.getNode(e).?.left, TestNullId);
    try t.expectEqual(tree.getNode(e).?.right, TestNullId);

    try tree.remove(d);
    try t.expectEqual(tree.getNode(c).?.color, .Black);
    try t.expectEqual(tree.getNode(c).?.left, TestNullId);
    try t.expectEqual(tree.getNode(e).?.color, .Red);

    const node_ids = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(node_ids);
    try t.expectEqualSlices(TestId, node_ids, &.{ b, a, c, e });
}

test "Remove non-root with double black case: Parent is red, left sibling is black, and sibling's children are black." {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(15);
    const b = try tree.insert(10);
    const c = try tree.insert(20);
    const d = try tree.insert(13);
    const e = try tree.insert(7);
    const f = try tree.insert(5);
    try tree.remove(f);
    try t.expectEqual(tree.getNode(b).?.color, .Red);
    try t.expectEqual(tree.getNode(d).?.color, .Black);
    try t.expectEqual(tree.getNode(e).?.color, .Black);
    try t.expectEqual(tree.getNode(e).?.left, TestNullId);
    try t.expectEqual(tree.getNode(e).?.right, TestNullId);

    try tree.remove(d);
    try t.expectEqual(tree.getNode(b).?.color, .Black);
    try t.expectEqual(tree.getNode(b).?.right, TestNullId);
    try t.expectEqual(tree.getNode(e).?.color, .Red);

    const node_ids = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(node_ids);
    try t.expectEqualSlices(TestId, node_ids, &.{ e, b, a, c });
}

test "Remove non-root with double black case: Right sibling is black, and sibling's right child is red." {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(10);
    const b = try tree.insert(5);
    const c = try tree.insert(15);
    const d = try tree.insert(20);
    try t.expectEqual(tree.getNode(b).?.color, .Black);
    try t.expectEqual(tree.getNode(c).?.color, .Black);
    try t.expectEqual(tree.getNode(c).?.right, d);
    try t.expectEqual(tree.getNode(d).?.color, .Red);

    try tree.remove(b);
    try t.expectEqual(tree.root, c);
    try t.expectEqual(tree.getNode(a).?.color, .Black);
    try t.expectEqual(tree.getNode(c).?.color, .Black);
    try t.expectEqual(tree.getNode(d).?.color, .Black);

    const node_ids = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(node_ids);
    try t.expectEqualSlices(TestId, node_ids, &.{ a, c, d });
}

test "Remove non-root with double black case: Left sibling is black, and sibling's left child is red." {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(10);
    const b = try tree.insert(5);
    const c = try tree.insert(15);
    const d = try tree.insert(4);
    try t.expectEqual(tree.getNode(c).?.color, .Black);
    try t.expectEqual(tree.getNode(b).?.color, .Black);
    try t.expectEqual(tree.getNode(b).?.left, d);
    try t.expectEqual(tree.getNode(d).?.color, .Red);

    try tree.remove(c);
    try t.expectEqual(tree.root, b);
    try t.expectEqual(tree.getNode(a).?.color, .Black);
    try t.expectEqual(tree.getNode(b).?.color, .Black);
    try t.expectEqual(tree.getNode(d).?.color, .Black);

    const node_ids = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(node_ids);
    try t.expectEqualSlices(TestId, node_ids, &.{ d, b, a });
}

test "Remove non-root case 5: parent is red, right sibling is black, sibling's left child is red, sibling's right child is black." {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(10);
    const b = try tree.insertDetached(10, .Red);
    const c = try tree.insertDetached(10, .Black);
    const d = try tree.insertDetached(10, .Black);
    const e = try tree.insertDetached(10, .Black);
    const f = try tree.insertDetached(10, .Red);
    try tree.attachLeft(a, b);
    try tree.attachRight(a, c);
    try tree.attachLeft(b, d);
    try tree.attachRight(b, e);
    try tree.attachLeft(e, f);
    try t.expectEqual(tree.isValid(), true);

    t.log_level = .debug;

    try tree.remove(d);
    try t.expectEqual(tree.isValid(), true);
    const node_ids = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(node_ids);
    try t.expectEqualSlices(TestId, node_ids, &.{ b, f, e, a, c });
}

test "Remove non-root with double black case: Parent is black, right sibling is red, sibling's children are black" {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(10);
    const b = try tree.insert(5);
    const c = try tree.insert(15);
    const d = try tree.insert(20);
    const e = try tree.insert(25);
    const f = try tree.insert(30);
    try t.expectEqual(tree.getNode(b).?.color, .Black);
    try t.expectEqual(tree.getNode(a).?.color, .Black);
    try t.expectEqual(tree.getNode(d).?.color, .Red);
    try t.expectEqual(tree.getNode(d).?.left, c);
    try t.expectEqual(tree.getNode(d).?.right, e);
    try t.expectEqual(tree.getNode(c).?.color, .Black);
    try t.expectEqual(tree.getNode(e).?.color, .Black);

    try tree.remove(b);
    try t.expectEqual(tree.root, d);
    try t.expectEqual(tree.getNode(c).?.color, .Red);
    try t.expectEqual(tree.getNode(a).?.color, .Black);
    try t.expectEqual(tree.getNode(d).?.color, .Black);
    try t.expectEqual(tree.getNode(e).?.color, .Black);

    const node_ids = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(node_ids);
    try t.expectEqualSlices(TestId, node_ids, &.{ a, c, d, e, f });
}

test "Remove non-root with double black case: Parent is black, left sibling is red, sibling's children are black" {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(10);
    const b = try tree.insert(5);
    const c = try tree.insert(15);
    const d = try tree.insert(4);
    const e = try tree.insert(3);
    const f = try tree.insert(2);
    try t.expectEqual(tree.getNode(b).?.color, .Black);
    try t.expectEqual(tree.getNode(a).?.color, .Black);
    try t.expectEqual(tree.getNode(d).?.color, .Red);
    try t.expectEqual(tree.getNode(d).?.left, e);
    try t.expectEqual(tree.getNode(d).?.right, b);
    try t.expectEqual(tree.getNode(e).?.color, .Black);
    try t.expectEqual(tree.getNode(b).?.color, .Black);

    try tree.remove(c);
    try t.expectEqual(tree.root, d);
    try t.expectEqual(tree.getNode(b).?.color, .Red);
    try t.expectEqual(tree.getNode(a).?.color, .Black);
    try t.expectEqual(tree.getNode(d).?.color, .Black);
    try t.expectEqual(tree.getNode(e).?.color, .Black);

    const node_ids = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(node_ids);
    try t.expectEqualSlices(TestId, node_ids, &.{ f, e, d, b, a });
}

test "Remove non-root with double black case: Parent is black, right sibling is black, sibling's left child is red, sibling's right child is black." {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(10);
    const b = try tree.insert(5);
    const c = try tree.insert(15);
    const d = try tree.insert(13);
    try t.expectEqual(tree.getNode(b).?.color, .Black);
    try t.expectEqual(tree.getNode(c).?.color, .Black);
    try t.expectEqual(tree.getNode(c).?.left, d);
    try t.expectEqual(tree.getNode(d).?.color, .Red);

    try tree.remove(b);
    try t.expectEqual(tree.root, d);
    try t.expectEqual(tree.getNode(c).?.color, .Black);
    try t.expectEqual(tree.getNode(d).?.color, .Black);

    const node_ids = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(node_ids);
    try t.expectEqualSlices(TestId, node_ids, &.{ a, d, c });
}

test "Remove non-root with double black case: Parent is black, left sibling is black, sibling's right child is red, sibling's left child is black." {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(10);
    const b = try tree.insert(5);
    const c = try tree.insert(15);
    const d = try tree.insert(7);
    try t.expectEqual(tree.getNode(c).?.color, .Black);
    try t.expectEqual(tree.getNode(b).?.color, .Black);
    try t.expectEqual(tree.getNode(b).?.right, d);
    try t.expectEqual(tree.getNode(d).?.color, .Red);

    try tree.remove(c);
    try t.expectEqual(tree.root, d);
    try t.expectEqual(tree.getNode(b).?.color, .Black);
    try t.expectEqual(tree.getNode(d).?.color, .Black);

    const node_ids = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(node_ids);
    try t.expectEqualSlices(TestId, node_ids, &.{ b, d, a });
}

test "Remove non-root with double black case: Parent is black, right sibling is black, and sibling's children are black." {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(15);
    const b = try tree.insert(10);
    const c = try tree.insert(20);
    const d = try tree.insert(5);
    try tree.remove(d);
    try t.expectEqual(tree.getNode(a).?.color, .Black);
    try t.expectEqual(tree.getNode(b).?.color, .Black);
    try t.expectEqual(tree.getNode(c).?.color, .Black);

    try tree.remove(b);
    try t.expectEqual(tree.root, a);
    try t.expectEqual(tree.getNode(a).?.color, .Black);
    try t.expectEqual(tree.getNode(c).?.color, .Red);

    const node_ids = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(node_ids);
    try t.expectEqualSlices(TestId, node_ids, &.{ a, c });
}

test "Remove non-root with double black case: Parent is black, left sibling is black, and sibling's children are black." {
    var tree = initTestTree();
    defer tree.deinit();

    const a = try tree.insert(15);
    const b = try tree.insert(10);
    const c = try tree.insert(20);
    const d = try tree.insert(5);
    try tree.remove(d);
    try t.expectEqual(tree.getNode(a).?.color, .Black);
    try t.expectEqual(tree.getNode(b).?.color, .Black);
    try t.expectEqual(tree.getNode(c).?.color, .Black);

    try tree.remove(c);
    try t.expectEqual(tree.root, a);
    try t.expectEqual(tree.getNode(a).?.color, .Black);
    try t.expectEqual(tree.getNode(b).?.color, .Red);

    const node_ids = tree.allocNodeIdsInOrder(t.allocator);
    defer t.allocator.free(node_ids);
    try t.expectEqualSlices(TestId, node_ids, &.{ b, a });
}

fn initTestTree() RbTree(TestId, u32, void, testCompare) {
    return RbTree(TestId, u32, void, testCompare).init(t.allocator, {});
}

const TestId = u32;
const TestNullId = @import("compact.zig").CompactNull(TestId);