#!/usr/bin/env python3

import argparse
import contextlib
import enum
import random
import sys
import time

import pygame

class PS3Controller(enum.Enum):
    CROSS = 0
    CIRCLE = 1
    TRIANGLE = 2
    SQUARE = 3
    DOWN = 14
    LEFT = 15
    RIGHT = 16


class WouldCollide(Exception):
    pass


class Piece():
    def __init__(self):
        self.square = square_template.copy()

    def rotate_clockwise(self):
        self.elements = list(zip(*self.elements[::-1]))

    def rotate_counter_clockwise(self):
        self.rotate_clockwise()
        self.rotate_clockwise()
        self.rotate_clockwise()


class ZPiece(Piece):
    def __init__(self):
        super().__init__()

        self.square.fill(blue)
        self.elements = (self.square, self.square, None), (None, self.square, self.square), (None, None, None)
        self.vertical = False

    def rotate_clockwise(self):
        self.rotate()

    def rotate_counter_clockwise(self):
        self.rotate()

    def rotate(self):
        if self.vertical:
            self.elements = (self.square, self.square, None), (None, self.square, self.square), (None, None, None)
        else:
            self.elements = (None, None, self.square), (None, self.square, self.square), (None, self.square, None)

        self.vertical = not self.vertical


class SPiece(Piece):
    def __init__(self):
        super().__init__()

        self.square.fill(green)
        self.elements = (None, None, None), (None, self.square, self.square), (self.square, self.square, None)
        self.vertical = False

    def rotate_clockwise(self):
        self.rotate()

    def rotate_counter_clockwise(self):
        self.rotate()

    def rotate(self):
        if self.vertical:
            self.elements = (None, self.square, self.square), (self.square, self.square, None), (None, None, None)
        else:
            self.elements = (None, self.square, None), (None, self.square, self.square), (None, None, self.square)

        self.vertical = not self.vertical


class SquarePiece(Piece):
    def __init__(self):
        super().__init__()

        self.square.fill(brown)
        self.elements = ((self.square, self.square), (self.square, self.square))


class IPiece(Piece):
    def __init__(self):
        super().__init__()

        self.square.fill(red)
        self.elements = (None, None, None, None), (None, None, None, None), (self.square, self.square, self.square, self.square), (None, None, None, None)
        self.vertical = False

    def rotate_clockwise(self):
        self.rotate()

    def rotate_counter_clockwise(self):
        self.rotate()

    def rotate(self):
        if self.vertical:
            self.elements = (None, None, None, None), (None, None, None, None), (self.square, self.square, self.square, self.square), (None, None, None, None)
        else:
            self.elements = (None, None, self.square, None), (None, None, self.square, None), (None, None, self.square, None), (None, None, self.square, None)

        self.vertical = not self.vertical


class LPiece(Piece):
    def __init__(self):
        super().__init__()

        self.square.fill(cyan)
        self.elements = (None, None, None), (self.square, self.square, self.square), (None, None, self.square)


class JPiece(Piece):
    def __init__(self):
        super().__init__()

        self.square.fill(purple)
        self.elements = (None, None, None), (self.square, self.square, self.square), (self.square, None, None)


class TPiece(Piece):
    def __init__(self):
        super().__init__()

        self.square.fill(yellow)
        self.elements = (None, None, None), (self.square, self.square, self.square), (None, self.square, None)


def refresh_game_canvas():
    game_canvas.fill(black)

    for row_idx, row in enumerate(grid):
        for col_idx, element in enumerate(row):
            if element is not None:
                game_canvas.blit(element, (col_idx * 50 + 1, row_idx * 50 + 1))

    # Display the current piece
    for row_idx, row in enumerate(current_piece.elements):
        for col_idx, element in enumerate(row):
            if element is not None:
                game_canvas.blit(element, ((col_idx + current_piece_position[1]) * 50 + 1, (row_idx + current_piece_position[0]) * 50 + 1))

def refresh_right_pane_canvas():
    level_canvas.fill(black)
    score_canvas.fill(black)

    if pygame.font:
        score_font = pygame.font.Font(None, 56)

        score_canvas.blit(score_font.render(f"{score:08d}", True, white), (0, 0))
        level_canvas.blit(score_font.render(f"{level:08d}", True, white), (0, 0))
        right_pane_canvas.blit(score_canvas, (2, 10))
        right_pane_canvas.blit(level_canvas, (2, 70))


