import pygame
import pygame.draw as draw
from collections import namedtuple
from typing import NamedTuple
import random
import math

Cell = namedtuple("Cell", ["x", "y"])
Size = namedtuple("Size", ["width", "height"])
ColorSet = namedtuple("ColorSet", ["norm", "light", "dark"])


def flip(mat: list[list]) -> list[list]:
    return list(reversed(mat))


def transpose(mat: list[list]) -> list[list]:
    if not mat:
        return []
    out = [list([None] * len(mat)) for _ in range(len(mat[0]))]
    for x in range(len(mat[0])):
        for y in range(len(mat)):
            out[x][y] = mat[y][x]
    return out


def mirror(mat: list[list]) -> list[list]:
    return [list(reversed(row)) for row in mat]


class Tetromino(NamedTuple):
    color: str
    shape: list[list[str]]

    @property
    def width(self):
        return len(self.shape[0])

    @property
    def height(self):
        return len(self.shape)

    def rotated(self, times: int):
        """Return self rotated TIMES times to the right."""
        times %= 4
        if times == 0:
            return self
        if times == 1:
            return Tetromino(self.color, mirror(transpose(self.shape)))
        if times == 2:
            return Tetromino(self.color, mirror(flip(self.shape)))
        if times == 3:
            return Tetromino(self.color, flip(transpose(self.shape)))

    def __repr__(self):
        return "\n".join(["".join(row) for row in self.shape])


