#!/usr/bin/env python3

import math
import sys
from time import sleep
from threading import Thread
from queue import Queue, Empty
from abc import abstractmethod
from typing import Tuple, Callable, TypedDict, List

import pygame
import pygame.freetype


class ButtonDef(TypedDict):
    text: str
    cbk: Callable[[], None]


class IBoard:
    @abstractmethod
    def draw(self, screen: pygame.Surface):
        raise NotImplementedError()

    @abstractmethod
    def handle_event(self, event: pygame.event.Event):
        raise NotImplementedError()


class Button:
    def __init__(self, position: pygame.Rect, text: str, cbk: Callable[[], None]):
        self.position: pygame.Rect
        self.set_position(position)
        self.text = text
        self.cbk = cbk
        self.font = pygame.freetype.SysFont(name="Sans", size=20)
        text_rect = self.font.get_rect(self.text)
        self.text_pos = pygame.Rect((self.position.left + (self.position.width - text_rect.width)/2),
                                    (self.position.top + (self.position.height - text_rect.height)/2),
                                    text_rect.width, text_rect.height)

    def set_position(self, position: pygame.Rect):
        self.position = position

    def draw(self, screen: pygame.Surface):
        pygame.draw.rect(screen, "black", self.position, 0)
        self.font.render_to(screen, self.text_pos, self.text, fgcolor="pink")

    def pos_is_inside(self, pos: Tuple) -> bool:
        return (self.position.left < pos[0] and (self.position.left + self.position.width) > pos[0]) and \
            (self.position.top < pos[1] and (self.position.top + self.position.height) > pos[1])

    def handle_event(self, event: pygame.event.Event):
        if event.type == pygame.MOUSEBUTTONUP:
            if event.button == 1 and self.pos_is_inside(event.pos):
                self.cbk()


class MessageBoard(IBoard):
    def __init__(self, screen_rect: pygame.Rect, text: str):
        self.text = text
        self.font = pygame.freetype.SysFont(name="Sans", size=30)
        text_rect = self.font.get_rect(self.text)
        self.text_pos = pygame.Rect((screen_rect.left + (screen_rect.width - text_rect.width)/2),
                                    (screen_rect.top + (screen_rect.height - text_rect.height)/2),
                                    text_rect.width, text_rect.height)

    def handle_event(self, _):
        pass

    def draw(self, screen: pygame.Surface):
        screen.fill("purple")
        self.font.render_to(screen, self.text_pos, self.text, fgcolor="black")


class MenuBoard(IBoard):
    BUTTON_MARGINS = 10

    def __init__(self, screen_rect: pygame.Rect, buttons: List[ButtonDef]):
        self.rows: int
        self.cols: int
        if len(buttons) == 2:
            # special case which doesn't seem to fit into the sqrt
            self.rows = 1
            self.cols = 2
        else:
            self.rows = math.ceil(math.sqrt(len(buttons)))
            self.cols = math.ceil(math.sqrt(len(buttons)))
        button_positions = self.generate_button_positions(screen_rect)
        self.buttons = list(map(lambda d: Button(next(button_positions),
                                                 f"Button {d['text']}",
                                                 d['cbk']), buttons))

    def generate_button_positions(self, screen_rect: pygame.Rect):
        current_button_row = 0
        current_button_col = 0
        button_width = math.floor((screen_rect.width - (self.cols + 1) * self.BUTTON_MARGINS) / self.cols)
        button_height = math.floor((screen_rect.height - (self.rows + 1) * self.BUTTON_MARGINS) / self.rows)
        while (current_button_row * current_button_col) < (self.cols * self.rows):
            top = self.BUTTON_MARGINS + (current_button_row * (button_height + self.BUTTON_MARGINS))
            left = self.BUTTON_MARGINS + (current_button_col * (button_width + self.BUTTON_MARGINS))
            current_button_col += 1
            if current_button_col >= self.cols:
                current_button_col = 0
                current_button_row += 1
            yield pygame.Rect(left, top, button_width, button_height)

    def handle_event(self, event: pygame.event.Event):
        for b in self.buttons:
            b.handle_event(event)

    def draw(self, screen: pygame.Surface):
        screen.fill("purple")
        for b in self.buttons:
            b.draw(screen)


class App:
    def __init__(self, btns: int):
        pygame.init()
        info = pygame.display.Info()
        self.screen = pygame.display.set_mode((info.current_w, info.current_h), flags=pygame.FULLSCREEN)
        self.clock = pygame.time.Clock()
        self.running = False
        buttons = list(map(lambda t: ButtonDef(text=f"{t}", cbk=self.button_press_handler(str(t))), list(range(btns))))
        buttons.append(ButtonDef(text="Exit", cbk=self.quit))
        self.board: IBoard = MenuBoard(self.get_screen_rect(), buttons)
        self.task_q = Queue()

    def get_screen_rect(self) -> pygame.Rect:
        return pygame.Rect(0, 0, self.screen.get_width(), self.screen.get_height())

    def button_press_handler(self, text: str) -> Callable[[], None]:
        def impl():
            previous_board = self.board

            def thr_fun():
                def end_thr():
                    self.board = previous_board

                sleep(5)
                self.task_q.put(end_thr)

            process_thr = Thread(target=thr_fun)
            process_thr.start()
            self.board = MessageBoard(self.get_screen_rect(), f"Wait: {text}")
        return impl

    def loop(self):
        self.running = True
        while self.running:
            try:
                self.task_q.get_nowait()()
            except Empty:
                pass
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    self.running = False
                self.board.handle_event(event)

            self.board.draw(self.screen)

            pygame.display.flip()

            self.clock.tick(30)
        pygame.quit()

    def quit(self):
        self.board = MessageBoard(self.get_screen_rect(), "Exiting...")
        self.running = False


if __name__ == '__main__':
    app = App(int(sys.argv[1]))
    app.loop()