from __future__ import annotations
import pygame
import sys
from dataclasses import dataclass, field
import time
import numpy as np
import typing
import json
import os
from random import uniform, choice
if typing.TYPE_CHECKING:
    from nptyping import NDArray


pygame.init()
# width and height of the main part of the screen + vertical offset for general info
W, H, OFFSET = 900, 800, 50
# ball's lifetime (seconds)
DEFAULT_TTL = 3
# ball's radius (default)
DEFAULT_RADIUS = 20
# time delta, ms
dt = 10
# load the font
font = pygame.font.Font('orbitron-medium.otf', 15)
font_large = pygame.font.Font('orbitron-medium.otf', 25)
font_even_larger = pygame.font.Font('orbitron-medium.otf', 40)
# Orbitron doesn't support unicode, so I have to get another font for the scoreboard
unicode_font_large = pygame.font.Font('IntroRust.otf', 25)

N_BALLS = 18
ENDGAME_BUTTON_RECT = (850, 15, 20, 20)
BGCOLOR = (40, 46, 70)

# List of ball types: [(radius, speed, color, points, shape, deviation_radius)]
BALL_TYPES = [
        (25, 0.035, (185, 20, 50), 10, 'circle', 0), 
        (20, 0.05, (60, 0, 160), 20, 'circle', 0),
        (18, 0.055, (220, 0, 0), 30, 'circle', 0), 
        (30, 0.055, (220, 0, 0), 40, 'triangle', 30), 
    ]


@dataclass
class Game:
    """
    Class for everything connected to the game

    :param score: user's score
    :param balls: the list of balls
    :param target_ball_number: the target number of balls on screen at a given time
    :param game_end_time: the timeout of the game
    :param game_ended: true if the game has ended
    :param user_name: user's name for the scoreboard
    """
    score: int = 0
    balls: typing.List[Ball] = field(default_factory=lambda: [Ball.new_random() for _ in range(N_BALLS)])
    target_ball_number: int = N_BALLS
    game_end_time: float = field(default_factory=lambda: time.time() + 30)
    game_ended: bool = False
    scoreboard: typing.List[typing.Tuple[str, int]] = field(default_factory=lambda: list() 
            if not os.path.exists('scoreboard.json') 
            else json.load(open('scoreboard.json')))
    user_name: str = ''

    def draw_everything(self, screen: pygame.Surface):
        """
        Draw everything (for example, balls) on screen
        :param screen: the screen
        """
        for ball in self.balls:
            ball.draw_the_ball(screen)
        if not self.game_ended:
            pygame.draw.rect(screen, (255, 255, 255), pygame.Rect(0, 0, W, OFFSET))
            text_surface = font_large.render(f'Score: {self.score}', True, (0, 0, 0))
            screen.blit(text_surface, (360, 15))
            time_left = self.game_end_time - time.time()
            time_text_surface = font_large.render(
                    f'{int(time_left / 60)}:{int(time_left % 60):02}', False, 
                    (0, 0, 0) if time_left > 10 else (220, 0, 0))
            screen.blit(time_text_surface, (760, 15))
            # draw the button to end game (which is a red rect)
            pygame.draw.rect(screen, (255, 0, 0), ENDGAME_BUTTON_RECT)
        else:
            screen.fill((*BGCOLOR, 150), None, pygame.BLEND_RGBA_MULT)
            game_over_text = font_even_larger.render(f'Game over', True, (255, 0, 0))
            score_text = font_even_larger.render(f'Score: {self.score}', True, (255, 255, 255))
            screen.blit(game_over_text, (W / 2 - game_over_text.get_width() / 2, 
                -180 + H / 2 - game_over_text.get_height() / 2))
            screen.blit(score_text, (W / 2 - score_text.get_width() / 2, -120 + H / 2 - score_text.get_height() / 2))
            y = H / 2
            if len(self.scoreboard):
                scoreboard_title = font_large.render('Scoreboard:', False, (255, 255, 255))
                screen.blit(scoreboard_title, (W / 2 - scoreboard_title.get_width() / 2, 
                    -40 + H / 2 - scoreboard_title.get_height() / 2))
                y = -30 + H / 2 + scoreboard_title.get_height() / 2
                # Zip with a range of 5 instead of enumerating so that we show only the first 5 records
                for i, record in zip(range(5), sorted(self.scoreboard, key=lambda record: -record[1])):
                    text = unicode_font_large.render(f'{i + 1}. {record[0]} - {record[1]}', True, (255, 255, 255))
                    y += 10 + text.get_height()
                    screen.blit(text, (200, y))
            username_prompt = font_large.render('Your name:', False, (255, 255, 255))
            screen.blit(username_prompt, (200, y + 100))
            username_current = unicode_font_large.render(self.user_name, False, (255, 255, 255))
            pygame.draw.rect(
                    screen, 
                    [int(v * 0.5 + 255 * 0.5) for v in BGCOLOR], 
                    (200 + username_prompt.get_width() + 30, y + 80, max(400, username_current.get_width() + 20),
                        username_current.get_height() + 20))
            screen.blit(username_current, (200 + username_prompt.get_width() + 40, y + 90))


    def process_click(self, x: float, y: float):
        """
        Process mouse click, check if the user clicked on a ball or the endgame button and process it
        :param x: click's x coordinate
        :param y: click's y coordinate
        """
        e_x = ENDGAME_BUTTON_RECT[0] <= x <= ENDGAME_BUTTON_RECT[0] + ENDGAME_BUTTON_RECT[2]
        e_y = ENDGAME_BUTTON_RECT[1] <= y <= ENDGAME_BUTTON_RECT[1] + ENDGAME_BUTTON_RECT[3]
        if e_x and e_y:
            self.game_ended = True
            self.user_name = ''
            return
        if self.game_ended:
            return
        for i in range(len(self.balls)):
            if self.balls[i].is_point_inside(x, y - OFFSET):
                self.score += self.balls[i].points
                self.balls.pop(i)
                break

    def process_keypress(self, event):
        """
        Process keypress for name input
        :param event: the keypress event
        """
        if event.key == pygame.K_BACKSPACE:
            self.user_name = self.user_name[:-1]
        else:
            self.user_name += event.unicode
        self.write_scoreboard()

    def tick(self, dt: float):
        """
        Progress for some time

        dt: time period, ms
        """
        if time.time() > self.game_end_time and not self.game_ended:
            self.game_ended = True
            self.user_name = ''
        if len(self.balls) < self.target_ball_number:
            n_balls = round((self.target_ball_number - len(self.balls))) # * (1 - 0.9 ** dt))
            self.balls.extend([Ball.new_random() for _ in range(n_balls)])
        for ball in self.balls:
            ball.tick(dt)

    def write_scoreboard(self):
        """
        Save scoreboard (`self.scoreboard`) to scoreboard.json as JSON with the new record
        """
        if not self.user_name:
            return
        with open('scoreboard.json', 'w') as f:
            f.write(json.dumps(sorted(self.scoreboard + [(self.user_name, self.score)], key=lambda record: -record[1])))


