import pygame
import sys
import pygame_menu
from pygame.locals import *

pygame.init()
DISPLAYWINDOWSIZE = pygame.display.set_mode((600, 400))
pygame.display.set_caption("Test")
DISPLAYWINDOWSIZE.fill((255, 255, 255))

difficulty = None  # Global variable to store difficulty

def set_difficulty_easy():
    global difficulty
    difficulty = 'Easy'
    print("Difficulty set to Easy")
    pygame_menu.events.BACK  # Go back to the main menu

def set_difficulty_hard():
    global difficulty
    difficulty = 'Hard'
    print("Difficulty set to Hard")
    pygame_menu.events.BACK  # Go back to the main menu

def start_the_game():
    if difficulty:
        print("Starting the game with difficulty:", difficulty)
    else:
        print("Please select a difficulty first.")

def level_menu():
    mainmenu._open(level)  # Open levels menu

mainmenu_theme = pygame_menu.themes.THEME_SOLARIZED  # Use the same theme for both menus
level_menu_theme = pygame_menu.themes.THEME_SOLARIZED

mainmenu = pygame_menu.Menu('Welcome', 600, 400, theme=mainmenu_theme)
mainmenu.add.text_input('Name: ', default='username', maxchar=20)
mainmenu.add.button('Play', start_the_game)
mainmenu.add.button('Levels', level_menu)
mainmenu.add.button('Quit', pygame_menu.events.EXIT)

level = pygame_menu.Menu('Select a Difficulty', 600, 400, theme=level_menu_theme)
level.add.button('Easy', set_difficulty_easy)
level.add.button('Hard', set_difficulty_hard)
level.add.button('Back', pygame_menu.events.BACK)

# Main loop
while True:
    for event in pygame.event.get():
        if event.type == QUIT:
            pygame.quit()
            sys.exit()

    DISPLAYWINDOWSIZE.fill((255, 255, 255))
    mainmenu.mainloop(DISPLAYWINDOWSIZE)
    pygame.display.update()