import pygame as pg

from pygame.math import Vector2 as vec2
from pygame.time import Clock
from queue import Queue
from random import randint


class Ball:
    RADIUS = 10

    def __init__(self, surface, position, speed, color):
        self.surface = surface
        self.position = position
        self.speed = speed

        self.color = color

    def render(self):
        pg.draw.circle(self.surface, self.color, self.position, self.RADIUS)


class Game:
    def __init__(self, surface):
        self.surface = surface

        self.width = self.surface.get_width()
        self.height = self.surface.get_height()

        self.bg_color = (30, 40, 60)
        self.balls = []
        self.maxsize = 5

        self.clock = Clock()
        self.fps = 60

        self.mouse_down_pos = None
        self.current_mouse_pos = (0, 0)

        self.ball_color = self.get_next_ball_color()

        self.working = True

    @staticmethod
    def get_next_ball_color():
        return (randint(100, 255), randint(100, 255), randint(0, 100))

    def add_ball(self, position, diff: vec2):
        length = max(0.3, (diff / 10).length())
        length = min(length, 15)

        speed = diff.normalize() * length

        self.balls.append(Ball(self.surface, position, speed, self.ball_color))
        self.ball_color = self.get_next_ball_color()

        while len(self.balls) > self.maxsize:
            self.balls.pop(0)

    def process_events(self):
        for e in pg.event.get():
            if e.type == pg.QUIT or (e.type == pg.KEYDOWN and e.key == pg.K_ESCAPE):
                self.working = False

        self.current_mouse_pos = pg.mouse.get_pos()

        pressed = pg.mouse.get_pressed()[0]
        if pressed and self.mouse_down_pos is None:
            self.mouse_down_pos = self.current_mouse_pos

        if not pressed and self.mouse_down_pos is not None:
            if self.mouse_down_pos != self.current_mouse_pos:
                self.add_ball(
                    vec2(self.mouse_down_pos),
                    vec2(self.mouse_down_pos) - vec2(self.current_mouse_pos))
            self.mouse_down_pos = None

    def update(self):
        for b in self.balls:
            b.position += b.speed

            if b.position.x < Ball.RADIUS or b.position.x > self.width - Ball.RADIUS:
                b.speed = b.speed.reflect(vec2(1, 0))
                if b.position.x < Ball.RADIUS:
                    b.position.x = Ball.RADIUS
                else:
                    b.position.x = self.width - Ball.RADIUS

            if b.position.y < Ball.RADIUS or b.position.y > self.height - Ball.RADIUS:
                b.speed = b.speed.reflect(vec2(0, 1))
                if b.position.y < Ball.RADIUS:
                    b.position.y = Ball.RADIUS
                else:
                    b.position.y = self.height - Ball.RADIUS

    def render(self):
        self.surface.fill(self.bg_color)
        for b in self.balls:
            b.render()

        if self.mouse_down_pos is not None:
            pg.draw.circle(self.surface, self.ball_color, self.mouse_down_pos, Ball.RADIUS)
            pg.draw.line(self.surface, self.ball_color, self.mouse_down_pos, self.current_mouse_pos)

        pg.display.update()

    def run(self):
        while self.working:
            self.process_events()
            self.update()
            self.render()
            self.clock.tick(self.fps)


def main():
    pg.init()

    display_surface = pg.display.set_mode((1920, 1080), pg.FULLSCREEN)
    Game(display_surface).run()

    pg.quit()


if __name__ == "__main__":
    main()