def refresh_piece_preview_canvas():
    piece_preview_canvas.fill(black)

    non_empty_rows = list()
    for row in next_piece.elements:
        if any(map(lambda element: element is not None, row)):
            non_empty_rows.append(row)

    non_empty_cols = set()
    for row in next_piece.elements:
        for col_id, element in enumerate(row):
            if element is not None:
                non_empty_cols.add(col_id)

    y_offset = (4 - len(non_empty_rows)) / 2
    x_offset = (4 - len(non_empty_cols)) / 2

    # Display the next piece
    for row_idx, row in enumerate(non_empty_rows):
        for col_idx, element in enumerate(row):
            if element is not None:
                piece_preview_canvas.blit(element, ((col_idx + x_offset) * 50 + 1, (row_idx + y_offset) * 50 + 1))

    right_pane_canvas.blit(piece_preview_canvas, (50, 200))

def refresh_screen():
    refresh_game_canvas()
    screen.blit(game_canvas, (0, 0))
    refresh_right_pane_canvas()
    screen.blit(right_pane_canvas, (501, 0))


def has_collision(y: int, x: int) -> bool:
    try:
        for row_id, row in enumerate(current_piece.elements):
            for col_id, element in enumerate(row):
                if element is None:
                    continue

                if row_id + y < 0:
                    continue

                if col_id + x < 0:
                    return True

                if grid[row_id + y][col_id + x] is not None:
                    return True

    except IndexError:
        return True

    return False

def move_piece_down():
    if not has_collision(current_piece_position[0] + 1, current_piece_position[1]):
        current_piece_position[0] += 1
    else:
        raise WouldCollide()

def move_piece_up():
    if not has_collision(current_piece_position[0] - 1, current_piece_position[1]):
        current_piece_position[0] -= 1
    else:
        raise WouldCollide()

def move_piece_left():
    if not has_collision(current_piece_position[0], current_piece_position[1] - 1):
        current_piece_position[1] -= 1
    else:
        raise WouldCollide()

def move_piece_right():
    if not has_collision(current_piece_position[0], current_piece_position[1] + 1):
        current_piece_position[1] += 1
    else:
        raise WouldCollide()

def rotate_piece_counter_clockwise():
    current_piece.rotate_counter_clockwise()

    if has_collision(current_piece_position[0], current_piece_position[1]):
        current_piece.rotate_clockwise()
        raise WouldCollide()

def rotate_piece_clockwise():
    current_piece.rotate_clockwise()

    if has_collision(current_piece_position[0], current_piece_position[1]):
        current_piece.rotate_counter_clockwise()
        raise WouldCollide()


