import sys
from random import choice
from PyQt5.QtWidgets import QApplication, QWidget, QPushButton, QGridLayout, QLabel
from PyQt5.QtCore import QObject

class DifficultySelector(QObject):
    def __init__(self):
        super().__init__()
        self.initUI()

    def initUI(self):
        self.app = QApplication(sys.argv)
        self.window = QWidget()
        self.window.setWindowTitle("Game Difficulty Selector")

        # Set button styles for a more modern look
        style = "QPushButton { font-size: 16px; border: 1px solid #555; border-radius: 10px; padding: 5px 10px; }" \
                "QPushButton:checked { background-color: #005f73; color: #fff; }"
        self.window.setStyleSheet(style)

        # Initialize grid layout
        self.grid_layout = QGridLayout()
        self.window.setLayout(self.grid_layout)

        # Set spacing between different sections
        self.grid_layout.setRowStretch(1, 10)
        self.grid_layout.setRowStretch(5, 10)

        # Difficulty selection
        label = QLabel("Choose a difficulty level:")
        self.grid_layout.addWidget(label, 0, 0, 1, 3)

        self.difficulties = [QPushButton("Easy"), QPushButton("Medium"), QPushButton("Hard")]
        for i, button in enumerate(self.difficulties):
            button.setCheckable(True)
            button.clicked.connect(self.difficulty_clicked)
            button.setAccessibleName(f"MCTS {button.text()}")  # Set the algorithm
            self.grid_layout.addWidget(button, 2, i)  # Place in a row with spacing

        # Color selection
        label = QLabel("Choose a color:")
        self.grid_layout.addWidget(label, 4, 0, 1, 3)

        self.colors = [QPushButton("White"), QPushButton("Black")]
        for i, button in enumerate(self.colors):
            button.setCheckable(True)
            button.clicked.connect(self.color_clicked)
            self.grid_layout.addWidget(button, 6, i)

        # Start game button spanning all columns
        start_game = QPushButton("Start Game")
        start_game.clicked.connect(self.get_difficulty_and_close)
        self.grid_layout.addWidget(start_game, 7, 0, 1, 3)

        self.window.show()
        self.app.exec_()

    def difficulty_clicked(self, checked):
        if checked:
            sender = self.sender()
            for btn in self.difficulties:
                if btn != sender:
                    btn.setChecked(False)

    def color_clicked(self, checked):
        if checked:
            sender = self.sender()
            for btn in self.colors:
                if btn != sender:
                    btn.setChecked(False)

    def get_difficulty_and_close(self):
        difficulty, color = self.get_difficulty()
        print(difficulty, color)
        self.window.close()

    def get_difficulty(self):
        difficulty, color = None, None
        for btn in self.difficulties:
            if btn.isChecked():
                difficulty = btn.accessibleName()
        for btn in self.colors:
            if btn.isChecked():
                color = btn.text()

        if difficulty is None:
            difficulty = f'{choice(["Minimax", "MCTS"])} {choice(["Easy", "Medium", "Hard"])}'
        if color == "Random" or color is None:
            color = choice(["White", "Black"])

        return difficulty, color

if __name__ == '__main__':
    dif = DifficultySelector()
