can you take a look at this code for me just make sure it has no bugs that I overlooked

class AVLNode:
    def __init__(self, key):
        self.key = key
        self.count = 1      # To track how many times this key has been inserted
        self.height = 1     # Height of node
        self.size = 1       # Size of subtree rooted at this node (including duplicates)
        self.left = None
        self.right = None

class AVLTree:
    def _height(self, node):
        return node.height if node else 0

    def _size(self, node):
        return node.size if node else 0

    def _update(self, node):
        """ Update the height and size of a node. """
        if node:
            node.height = 1 + max(self._height(node.left), self._height(node.right))
            node.size = node.count + self._size(node.left) + self._size(node.right)

    def _balance_factor(self, node):
        return self._height(node.left) - self._height(node.right)

    def _rotate_right(self, y):
        """ Perform right rotation. """
        x = y.left
        T2 = x.right
        
        x.right = y
        y.left = T2

        # Update heights and sizes
        self._update(y)
        self._update(x)

        return x

    def _rotate_left(self, x):
        """ Perform left rotation. """
        y = x.right
        T2 = y.left

        y.left = x
        x.right = T2

        # Update heights and sizes
        self._update(x)
        self._update(y)

        return y

    def _rebalance(self, node):
        """ Rebalance the node if unbalanced. """
        self._update(node)

        # Left heavy
        if self._balance_factor(node) > 1:
            if self._balance_factor(node.left) < 0:
                node.left = self._rotate_left(node.left)
            return self._rotate_right(node)

        # Right heavy
        if self._balance_factor(node) < -1:
            if self._balance_factor(node.right) > 0:
                node.right = self._rotate_right(node.right)
            return self._rotate_left(node)

        return node

    def insert(self, node, key):
        """ Insert a key into the AVL tree rooted at node. """
        if not node:
            return AVLNode(key)

        if key == node.key:
            node.count += 1  # Increment count of the key (handle duplicate)
        elif key < node.key:
            node.left = self.insert(node.left, key)
        else:
            node.right = self.insert(node.right, key)

        return self._rebalance(node)

    def _min_value_node(self, node):
        """ Return the node with the minimum key value found in the tree. """
        current = node
        while current.left is not None:
            current = current.left
        return current

    def delete(self, node, key):
        """ Delete a key from the AVL tree rooted at node. """
        if not node:
            return node

        if key < node.key:
            node.left = self.delete(node.left, key)
        elif key > node.key:
            node.right = self.delete(node.right, key)
        else:
            if node.count > 1:
                node.count -= 1  # Decrease the count instead of deleting the node
            else:
                if node.left is None:
                    return node.right
                elif node.right is None:
                    return node.left

                temp = self._min_value_node(node.right)
                node.key = temp.key
                node.count = temp.count
                node.right = self.delete(node.right, temp.key)

        return self._rebalance(node)

    def count_less_than(self, node, key):
        """ Count the number of elements in the tree rooted at node that are less than key. """
        if not node:
            return 0

        if key <= node.key:
            return self.count_less_than(node.left, key)
        else:
            return node.count + self._size(node.left) + self.count_less_than(node.right, key)

class AVLSet:
    """ AVLTree wrapper that provides a set-like interface for inserting, removing, and counting. """
    def __init__(self):
        self.root = None
        self.tree = AVLTree()

    def insert(self, key):
        self.root = self.tree.insert(self.root, key)

    def remove(self, key):
        self.root = self.tree.delete(self.root, key)

    def count_less_than(self, key):
        return self.tree.count_less_than(self.root, key)