import pygame
import pygame.draw as draw
from collections import namedtuple, defaultdict
from typing import NamedTuple, Any
import random
import math
from enum import Enum, auto

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


class Direction(int, Enum):
    LEFT = -1
    RIGHT = 1


def make_matrix(w: int, h: int, elt: Any = None) -> list[list]:
    if not w:
        return []
    return [[elt] * w for _ in range(h)]


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


def transpose(mat: list[list]) -> list[list]:
    if not mat:
        return []
    out = make_matrix(len(mat), 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):
        out = []
        for row in self.shape:
            for cell in row:
                out.append("*" if cell else " ")
            out.append("\n")
        return "".join(out)


class Action(NamedTuple):
    key: int
    type: Type
    repeat: bool
    desc: str

    class Type(Enum):
        ROTATE_LEFT = auto()
        ROTATE_RIGHT = auto()
        DROP = auto()
        SWAP_HOLD = auto()
        MOVE_LEFT = auto()
        MOVE_RIGHT = auto()

    @staticmethod
    def make_map(*acts):
        return {act.key: act for act in acts}


class Game:
    MOVE_SPEED = 20  # cells per second
    NORMAL_DROP_SPEED = 3  # cells per second
    FAST_DROP_SPEED = 10  # cells per second
    CLEAR_SPEED = 5  # frames per block
    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", [[1, 1, 1, 1]]),
        Tetromino("yellow", [[1, 1], [1, 1]]),
        Tetromino("purple", [[1, 1, 1], [0, 1, 0]]),
        Tetromino("blue", [[0, 1], [0, 1], [1, 1]]),
        Tetromino("orange", [[1, 0], [1, 0], [1, 1]]),
        Tetromino("green", [[0, 1, 1], [1, 1, 0]]),
        Tetromino("red", [[1, 1, 0], [0, 1, 1]]),
    ]

    ACTION_MAP = Action.make_map(
        Action(
            pygame.K_z,
            Action.Type.ROTATE_LEFT,
            False,
            "Rotate the current piece left.",
        ),
        Action(
            pygame.K_x,
            Action.Type.ROTATE_RIGHT,
            False,
            "Rotate the current piece right.",
        ),
        Action(
            pygame.K_s,
            Action.Type.SWAP_HOLD,
            False,
            "Swap the current piece and the held piece.",
        ),
        Action(
            pygame.K_DOWN, Action.Type.DROP, True, "Drop the current piece."
        ),
        Action(
            pygame.K_LEFT,
            Action.Type.MOVE_LEFT,
            True,
            "Move the current piece left.",
        ),
        Action(
            pygame.K_RIGHT,
            Action.Type.MOVE_RIGHT,
            True,
            "Move the current piece right.",
        ),
    )

    @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_key_event(self, typ, key):
        if key not in Game.ACTION_MAP:
            return
        elif typ == pygame.KEYDOWN:
            act = Game.ACTION_MAP[key]
            self.pending_actions.append(act.type)
            self.key_states[key] = True
        elif typ == pygame.KEYUP:
            self.key_states[key] = False

    def handle_events(self):
        # process already pushed keys
        for key, state in self.key_states.items():
            if (
                state
                and key in self.ACTION_MAP
                and self.ACTION_MAP[key].repeat
            ):
                self.pending_actions.append(self.ACTION_MAP[key].type)
        # now process new events
        for event in pygame.event.get():
            match event.type:
                case pygame.QUIT:
                    self.running = False
                case pygame.KEYDOWN | pygame.KEYUP:
                    self.handle_key_event(event.type, event.key)

    def swap_with_hold(self):
        pass

    def intersects_board(self, x: int, y: int, piece: Tetromino):
        pw, ph = piece.width, piece.height
        if x < 0 or y < 0:
            return True
        elif x + pw > Game.BOARD_SIZE.width or y + ph > Game.BOARD_SIZE.height:
            return True
        shape = piece.shape
        for cy, row in enumerate(shape):
            for cx, cell in enumerate(row):
                if cell and self.board[y + cy][x + cx]:
                    return True
        return False

    def rotate_current_piece(self, times: int):
        new_piece = self.current_block.rotated(times)
        cx, cy = self.current_block_pos
        for dx in [0, -1]:
            for dy in [1, 0, -1]:
                np = Cell(cx + dx, cy + dy)
                if not self.intersects_board(*np, new_piece):
                    self.current_block = new_piece
                    self.current_block_pos = np
                    return

    def move_current_piece(self, dir: Direction):
        dx = dir * Game.MOVE_SPEED / Game.FRAMERATE
        self.subcell_move += dx
        move_cell = int(self.subcell_move)
        self.subcell_move -= move_cell
        if move_cell:
            new_pos = Cell(
                self.current_block_pos.x + move_cell, self.current_block_pos.y
            )
            if not self.intersects_board(*new_pos, self.current_block):
                self.current_block_pos = new_pos

    def process_actions(self):
        if not (
            {Action.Type.MOVE_LEFT, Action.Type.MOVE_RIGHT}
            & set(self.pending_actions)
        ):
            self.subcell_move = 0
        while self.pending_actions:
            match self.pending_actions.pop(0):
                case Action.Type.ROTATE_LEFT:
                    self.rotate_current_piece(Direction.LEFT)
                case Action.Type.ROTATE_RIGHT:
                    self.rotate_current_piece(Direction.RIGHT)
                case Action.Type.SWAP_HOLD:
                    self.swap_with_hold()
                case Action.Type.DROP:
                    self.doing_drop = True
                case Action.Type.MOVE_LEFT:
                    self.move_current_piece(Direction.LEFT)
                case Action.Type.MOVE_RIGHT:
                    self.move_current_piece(Direction.RIGHT)

    def place_piece(self, x: int, y: int, piece: Tetromino):
        for dy, row in enumerate(piece.shape):
            for dx, cell in enumerate(row):
                if cell:
                    self.board[y + dy][x + dx] = piece.color

    def advance_piece(self):
        speed = (
            Game.FAST_DROP_SPEED if self.doing_drop else Game.NORMAL_DROP_SPEED
        )
        self.subcell_drop += speed / Game.FRAMERATE
        move_cell = int(self.subcell_drop)
        self.subcell_drop -= move_cell
        cp = self.current_block_pos
        if not self.intersects_board(
            cp.x, cp.y + move_cell, self.current_block
        ):
            self.current_block_pos = Cell(cp.x, cp.y + move_cell)
        else:
            for dy in range(move_cell, -1, -1):
                if not self.intersects_board(
                    cp.x, cp.y + dy, self.current_block
                ):
                    self.place_piece(cp.x, cp.y + dy, self.current_block)
                    self.current_block = None
                    return
            # TODO Game over
            print("Game over!")

    @staticmethod
    def generate_clear_cells():
        bw = Game.BOARD_SIZE.width
        if bw % 2 == 0:
            return zip(range(bw // 2 - 1, -1, -1), range(bw // 2, bw))
        else:
            return zip(range(bw // 2, -1, -1), range(bw // 2, bw))

    def maybe_clear_rows(self):
        def do_single_clear(y: int):
            row = self.board[y]
            for x1, x2 in Game.generate_clear_cells():
                if row[x1]:
                    row[x1] = None
                    row[x2] = None
                    break

        def drop_rows_above_for_clear(y: int):
            self.board[: y + 1] = self.board[:y]
            self.board.insert(0, [None] * Game.BOARD_SIZE.width)

        if not self.clearing_rows:
            for y, row in enumerate(self.board):
                if all(row):
                    self.clearing_rows.add(y)
        elif not any(self.board[next(iter(self.clearing_rows))]):
            for y in self.clearing_rows:
                drop_rows_above_for_clear(y)
            self.clearing_rows = set()
            self.clearing_frames = self.CLEAR_SPEED
        else:
            if self.clearing_frames == self.CLEAR_SPEED:
                self.clearing_frames = 0
                for y in self.clearing_rows:
                    do_single_clear(y)
            else:
                self.clearing_frames += 1

    def loop(self):
        self.running = True
        while self.running:
            self.doing_drop = False

            self.handle_events()

            self.maybe_clear_rows()
            if not self.clearing_rows:
                self.process_actions()
                self.advance_piece()
            else:
                # ignore input while clearing
                self.pending_actions = []

            if not self.current_block:
                self.generate_new_block()

            self.screen.fill("black")
            if not self.clearing_rows:
                self.draw_current_block()
            self.draw_board_content()
            self.draw_board_border()

            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.SCALED)
        pygame.display.set_caption("Tetris")
        self.clock = pygame.time.Clock()

        pygame.key.stop_text_input()
        self.key_states = defaultdict(lambda: False)
        self.pending_actions = []

        # Game state
        self.board = make_matrix(*Game.BOARD_SIZE)
        self.generate_new_block()
        self.held_piece = None
        self.subcell_move = 0
        self.subcell_drop = 0
        self.clearing_rows = set()
        # immediately clear the first block
        self.clearing_frames = Game.CLEAR_SPEED

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


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