# Implement left rotates and right rotates for trees.
#
# Example right rotate on 10:
#
#       10
#      /  \
#     5   15
#    / \
#   3   7
#  /
# 2
#
# Turns into:
#
#        5
#      /  \
#     3   10
#    /   /  \
#   2   7   15
#
# Notice how the right child of 10's left child (7) becomes the left child of 10.
#
# Left rotates are the reverse. See this example on 16:
#     16
#       \
#        33
#       /  \
#     22    55
#
# Turns into:
#
#      33
#     /  \
#    16   55
#     \
#      22
#
# where the left child of 16's right child (22) turns into 16's right child.
#
# Copy this entire file, implement rotate_left and rotate_right,
# and run the tests using `python3 <name-of-file>.py`.


class Node:
    def __init__(self, value, left=None, right=None, parent=None):
        self.value = value
        self.left: Node | None = left
        self.right: Node | None = right

        # required to re-set the child later
        self.parent: Node | None = parent


def rotate_left(node: Node) -> Node:
    if node.right is None:
        return node

    original_parent = node.parent
    new_parent = node.right
    new_right_child = node.right.left

    node.right = new_right_child
    if new_right_child is not None:
        new_right_child.parent = node

    new_parent.left = node
    node.parent = new_parent
    new_parent.parent = original_parent

    if original_parent is not None:
        if original_parent.left == node:
            original_parent.left = new_parent
        elif original_parent.right == node:
            original_parent.right = new_parent

    return new_parent


def rotate_right(node: Node) -> Node:
    if node.left is None:
        return node

    original_parent = node.parent
    new_parent = node.left
    new_left_child = node.left.right

    node.left = new_left_child
    if new_left_child is not None:
        new_left_child.parent = node

    new_parent.right = node
    node.parent = new_parent
    new_parent.parent = original_parent

    if original_parent is not None:
        if original_parent.left == node:
            original_parent.left = new_parent
        elif original_parent.right == node:
            original_parent.right = new_parent

    return new_parent


# --- Tests ---------------------------------------------------------------
# Both rotate functions are expected to return the new root of the rotated
# subtree (the parent would then point its child link at that new root).


def to_tuple(node: Node | None):
    """Turn a tree into nested (value, left, right) tuples for easy comparison."""
    if node is None:
        return None
    return (
        node.value,
        to_tuple(node.left),
        to_tuple(node.right),
        node.parent.value if node.parent is not None else None,
    )


def trees_equal(a: Node | None, b: Node | None) -> bool:
    return to_tuple(a) == to_tuple(b)


if __name__ == "__main__":
    # note that matching is done by value for simplicity here

    def test_rotate_right_example():
        # The worksheet example: right rotate on 10.
        root = Node(
            10,
            left=Node(
                5,
                left=Node(3, left=Node(2, parent=Node(3)), parent=Node(5)),
                right=Node(7, parent=Node(5)),
                parent=Node(10),
            ),
            right=Node(15, parent=Node(10)),
        )
        expected = Node(
            5,
            left=Node(3, left=Node(2, parent=Node(3)), parent=Node(5)),
            right=Node(
                10,
                left=Node(7, parent=Node(10)),
                right=Node(15, parent=Node(10)),
                parent=Node(5),
            ),
        )
        return rotate_right(root), expected


    def test_rotate_left_example():
        # The worksheet example: left rotate on 16.
        root = Node(
            16,
            right=Node(
                33,
                left=Node(22, parent=Node(33)),
                right=Node(55, parent=Node(33)),
                parent=Node(16),
            ),
        )
        expected = Node(
            33,
            left=Node(16, right=Node(22, parent=Node(16)), parent=Node(33)),
            right=Node(55, parent=Node(33)),
        )
        return rotate_left(root), expected


    def test_rotate_right_two_nodes():
        # Smallest case: a root with only a left child.
        return rotate_right(Node(2, left=Node(1, parent=Node(2)))), Node(1, right=Node(2, parent=Node(1)))


    def test_rotate_left_two_nodes():
        # Smallest case: a root with only a right child.
        return rotate_left(Node(1, right=Node(2, parent=Node(1)))), Node(2, left=Node(1, parent=Node(2)))


    def test_round_trip():
        # A left rotate undone by a right rotate restores the original tree.
        original = Node(1, left=Node(0, parent=Node(1)),
                        right=Node(3, left=Node(2, parent=Node(3)), right=Node(4, parent=Node(3)), parent=Node(1)))
        expected = Node(1, left=Node(0, parent=Node(1)),
                        right=Node(3, left=Node(2, parent=Node(3)), right=Node(4, parent=Node(3)), parent=Node(1)))
        return rotate_right(rotate_left(original)), expected


    def test_rotate_with_parent():
        # run it on 33 instead of the root
        root = Node(16, right=Node(33, left=Node(22), right=Node(55)))
        root.right.parent = root
        root.right.left.parent = root.right
        root.right.right.parent = root.right
        original = root.right
        expected = Node(
            55,
            left=Node(33, left=Node(22, parent=Node(33)), parent=Node(55)),
            parent=Node(16),
        )
        return rotate_left(original), expected  # pyright: ignore[reportArgumentType]


    tests = [
        ("rotate_right on worksheet example (10)", test_rotate_right_example),
        ("rotate_left on worksheet example (16)", test_rotate_left_example),
        ("rotate_right on two nodes", test_rotate_right_two_nodes),
        ("rotate_left on two nodes", test_rotate_left_two_nodes),
        ("rotate_left then rotate_right restores the tree", test_round_trip),
        ("rotate_left works when there is a parent", test_rotate_with_parent),
    ]

    passed = 0
    for name, test in tests:
        try:
            result, expected = test()
            if trees_equal(result, expected):
                print(f"PASS: {name}")
                passed += 1
            else:
                print(f"FAIL: {name}")
                print(f"      expected {to_tuple(expected)}")
                print(f"      got      {to_tuple(result)}")
        except Exception as e:  # noqa: BLE001 - surface any bug as a failed test
            print(f"ERROR: {name}: {type(e).__name__}: {e}")

    print(f"\n{passed}/{len(tests)} tests passed")