class Game:
    BOARD_SIZE = Size(10, 20)  # Cell count
    CELL_SIZE = Size(30, 30)
    CELL_BORDER_SIZE = CELL_SIZE.width // 10
    CONTROLS_WIDTH = 160
    FRAMERATE = 120
    WINDOW_SIZE = Size(
        (BOARD_SIZE.width + 2) * CELL_SIZE.width + CONTROLS_WIDTH,
        (BOARD_SIZE.height + 2) * CELL_SIZE.height,
    )
    PALETTE = {
        "gray": ColorSet("#787878", "#9a9a9a", "#303030"),
        "purple": ColorSet("#9a00cd", "#cd00ff", "#66009a"),
        "yellow": ColorSet("#cdcd00", "#ffff00", "#9a9a00"),
        "red": ColorSet("#cd0000", "#ff0000", "#9a0000"),
        "orange": ColorSet("#cd6600", "#ff8900", "#9a4200"),
        "blue": ColorSet("#0000cd", "#0000ff", "#00009a"),
        "aqua": ColorSet("#00cdcd", "#00ffff", "#009a9a"),
        "green": ColorSet("#00cd00", "#00ff00", "#009a00"),
    }
    TETROMINOS = [
        Tetromino("aqua", [["*", "*", "*", "*"]]),
        Tetromino("yellow", [["*", "*"], ["*", "*"]]),
        Tetromino("purple", [["*", "*", "*"], [" ", "*", " "]]),
        Tetromino("blue", [[" ", "*"], [" ", "*"], ["*", "*"]]),
        Tetromino("orange", [["*", " "], ["*", " "], ["*", "*"]]),
        Tetromino("green", [[" ", "*", "*"], ["*", "*", " "]]),
        Tetromino("red", [["*", "*", " "], [" ", "*", "*"]]),
    ]

    @staticmethod
    def random_tetromino():
        return random.choice(Game.TETROMINOS)

    def __init__(self):
        self.screen = None
        self.clock = None
        self.running = False

    def draw_block(self, x: int, y: int, color):
        """Draw a block at (X, Y). Coordinates are for the top left corner."""
        draw.rect(
            self.screen,
            Game.PALETTE[color].norm,
            (x, y, Game.CELL_SIZE.width, Game.CELL_SIZE.height),
        )
        draw.polygon(
            self.screen,
            Game.PALETTE[color].light,
            [
                (x, y),
                (x + Game.CELL_SIZE.width, y),
                (
                    x + Game.CELL_SIZE.width - Game.CELL_BORDER_SIZE,
                    y + Game.CELL_BORDER_SIZE,
                ),
                (x + Game.CELL_BORDER_SIZE, y + Game.CELL_BORDER_SIZE),
                (
                    x + Game.CELL_BORDER_SIZE,
                    y + Game.CELL_SIZE.width - Game.CELL_BORDER_SIZE,
                ),
                (x, y + Game.CELL_SIZE.height),
                (x, y),
            ],
        )
        draw.polygon(
            self.screen,
            Game.PALETTE[color].dark,
            [
                (x + Game.CELL_SIZE.width, y),
                (
                    x + Game.CELL_SIZE.width - Game.CELL_BORDER_SIZE,
                    y + Game.CELL_BORDER_SIZE,
                ),
                (
                    x + Game.CELL_SIZE.width - Game.CELL_BORDER_SIZE,
                    y + Game.CELL_SIZE.height - Game.CELL_BORDER_SIZE,
                ),
                (
                    x + Game.CELL_BORDER_SIZE,
                    y + Game.CELL_SIZE.width - Game.CELL_BORDER_SIZE,
                ),
                (x, y + Game.CELL_SIZE.height),
                (x + Game.CELL_SIZE.width, y + Game.CELL_SIZE.height),
                (x + Game.CELL_SIZE.width, y),
            ],
        )

    def draw_board_border(self):
        bwidth = (Game.BOARD_SIZE.width + 2) * Game.CELL_SIZE.width
        bheight = Game.WINDOW_SIZE.height
        for x in range(Game.BOARD_SIZE.width + 2):
            self.draw_block(x * Game.CELL_SIZE.width, 0, "gray")
            self.draw_block(
                x * Game.CELL_SIZE.width,
                bheight - Game.CELL_SIZE.height,
                "gray",
            )
        for y in range(1, Game.BOARD_SIZE.height + 1):
            self.draw_block(0, y * game.CELL_SIZE.height, "gray")
            self.draw_block(
                bwidth - Game.CELL_SIZE.width,
                y * game.CELL_SIZE.height,
                "gray",
            )

    def draw_board_content(self):
        for x in range(Game.BOARD_SIZE.width):
            for y in range(Game.BOARD_SIZE.height):
                if self.board[y][x]:
                    self.draw_block(
                        (x + 1) * Game.CELL_SIZE.width,
                        (y + 1) * Game.CELL_SIZE.height,
                        self.board[y][x],
                    )

    def draw_current_block(self):
        start_x, start_y = self.current_block_pos
        for x in range(self.current_block.width):
            for y in range(self.current_block.height):
                if self.current_block.shape[y][x] == "*":
                    screen_x = (start_x + x + 1) * Game.CELL_SIZE.width
                    screen_y = (start_y + y + 1) * Game.CELL_SIZE.height
                    self.draw_block(
                        screen_x, screen_y, self.current_block.color
                    )

    def handle_events(self):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                self.running = False

    def loop(self):
        self.running = True
        while self.running:
            self.handle_events()
            self.screen.fill("black")
            self.draw_board_border()
            self.draw_board_content()
            self.draw_current_block()
            pygame.display.flip()
            self.clock.tick(Game.FRAMERATE)

    def generate_new_block(self):
        self.current_block = Game.random_tetromino()
        # top left corner
        self.current_block_pos = Cell(
            self.BOARD_SIZE.width // 2
            - math.ceil(self.current_block.width / 2),
            0,
        )

    def init(self):
        pygame.init()
        self.screen = pygame.display.set_mode(Game.WINDOW_SIZE)
        pygame.display.set_caption("Tetris")
        self.clock = pygame.time.Clock()

        # Game state
        self.board = [[None] * Game.BOARD_SIZE.width] * Game.BOARD_SIZE.height
        self.generate_new_block()

    def run(self):
        self.init()
        self.loop()


if __name__ == "__main__":
    game = Game()
    game.run()