# Imports
import pygame, random
from button import Button
from game_functions import animate_X, animate_0, draw_board, check_winner, ai_move

# Setup
pygame.init()
screen = pygame.display.set_mode((300, 300))
pygame.display.set_caption("TicTacToe")
clock = pygame.time.Clock()
pygame.font.init()
font = pygame.font.Font("./assets/font.ttf", 70)

# Buttons
button_human = Button(pygame.image.load("./assets/vshuman.png"), 50, 135)
button_ai = Button(pygame.image.load("./assets/vsai.png"), 50, 185)
button_menu = Button(pygame.image.load("./assets/menu.png"), 50, 135)
button_restart = Button(pygame.image.load("./assets/restart.png"), 50, 185)

# Title Screen
def screen_title():
    # Game Loop
    counter = 0
    run = True
    while run:
        # Tick
        clock.tick(30)
        # Event Handler
        for event in pygame.event.get():
            # Quit
            if event.type == pygame.QUIT:
                pygame.quit()
                exit()
        counter += 1
        if counter == 120:
            run = False
        # Display Title Screen
        if run:
            screen.fill((255, 255, 200))
            text = font.render("Tic", 1, (255, 100, 100))
            screen.blit(text, (25, 15))
            text = font.render("Tac", 1, (0, 0, 0))
            screen.blit(text, (95, 15))
            text = font.render("Toe", 1, (100, 100, 255))
            screen.blit(text, (180, 15))
            text = font.render("by", 1, (0, 0, 0))
            screen.blit(text, (120, 107))
            text = font.render("Tylan", 1, (255, 100, 100))
            screen.blit(text, (5, 210))
            text = font.render("Tyson", 1, (100, 100, 255))
            screen.blit(text, (140, 210))
            # Update Display
            pygame.display.update()

# Menu Screen
def screen_menu():
    # Game Loop
    option = None
    run = True
    while run:
        # Tick
        clock.tick(30)
        # Event Handler
        for event in pygame.event.get():
            # Quit
            if event.type == pygame.QUIT:
                pygame.quit()
                exit()
        # Display Title
        screen.fill((255, 255, 200))
        text = font.render("Tic", 1, (255, 100, 100))
        screen.blit(text, (25, 15))
        text = font.render("Tac", 1, (0, 0, 0))
        screen.blit(text, (95, 15))
        text = font.render("Toe", 1, (100, 100, 255))
        screen.blit(text, (180, 15))
        # Buttons
        if button_human.draw(screen):
            option = "human"
            run = False
        if button_ai.draw(screen):
            option = "ai"
            run = False
        # Update Display
        pygame.display.update()
    # Return Option
    return option

# Winner Screen
def screen_winner(winner, restart_option):
    # Game Loop
    option = None
    run = True
    while run:
        # Tick
        clock.tick(30)
        # Event Handler
        for event in pygame.event.get():
            # Quit
            if event.type == pygame.QUIT:
                pygame.quit()
                exit()
        # Display Title
        screen.fill((255, 255, 200))
        if winner == 0:
            text = font.render("Draw!", 1, (0, 0, 0))
            screen.blit(text, (75, 15))
        if winner == 1:
            text = font.render("X Won!", 1, (255, 100, 100))
            screen.blit(text, (60, 15))
        if winner == -1:
            text = font.render("0 Won!", 1, (100, 100, 255))
            screen.blit(text, (60, 15))
        # Buttons
        if button_menu.draw(screen):
            option = None
            run = False
        if button_restart.draw(screen):
            option = restart_option
            run = False
        # Update Display
        pygame.display.update()
    # Return Option
    return option

# Title Screen
screen_title()

# Game Loop
option = None
human_first_run = True
ai_first_run = True
run = True
while run:
    # Tick
    clock.tick(30)
    # Menu
    if option == None:
        option = screen_menu()
    # VS Human
    if option == "human":
        if human_first_run:
            start = random.randint(1, 2)
            if start == 1:
                player = 1
            else:
                player = -1
            mouse_position = None
            clicked = False
            human_first_run = False
            board = [
                [0, 0, 0],
                [0, 0, 0],
                [0, 0, 0]
            ]
        draw_board(screen, board)
        for event in pygame.event.get():
            if event.type == pygame.MOUSEBUTTONDOWN and clicked == False:
                clicked = True
            if event.type == pygame.MOUSEBUTTONUP and clicked == True:
                clicked = False
                mouse_position = pygame.mouse.get_pos()
                cell_x = mouse_position[0]
                cell_y = mouse_position[1]
                if board[cell_y // 100][cell_x // 100] == 0:
                    board[cell_y // 100][cell_x // 100] = player
                    if player == 1:
                        animate_X(screen, cell_x // 100, cell_y // 100)
                    elif player == -1:
                        animate_0(screen, cell_x // 100, cell_y // 100)
                    player *= - 1
            if event.type == pygame.QUIT:
                run = False
        winner = check_winner(board)
        if winner != None:
            human_first_run = True
            option = screen_winner(winner, "human")
        pygame.display.update()
    # VS AI
    if option == "ai":
        if ai_first_run:
            start = random.randint(1, 2)
            player = 1
            mouse_position = None
            clicked = False
            ai_first_run = False
            board = [
                [0, 0, 0],
                [0, 0, 0],
                [0, 0, 0]
            ]
            if start == 1:
                draw_board(screen, board)
                pygame.display.update()
                row, column = ai_move(board)
                board[row][column] = player
                animate_X(screen, column, row)
                player *= -1
                start = 0
        draw_board(screen, board)
        for event in pygame.event.get():
            if event.type == pygame.MOUSEBUTTONDOWN and clicked == False:
                clicked = True
            if event.type == pygame.MOUSEBUTTONUP and clicked == True:
                clicked = False
                mouse_position = pygame.mouse.get_pos()
                mouse_position = pygame.mouse.get_pos()
                cell_x = mouse_position[0]
                cell_y = mouse_position[1]
                if board[cell_y // 100][cell_x // 100] == 0:
                    board[cell_y // 100][cell_x // 100] = player
                    if player == 1:
                        animate_X(screen, cell_x // 100, cell_y // 100)
                    elif player == -1:
                        animate_0(screen, cell_x // 100, cell_y // 100)
                    player *= - 1
                    winner = check_winner(board)
                    if winner == None:
                        row, column = ai_move(board)
                        board[row][column] = player
                        if player == 1:
                            animate_X(screen, column, row)
                        elif player == -1:
                            animate_0(screen, column, row)
                        player *= -1
            if event.type == pygame.QUIT:
                run = False
        winner = check_winner(board)
        if winner != None:
            ai_first_run = True
            option = screen_winner(winner, "ai")
        pygame.display.update()

# Quit Pygame
pygame.quit()