import random
import pygame
from utils.game import Game
from avl_tree import AVLTree

WHITE = (255, 255, 255)
BLUE = (0, 0, 255)
BLACK = (0, 0, 0)
WIDTH = 800
HEIGHT = 400
MARGIN = 20

class AVLTreeGame(Game):

    def __init__(self):
        super().__init__("AVLTree Game", fps=10, size=(WIDTH, HEIGHT))
        random.seed()
        self.z = list(range(1, 501))
        random.shuffle(self.z)
        self.finished = False
        self.tree = AVLTree()
        self.tree.get_height = lambda node: 0 if node is None else 1 + max(self.tree.get_height(node.left), self.tree.get_height(node.right))
        self.height = self.tree.get_height(self.tree.root)
        self.generator = None

    def update_game(self):
        if not self.finished:
            if self.generator is None:
                self.generator = self.tree.insert_stepwise(self.z.pop())
            try:
                next(self.generator)
            except StopIteration:
                self.generator = None
            if self.generator is None and len(self.z) == 0:
                self.finished = True
            self.height = self.tree.get_height(self.tree.root)
        return True

    def draw_game(self):
        self.screen.fill(WHITE)
        if self.height > 0:
            self.draw_tree(self.tree.root, WIDTH // 2, MARGIN, WIDTH // 4 - MARGIN)
        super().draw_game()

    def draw_tree(self, node, x, y, x_offset):
        y_offset = (HEIGHT - (2 * MARGIN)) / self.height
        if node is not None:
            pygame.draw.circle(self.screen, BLUE, (x, y), 2)
            if node.left is not None:
                pygame.draw.line(self.screen, BLACK, (x, y), (x - x_offset, y + y_offset))
                self.draw_tree(node.left, x - x_offset, y + y_offset, x_offset // 2)
            if node.right is not None:
                pygame.draw.line(self.screen, BLACK, (x, y), (x + x_offset, y + y_offset))
                self.draw_tree(node.right, x + x_offset, y + y_offset, x_offset // 2)

if __name__ == "__main__":
    tree_game = AVLTreeGame()
    tree_game.run()