def generate_piece():
    piece = random.choice((TPiece, SPiece, IPiece, ZPiece, SquarePiece, LPiece, JPiece))()

    for row_id, row in enumerate(piece.elements):
        if list(filter(lambda x: x is not None, row)):
            break

    initial_y_position = -row_id
    initial_x_position = (len(grid[0]) // 2) - (len(piece.elements[0]) // 2)

    return (piece, [initial_y_position, initial_x_position])


def burn_rows():
    rows_to_burn = list()

    for row in grid:
        if all(map(lambda element: element is not None, row)):
            rows_to_burn.append(row)

    for row in rows_to_burn:
        grid.insert(0, [None for _ in range(10)])
        grid.remove(row)

    return len(rows_to_burn)


def stick_piece():
    global current_piece, current_piece_position, next_piece, next_piece_position, lines_burnt, level, score

    for row_id, row in enumerate(current_piece.elements):
        for col_id, element in enumerate(row):
            if element is None:
                continue
            grid[row_id + current_piece_position[0]][col_id + current_piece_position[1]] = element

    count = burn_rows()

    if count == 1:
        print("Single")
        rate = 1
    elif count == 2:
        print("Double")
        rate = 2.5
    elif count == 3:
        print("Triple")
        rate = 7.5
    elif count == 4:
        print("Tetris!")
        rate = 30
    else:
        rate = 0

    lines_burnt += count

    score += int(level * 40 * rate)

    if lines_burnt >= level * 10:
        level += 1
    current_piece, current_piece_position = next_piece, next_piece_position
    next_piece, next_piece_position = generate_piece()
    refresh_piece_preview_canvas()

PARSER = argparse.ArgumentParser()

PARSER.add_argument("--starting-level", type=int, choices=list(range(1, 30)), default=1)
PARSER.add_argument("--joystick", type=int, dest="joystick_id", metavar="Joystick ID")

ARGS = vars(PARSER.parse_args())

print(ARGS)

pygame.init()

black = (0, 0, 0)
white = (0xff, 0xff, 0xff)
brown = (163, 75, 31)
blue = (30, 34, 164)
green = (30, 164, 59)
red = (164, 30, 30)
purple = (126, 30, 164)
yellow = (164, 164, 30)
cyan = (30, 164, 150)

square_template = pygame.Surface((48, 48))

screen = pygame.display.set_mode((801, 1000))

game_canvas = pygame.Surface((500, 1000))
right_pane_canvas = pygame.Surface((300, 1000))
right_pane_canvas.fill((255, 255, 255))

piece_preview_canvas = pygame.Surface((200, 200))
score_canvas = pygame.Surface((296, 50))
level_canvas = pygame.Surface((296, 50))

grid = [[None for _ in range(10)] for _ in range(20)]

current_piece, current_piece_position = generate_piece()
next_piece, next_piece_position = generate_piece()

refresh_piece_preview_canvas()

# Number of frames
frames_per_gridcell = [48, 43, 38, 33, 28, 23, 18, 13, 8, 6, 5, 5, 5, 4, 4, 4, 3, 3, 3, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1]
piece_drop_frames = 0

starting_level = level = 1
score = 0

clock = pygame.time.Clock()

lines_burnt = 0

das = 0

pressing_down_countdown = None

if ARGS["joystick_id"] is not None:
    joystick = pygame.joystick.Joystick(ARGS["joystick_id"])
    joystick.init()

while True:
    piece_drop_frames += 1

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            sys.exit()
        elif event.type == pygame.KEYDOWN:
            if event.key in (pygame.K_ESCAPE, pygame.K_q):
                sys.exit()
            with contextlib.suppress(WouldCollide):
                if event.key in (pygame.K_RIGHT, pygame.K_l):
                    move_piece_right()
                    das = 0
                if event.key in (pygame.K_LEFT, pygame.K_h):
                    move_piece_left()
                    das = 0
                if event.key == pygame.K_s:
                    rotate_piece_clockwise()
                if event.key == pygame.K_d:
                    rotate_piece_counter_clockwise()
                if event.key in (pygame.K_DOWN, pygame.K_j):
                    piece_drop_frames = 0
                    pressing_down_countdown = 3
                    try:
                        move_piece_down()
                    except WouldCollide:
                        stick_piece()
        elif event.type == pygame.KEYUP:
            if event.key in (pygame.K_DOWN, pygame.K_j):
                pressing_down_countdown = None

        elif event.type == pygame.JOYBUTTONDOWN:
            with contextlib.suppress(WouldCollide):
                if event.button == PS3Controller.RIGHT.value:
                    move_piece_right()
                    das = 0
                if event.button == PS3Controller.LEFT.value:
                    move_piece_left()
                    das = 0
                if event.button == PS3Controller.CROSS.value:
                    rotate_piece_clockwise()
                if event.button == PS3Controller.CIRCLE.value:
                    rotate_piece_counter_clockwise()
                if event.button == PS3Controller.DOWN.value:
                    piece_drop_frames = 0
                    pressing_down_countdown = 3
                    try:
                        move_piece_down()
                    except WouldCollide:
                        stick_piece()
        elif event.type == pygame.JOYBUTTONUP:
            if event.button == PS3Controller.DOWN.value:
                pressing_down_countdown = None

    das += 1
    if das == 16:
        pressed_keys = pygame.key.get_pressed()
        pressed_buttons = [joystick.get_button(x) for x in range(joystick.get_numbuttons())]
        with contextlib.suppress(WouldCollide):
            if pressed_keys[pygame.K_RIGHT] or pressed_keys[pygame.K_l] or pressed_buttons[PS3Controller.RIGHT.value]:
                move_piece_right()
            if pressed_keys[pygame.K_LEFT] or pressed_keys[pygame.K_h] or pressed_buttons[PS3Controller.LEFT.value]:
                move_piece_left()
        das = 10

    if pressing_down_countdown == 0:
        try:
            move_piece_down()
        except WouldCollide:
            stick_piece()
        pressing_down_countdown = 2
    elif pressing_down_countdown is not None:
        pressing_down_countdown -= 1

    if piece_drop_frames >= frames_per_gridcell[level - 1]:
        piece_drop_frames = 0
        try:
            move_piece_down()
        except WouldCollide:
            stick_piece()

    refresh_screen()
    pygame.display.update()

    clock.tick(60)