#!/usr/bin/env python3

import argparse
import contextlib
from functools import partial
import random
import sys

import pygame

from pycameltris.controller import KeyboardController, JoystickController, Input


class WouldCollide(Exception):
    pass


class PlayerQuit(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_right_pane_canvas():
    for player in players:
        player.level_canvas.fill(black)
        player.score_canvas.fill(black)

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

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

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


def refresh_screen():
    for player in players:
        player.refresh_grid_canvas()
        screen.blit(player.grid_canvas, (0, 0))
        refresh_right_pane_canvas()
        screen.blit(right_pane_canvas, (501, 0))


class Player:
    def __init__(self, controller):
        self.controller = controller
        self.grid = [[None for _ in range(10)] for _ in range(20)]

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

        self.starting_level = self.level = 1
        self.score = 0

        self.lines_burnt = 0

        self.das = 0
        self.pressing_down_countdown = None

        self.piece_drop_frames = 0

        self.grid_canvas = pygame.Surface((500, 1000))
        self.piece_preview_canvas = pygame.Surface((200, 200))
        self.score_canvas = pygame.Surface((296, 50))
        self.level_canvas = pygame.Surface((296, 50))


    def generate_piece(self):
        # We may want to make this a function outside the class
        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(self.grid[0]) // 2) - (len(piece.elements[0]) // 2)

        return (piece, [initial_y_position, initial_x_position])

    def burn_rows(self):
        rows_to_burn = list()

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

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

        return len(rows_to_burn)

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

        count = self.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

        self.lines_burnt += count

        self.score += int(self.level * 40 * rate)

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

    def has_collision(self, y: int, x: int) -> bool:
        try:
            for row_id, row in enumerate(self.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 self.grid[row_id + y][col_id + x] is not None:
                        return True

        except IndexError:
            return True

        return False

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

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

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

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

    def rotate_piece_counter_clockwise(self):
        self.current_piece.rotate_counter_clockwise()

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

    def rotate_piece_clockwise(self):
        self.current_piece.rotate_clockwise()

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

    def handle_input_pressed(self, event):
        if self.controller.get_input_down(event) == Input.QUIT:
            raise PlayerQuit()

        with contextlib.suppress(WouldCollide):
            if self.controller.get_input_down(event) == Input.MOVE_RIGHT:
                self.move_piece_right()
                self.das = 0
            if self.controller.get_input_down(event) == Input.MOVE_LEFT:
                self.move_piece_left()
                self.das = 0
            if self.controller.get_input_down(event) == Input.ROTATE_CLOCKWISE:
                self.rotate_piece_clockwise()
            if self.controller.get_input_down(event) == Input.ROTATE_COUNTER_CLOCKWISE:
                self.rotate_piece_counter_clockwise()
            if self.controller.get_input_down(event) == Input.MOVE_DOWN:
                self.piece_drop_frames = 0
                self.pressing_down_countdown = 3
                try:
                    self.move_piece_down()
                except WouldCollide:
                    self.stick_piece()

    def handle_input_released(self, event):
        if self.controller.get_input_up(event) == Input.MOVE_DOWN:
            self.pressing_down_countdown = None

    def refresh_piece_preview_canvas(self):
        self.piece_preview_canvas.fill(black)

        non_empty_rows = list()
        for row in self.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 self.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:
                    self.piece_preview_canvas.blit(element, ((col_idx + x_offset) * 50 + 1, (row_idx + y_offset) * 50 + 1))

    def refresh_grid_canvas(self):
        self.grid_canvas.fill(black)

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

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


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

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


if ARGS["joystick_id"] is not None:
    joystick = pygame.joystick.Joystick(ARGS["joystick_id"])
    joystick.init()
    controller = JoystickController(joystick)
else:
    controller = KeyboardController(pygame.key)

# Just one player
players = [Player(controller)]

for player in players:
    player.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]

clock = pygame.time.Clock()

def handle_input_pressed(instance, event):
    for player in players:
        if isinstance(player.controller, instance):
            player.handle_input_pressed(event)


def handle_input_released(instance, event):
    for player in players:
        if isinstance(player.controller, instance):
            player.handle_input_released(event)

event_handler = dict()
event_handler[pygame.QUIT] = lambda _: sys.exit()
event_handler[pygame.KEYDOWN] = partial(handle_input_pressed, KeyboardController)
event_handler[pygame.KEYUP] = partial(handle_input_released, KeyboardController)
event_handler[pygame.JOYBUTTONDOWN] = partial(handle_input_pressed, JoystickController)
event_handler[pygame.JOYBUTTONUP] = partial(handle_input_released, JoystickController)

while True:
    for player in players:
        player.piece_drop_frames += 1

    for event in pygame.event.get():
        with contextlib.suppress(KeyError):
            event_handler[event.type](event)

    for player in players:
        player.das += 1
        if player.das == 16:
            with contextlib.suppress(WouldCollide):
                if controller.is_pressed(Input.MOVE_RIGHT):
                    player.move_piece_right()
                if controller.is_pressed(Input.MOVE_LEFT):
                    player.move_piece_left()
            player.das = 10

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

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

    refresh_screen()
    pygame.display.update()

    clock.tick(60)