@dataclass
class Ball:
    """
    The ball the user has to catch

    :param points: the number of points the user gets for catching this ball
    :param position: ball's position on the screen
    :param velocity: ball's velocity, px/ms
    :param time_of_death: the time when this ball will be destroyed. Its death is pre-determined, which is probably
    very sad and depressing
    :param radius: the radius of the ball on screen in pixels
    :param color: ball's color
    :param shape: shape of the ball (circle, triangle)
    :param deviation_radius: if the radius is >0, the ball will be rotating arond its actual trajectory
    :param deviation_angle: the phase of the ball's rotation
    """
    points: int = 10
    position: NDArray[(2,), float] = field(default_factory=lambda: np.array([uniform(0, W), uniform(0, H)]))
    velocity: NDArray[(2,), float] = field(default_factory=lambda: np.array([0, 0]))
    time_of_death: float = field(default_factory=lambda: time.time() + DEFAULT_TTL)
    radius: float = DEFAULT_RADIUS
    color: typing.Tuple[int, int, int] = (255, 0, 0)
    shape: str = 'circle'
    deviation_radius: float = 0
    deviation_angle: float = 0

    @classmethod
    def new_random(cls) -> Ball:
        radius, speed, color, points, shape, deviation_radius = choice(BALL_TYPES)
        angle = uniform(0, np.pi)
        return cls(points, radius=radius, color=color, velocity=speed * np.array([np.sin(angle), np.cos(angle)]),
                shape=shape, deviation_radius=deviation_radius)

    def draw_the_ball(self, screen: pygame.Surface):
        """
        Just draw the ball on the screen
        :param screen: the screen, obviously
        """
        x, y = self.position_with_shift()[0], self.position_with_shift()[1] + OFFSET
        if self.shape == 'triangle':
            angles = [0, np.pi * 2 / 3, np.pi * 4 / 3]
            pygame.draw.polygon(screen, self.color, [
                (
                    x + self.radius * np.sin(angle + self.deviation_angle), 
                    y + self.radius * np.cos(angle + self.deviation_angle)
                ) for angle in angles])
        else:
            pygame.draw.circle(screen, self.color, [x, y], self.radius)
        text_surface = font.render(str(self.points), False, (0, 0, 0))
        screen.blit(text_surface, (
            round(x - text_surface.get_rect().width / 2),
            round(y - text_surface.get_rect().height / 2)
        ))

    def position_with_shift(self) -> typing.Tuple[float, float]:
        """
        Get position including deviation (rotation arond the true center)
        """
        return (
                self.position[0] + self.deviation_radius * np.cos(self.deviation_angle), 
                self.position[1] + self.deviation_radius * np.sin(self.deviation_angle)
            )

    def is_point_inside(self, x: float, y: float) -> bool:
        """
        Check if a given point is inside this ball (to handle click events)
        :param x: point's x coordinate
        :param y: point's y coordinate
        """
        return (self.position_with_shift()[0] - x) ** 2 + (self.position_with_shift()[1] - y) ** 2 <= self.radius ** 2

    def tick(self, dt: float):
        """
        Progress for time dt (move the ball according to speed)
        dt: time in ms
        """
        self.position += self.velocity * dt
        if self.position[0] - self.radius < 0:
            self.position[0] = self.radius
            self.velocity[0] *= -1
        if self.position[0] >= W - self.radius:
            self.position[0] = W - self.radius
            self.velocity[0] *= -1
        if self.position[1] < self.radius:
            self.position[1] = self.radius
            self.velocity[1] *= -1
        if self.position[1] >= H - self.radius:
            self.position[1] = H - self.radius
            self.velocity[1] *= -1
        self.deviation_angle += dt * 2 * np.pi / 2000


# create the display surface object  
screen = pygame.display.set_mode((W, H + OFFSET))

game = Game()
  
# set the window name 
pygame.display.set_caption("Capture the ball!")
    
# the main loop
while True:
    # time delay of dt
    pygame.time.delay(dt)
      
    # iterate over the list of Event objects  
    # that was returned by pygame.event.get() method.  
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit(0)
        if event.type == pygame.MOUSEBUTTONUP:
            game.process_click(*pygame.mouse.get_pos())
        if event.type == pygame.KEYDOWN:
            game.process_keypress(event)

    # clear the screen
    screen.fill(BGCOLOR)
    game.tick(dt)
    game.draw_everything(screen)
    
    pygame.display.update()