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.

    We use the height convention from the lecture:
    height(None) = -1, hence a leaf has height 0.
    """

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


class AVLTree:
    """AVL tree with insertion and successor.

    The implementation maintains parent pointers and node heights.  Duplicate
    keys are not inserted again; instead, the stored value is updated.
    """

    def __init__(self) -> None:
        self.root: Optional[AVLNode] = None
        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) = -1."""
        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)."""
        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.

                x                         y
               / \                       / \
              A   y        ---->        x   C
                 / \                   / \
             mid   C                  A  mid
        """
        y = x.right
        assert y is not None

        middle_subtree = y.left
        old_parent = x.parent

        y.left = x
        x.parent = y

        x.right = middle_subtree
        if middle_subtree is not None:
            middle_subtree.parent = x

        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

        self._update_height(x)
        self._update_height(y)
        return y

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

                  x                       y
                 / \                     / \
                y   C      ---->        A   x
               / \                         / \
              A  mid                    mid C
        """
        y = x.left
        assert y is not None

        middle_subtree = y.right
        old_parent = x.parent

        y.right = x
        x.parent = y

        x.left = middle_subtree
        if middle_subtree is not None:
            middle_subtree.parent = x

        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

        self._update_height(x)
        self._update_height(y)
        return y

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

        if balance > 1:
            assert node.left is not None
            if self._balance_factor(node.left) < 0:
                self._rotate_left(node.left)
            return self._rotate_right(node)

        if balance < -1:
            assert node.right is not None
            if self._balance_factor(node.right) > 0:
                self._rotate_right(node.right)
            return self._rotate_left(node)

        return node

    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 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.
        """
        if self.root is None:
            self.root = AVLNode(key=key, value=value)
            self._size = 1
            return self.root

        current = self.root
        parent: Optional[AVLNode] = None

        while current is not None:
            parent = current
            if key == current.key:
                current.value = value
                return current
            if key < current.key:
                current = current.left
            else:
                current = current.right

        assert parent is not None
        new_node = AVLNode(key=key, value=value, parent=parent)
        if key < parent.key:
            parent.left = new_node
        else:
            parent.right = new_node
        self._size += 1

        current = parent
        while current is not None:
            new_subtree_root = self._rebalance(current)
            current = new_subtree_root.parent

        return new_node

    def successor(self, node: AVLNode) -> Optional[AVLNode]:
        """Return the node with the smallest key strictly larger than node.key."""
        if node.right is not None:
            return self.minimum(node.right)

        current = node
        parent = current.parent
        while parent is not None and current is parent.right:
            current = parent
            parent = parent.parent
        return parent