from typing import Any, List, Tuple
import pygame as pg
from effects.effect import Effect, Colors
import random
from typing import Union, Generator


def fade_statemachine(
    fade_in: int, hold: int, fade_out: int
) -> Generator[Tuple[int, bool, bool], None, None]:
    for t in range(fade_in):
        yield 255 * t // fade_in, False, False
    for t in range(hold):
        yield 255, False, False
    for t in range(fade_out):
        yield 255 * (fade_out - t) // fade_out, True, False
    yield 0, True, True


class Spot(pg.sprite.Sprite):
    def __init__(
        self,
        color: pg.Color,
        position: Tuple[int, int],
        radius: int,
        fade_in: int,
        hold: int,
        fade_out: int,
    ) -> None:
        super().__init__()
        self.rect = pg.Rect(
            position[0] - radius, position[1] - radius, radius * 2, radius * 2
        )
        self.image = pg.Surface(self.rect.size)
        self.image.set_colorkey(Colors.Black)
        self.image.fill(Colors.Black)
        pg.draw.ellipse(
            self.image,
            color,
            ((0, 0), self.rect.size),
        )
        self.start_next = False
        self.finished = False
        self.state = fade_statemachine(fade_in, hold, fade_out)

    def update(self) -> None:
        if not self.finished:
            alpha, self.start_next, self.finished = next(self.state)
            self.image.set_alpha(alpha)

    def draw(self, dest: pg.Surface) -> None:
        dest.blit(self.image, self.rect)


class Starfield(Effect):
    def __init__(
        self,
        bounds: pg.Rect,
        color: Union[pg.Color, Generator[pg.Color, None, None]],
        radius: int = 20,
        star_factor: float = 0.2,
        hold: int = 60 * 3,
        fade_out: bool = False,
        fade_in: bool = False,
        beat_adapt: bool = False,
        *groups: pg.sprite.Group
    ) -> None:
        self.color = color
        self.radius: int = radius * 2
        self.fade_in_time: int = hold // 10 if fade_in else 0
        self.fade_out_time: int = int(hold * 0.8 if fade_out else 0)
        self.hold_time: int = hold - (self.fade_in_time + self.fade_out_time)
        self.beat_adapt: bool = beat_adapt
        self.stars: List[Spot] = []

        self.randrange = (
            (bounds.width - 2 * self.radius) // (self.radius * 3),
            (bounds.height - 2 * self.radius) // (self.radius * 2),
        )

        self.num_stars = int(self.randrange[0] * self.randrange[1] * star_factor)

        # self.time_to_next = (self.hold_time + self.fade_in_time) // 3
        self.time_to_next = hold // 5

        image = pg.Surface(size=bounds.size)
        image.set_colorkey(Colors.Black)
        super().__init__(image, bounds, *groups)

        self.update(is_beat=False, frames_per_beat=0.0)

    def add_star(self, spot_color):
        loc_y = random.randint(0, self.randrange[1])
        pos_y = loc_y * self.radius * 2 + self.radius
        pos_x = (
            random.randint(0, self.randrange[0]) * self.radius * 2
            + self.radius
            + (self.radius if (loc_y & 1 != 0) else 0)
        )

        while any(
            ((s.rect.centerx == pos_x and s.rect.centery == pos_y) for s in self.stars)
        ):
            loc_y = random.randint(0, self.randrange[1])
            pos_y = loc_y * self.radius * 2 + self.radius
            pos_x = (
                random.randint(0, self.randrange[0]) * self.radius * 2
                + self.radius
                + (self.radius if (loc_y & 1 != 0) else 0)
            )

        self.stars.append(
            Spot(
                color=spot_color,
                position=(pos_x, pos_y),
                radius=self.radius // 2,
                fade_in=random.randint(0, self.fade_in_time * 2),
                hold=random.randint(0, self.hold_time * 2),
                fade_out=random.randint(0, self.fade_out_time * 2),
            )
        )

    def update(self, *args: Any, **kwargs: Any) -> None:
        self.image.fill(Colors.Black)

        if self.beat_adapt and kwargs["frames_per_beat"]:
            hold = 3*kwargs["frames_per_beat"] - 2
            self.fade_in_time = int(hold // 10 if self.fade_in_time else 0)
            self.fade_out_time = int(hold * 0.8 if self.fade_out_time else 0)
            self.hold_time = int(hold - self.fade_in_time)
        else:
            hold = self.fade_in_time + self.hold_time + self.fade_out_time

        if (
            self.time_to_next == 0 and (not self.beat_adapt or kwargs["is_beat"])
        ) and len(self.stars) < int(self.num_stars * 0.8):
            missing_stars = self.num_stars - len(self.stars)
            for _ in range(random.randint(missing_stars // 3, missing_stars//2)):
                star_color = (
                    self.color if isinstance(self.color, pg.Color) else next(self.color)
                )
                self.add_star(star_color)

            self.time_to_next = hold // 5
        else:
            self.time_to_next -= 1
            if self.time_to_next < 0:
                self.time_to_next = 0

        for spot in self.stars:
            if not spot.finished:
                spot.update()
                spot.draw(self.image)

        if self.stars:
            first_spot = self.stars[0]
            while self.stars and first_spot.finished:
                first_spot = self.stars.pop(0)
                del first_spot
                if self.stars:
                    first_spot = self.stars[0]