from __future__ import annotations

from dataclasses import dataclass
from typing import Any, List, Optional, Tuple


@dataclass
class AVLNode:
    """A single node of an AVL tree.

    The tree stores keys and optional values.  The key determines the position
    of the node in the binary-search-tree order.  The value is arbitrary data
    associated with the key.

    The parent pointer is not strictly necessary for AVL insertion, but it makes
    operations such as successor easier: from a node we can walk upwards without
    starting again at the root.

    The height field stores the height of the subtree rooted at this node.
    We use height(None) = 0, hence a leaf has height 1.
    """

    key: int
    value: Any = None
    left: Optional["AVLNode"] = None
    right: Optional["AVLNode"] = None
    parent: Optional["AVLNode"] = None
    height: int = 0


class AVLTree:
    """AVL tree boilerplate.

    Students are expected to implement only insert and successor.  The helper
    methods below maintain heights, perform rotations, rebalance nodes, and
    provide utilities for testing.
    """

    def __init__(self) -> None:
        # The root is None exactly when the tree is empty.
        self.root: Optional[AVLNode] = None

        # Number of nodes currently stored in the tree.
        self._size = 0

    def size(self) -> int:
        """Return the number of keys stored in the tree."""
        return self._size

    def height(self, node: Optional[AVLNode]) -> int:
        """Return the height of node, using height(None) = 0."""
        if node is None:
            return -1
        return node.height

    def _update_height(self, node: AVLNode) -> None:
        """Recompute node.height from the heights of its children."""
        node.height = 1 + max(self.height(node.left), self.height(node.right))

    def _balance_factor(self, node: AVLNode) -> int:
        """Return height(left subtree) - height(right subtree).

        In an AVL tree this value must always be -1, 0, or 1.  A value greater
        than 1 means the node is left-heavy.  A value smaller than -1 means the
        node is right-heavy.
        """
        return self.height(node.left) - self.height(node.right)

    def _rotate_left(self, x: AVLNode) -> AVLNode:
        r"""Perform a left rotation at node x and return the new subtree root.

        Before the rotation, x must have a right child y.

                x                         y
               / \                       / \
              A   y        ---->        x   C
                 / \                   / \
              beta  C                 A  beta

        The subtree beta is important: it lies between x and y in sorted order,
        so after y moves above x, beta must become the right child of x.
        """
        y = x.right
        assert y is not None

        # beta is y's left subtree before the rotation.  It contains keys that
        # are larger than x.key and smaller than y.key.  Therefore it becomes
        # the right subtree of x after the rotation.
        beta = y.left

        # Save x's old parent so that the rotated subtree can be reconnected to
        # the rest of the tree after y becomes the local root.
        old_parent = x.parent

        # Put x below y.
        y.left = x
        x.parent = y

        # Move beta from y.left to x.right and fix its parent pointer.
        x.right = beta
        if beta is not None:
            beta.parent = x

        # Reconnect y to x's old parent.  If x used to be the global root, y is
        # now the global root.
        y.parent = old_parent
        if old_parent is None:
            self.root = y
        elif old_parent.left is x:
            old_parent.left = y
        else:
            old_parent.right = y

        # Heights must be updated bottom-up: x is now below y, so x first.
        self._update_height(x)
        self._update_height(y)
        return y

    def search(self, key: int) -> Optional[AVLNode]:
        """Return the node with key, or None if key is not present."""
        current = self.root
        while current is not None:
            if key == current.key:
                return current
            if key < current.key:
                current = current.left
            else:
                current = current.right
        return None

    def minimum(self, node: AVLNode) -> AVLNode:
        """Return the node with minimum key in the subtree rooted at node."""
        current = node
        while current.left is not None:
            current = current.left
        return current

    def inorder_keys(self) -> List[int]:
        """Return all keys in sorted order."""
        result: List[int] = []

        def visit(node: Optional[AVLNode]) -> None:
            if node is None:
                return
            visit(node.left)
            result.append(node.key)
            visit(node.right)

        visit(self.root)
        return result

    def check_avl(self) -> bool:
        """Return True iff all stored structure invariants hold.

        This is intended for tests. It checks four conditions:
        1. the binary-search-tree order,
        2. the stored heights,
        3. the AVL balance condition, and
        4. the parent pointers.
        """
        def check(node: Optional[AVLNode], lo: float, hi: float) -> Tuple[bool, int]:
            if node is None:
                return True, -1

            if not (lo < node.key < hi):
                return False, -1

            left_ok, left_h = check(node.left, lo, node.key)
            right_ok, right_h = check(node.right, node.key, hi)

            expected_height = 1 + max(left_h, right_h)
            height_ok = node.height == expected_height
            balance_ok = abs(left_h - right_h) <= 1

            parent_ok = True
            if node.left is not None:
                parent_ok = parent_ok and node.left.parent is node
            if node.right is not None:
                parent_ok = parent_ok and node.right.parent is node

            ok = left_ok and right_ok and height_ok and balance_ok and parent_ok
            return ok, expected_height

        ok, _ = check(self.root, float("-inf"), float("inf"))
        if self.root is not None and self.root.parent is not None:
            return False
        return ok

    def _rebalance(self, node: AVLNode) -> AVLNode:
        """Restore the AVL condition at node and return the new subtree root.

        This method assumes that the subtrees of node are already AVL trees, but
        node itself may have become unbalanced after an insertion below it.
        """
        self._update_height(node)
        bf = self._balance_factor(node)

        # Left-heavy case: node's left subtree is too tall.
        if bf > 1:
            # TODO: Implement this function.
            raise NotImplementedError

        # Right-heavy case: node's right subtree is too tall.
        if bf < -1:
            # TODO: Implement this function.
            raise NotImplementedError

        # Already balanced; only the height may have changed.
        return node

    def _rotate_right(self, x: AVLNode) -> AVLNode:
        r"""Perform a right rotation at node x and return the new subtree root.

        Before the rotation, x must have a left child y.

                  x                       y
                 / \                     / \
                y   C      ---->        A   x
               / \                         / \
              A  beta                   beta C

        The subtree beta is between y and x in sorted order, so after y moves
        above x, beta must become the left child of x.
        """
        # TODO: Implement this function.
        raise NotImplementedError

    def insert(self, key: int, value: Any = None) -> AVLNode:
        """Insert key into the AVL tree and return the corresponding node.

        If key is already present, update its value and return the existing node.
        Otherwise create a new node, insert it as in a binary search tree,
        and rebalance all ancestors on the path back to the root.
        """
        # TODO: Implement this function.
        raise NotImplementedError

    def successor(self, node: AVLNode) -> Optional[AVLNode]:
        """Return the node with the smallest key strictly larger than node.key.

        Return None if node has no successor.
        The function should run in O(h), where h is the height of the tree.
        """
        # TODO: Implement this function.
        raise NotImplementedError