#!/usr/bin/env python3

import math
import sys
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 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=lambda: print(t)), list(range(btns))))
        buttons.append(ButtonDef(text="Exit", cbk=self.quit))
        self.board: IBoard = MenuBoard(pygame.Rect(0, 0, info.current_w, info.current_h), buttons)

    def loop(self):
        self.running = True
        while self.running:
            # poll for events
            # pygame.QUIT event means the user clicked X to close your window
            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.running = False


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