#!/usr/bin/env python3

import contextlib
import random
import sys
import time

import pygame


class WouldCollide(Exception):
    pass


class Piece():
    def rotate(self):
        pass

    def rotate_clockwise(self):
        self.rotate()

    def rotate_counter_clockwise(self):
        self.rotate()


class ZPiece(Piece):
    def __init__(self):
        self.elements = (blue_square, blue_square, None), (None, blue_square, blue_square), (None, None, None)
        self.vertical = False

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

        self.vertical = not self.vertical


class SPiece(Piece):
    def __init__(self):
        self.elements = (None, None, None), (None, green_square, green_square), (green_square, green_square, None)
        self.vertical = False

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

        self.vertical = not self.vertical


class SquarePiece(Piece):
    def __init__(self):
        self.elements = ((white_square, white_square), (white_square, white_square))


class IPiece(Piece):
    def __init__(self):
        self.elements = (None, None, None, None), (None, None, None, None), (red_square, red_square, red_square, red_square), (None, None, None, None)
        self.vertical = False

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

        self.vertical = not self.vertical


class TPiece(Piece):
    def __init__(self):
        self.elements = [[yellow_square, yellow_square, yellow_square], [None, yellow_square, None], [None, None, None]]

    def rotate_clockwise(self):
        # Set the corners
        self.elements[0][0], self.elements[0][2], self.elements[2][2], self.elements[2][0] = self.elements[2][0], self.elements[0][0], self.elements[0][2], self.elements[2][2]
        # Set the middle squares
        self.elements[0][1], self.elements[1][2], self.elements[2][1], self.elements[1][0] = self.elements[1][0], self.elements[0][1], self.elements[1][2], self.elements[2][1]

    def rotate_counter_clockwise(self):
        # Set the corners
        self.elements[2][0], self.elements[0][0], self.elements[0][2], self.elements[2][2] = self.elements[0][0], self.elements[0][2], self.elements[2][2], self.elements[2][0]
        # Set the middle squares
        self.elements[1][0], self.elements[0][1], self.elements[1][2], self.elements[2][1] = self.elements[0][1], self.elements[1][2], self.elements[2][1], self.elements[1][0]


def refresh_screen():
    screen.fill(black)

    for row_idx, row in enumerate(grid):
        for col_idx, element in enumerate(row):
            if element is not None:
                screen.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:
                screen.blit(element, ((col_idx + current_piece_position[1]) * 50 + 1, (row_idx + current_piece_position[0]) * 50 + 1))


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))()

    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)


def stick_piece():
    global current_piece, current_piece_position

    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

    burn_rows()
    current_piece, current_piece_position = generate_piece()

pygame.init()

black = (0, 0, 0)
white = (0xff, 0xff, 0xff)
blue = (0x10, 0x20, 0xbb)
green = (0x20, 0xbb, 0x10)
red = (0xbb, 0x10, 0x20)
yellow = (0xab, 0xd0, 0x20)

white_square = pygame.Surface((48, 48))
white_square.fill(white)

blue_square = pygame.Surface((48, 48))
blue_square.fill(blue)

green_square = pygame.Surface((48, 48))
green_square.fill(green)

red_square = pygame.Surface((48, 48))
red_square.fill(red)

yellow_square = pygame.Surface((48, 48))
yellow_square.fill(yellow)

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

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

current_piece, current_piece_position = generate_piece()

while True:
    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_UP, pygame.K_k):
                    move_piece_up()
                if event.key in (pygame.K_RIGHT, pygame.K_l):
                    move_piece_right()
                if event.key in (pygame.K_LEFT, pygame.K_h):
                    move_piece_left()
                if event.key == pygame.K_s:
                    rotate_piece_clockwise()
                if event.key == pygame.K_d:
                    rotate_piece_counter_clockwise()
            try:
                if event.key in (pygame.K_DOWN, pygame.K_j):
                    move_piece_down()
            except WouldCollide:
                stick_piece()

    refresh_screen()
    pygame.display.flip()
    time.sleep(.1)