from argparse import ArgumentParser
from typing import Iterable, Tuple
from pygame.locals import *
import math
import pygame as pg
import random
import sys
import time

from effects.bouncingspot import BouncingSpot
from effects.effect import Effect, color_wheel


def print_displays() -> None:
    pg.display.init()
    display = pg.display.get_desktop_sizes()
    for i, d in enumerate(display):
        print(f"#{i} - {d[0]}x{d[1]}")
    sys.exit(0)


def initialize(
    display_id: int, windowed: bool, trails: bool
) -> Tuple[pg.Surface, pg.Surface]:
    pg.display.init()
    displays = pg.display.get_desktop_sizes()
    if not 0 <= display_id < len(displays):
        raise ValueError(
            f"Display ID {display_id} invalid. Must be between 0 and {len(displays)}!"
        )

    win = pg.display.set_mode(
        size=displays[display_id]
        if not windowed
        else (displays[display_id][0] // 2, displays[display_id][1] // 2),
        flags=pg.FULLSCREEN if not windowed else 0,
        display=display_id,
    )
    if trails:
        background = pg.Surface(win.get_size(), flags=pg.SRCALPHA)
        background.fill(pg.Color(0, 0, 0, 5))
    else:
        background = pg.Surface(win.get_size())
        background.fill(pg.Color(0, 0, 0))
    return win, background


def render_loop_normal(
    window: pg.Surface,
    background: pg.surface,
    effects: Iterable[Effect],
    clock: pg.time.Clock,
) -> None:
    blackout = False

    while True:
        for event in pg.event.get():
            if event.type == K_SPACE:
                blackout = not blackout
            if event.type == QUIT:
                pg.quit()
                sys.exit()

        window.blit(background, (0, 0))
        if not blackout:
            for e in effects:
                e.update()
                e.draw(window)
        yield


def render_loop_3d(
    window: pg.Surface,
    background: pg.surface,
    effects: Iterable[Effect],
    clock: pg.time.Clock,
) -> None:
    stage = pg.Surface(size=window.get_size())
    stage.fill(Color(0, 0, 0))

    full_size = stage.get_size()
    scaled_sizes = [
        (((full_size[0] * i) // 100), ((full_size[1] * i) // 100)) for i in range(101)
    ]

    scaled_positions = [
        ((full_size[0] - s[0]) // 2, (full_size[1] - s[1]) // 2) for s in scaled_sizes
    ]

    while True:
        for event in pg.event.get():
            if event.type == QUIT:
                pg.quit()
                sys.exit()

        stage.blit(background, (0, 0))
        for e in effects:
            e.update()
            e.draw(stage)

        window.fill(Color(0, 0, 0))

        stage.set_colorkey(Color(0, 0, 0))
        for i in range(0, 101, 5):
            stage.set_alpha(192 - 192 * ((i / 100) ** 0.5))
            window.blit(
                pg.transform.scale(stage, scaled_sizes[int(100 * (i/100)**2)]),
                scaled_positions[int(100 * (i/100)**2)],
            )
        yield


def main() -> None:
    argparser = ArgumentParser(
        description="beamshow - Render a light show for a video projector"
    )
    argparser.add_argument(
        "--3d",
        dest="render3d",
        action="store_true",
        help="Render a 3D preview instead of the normal output",
    )
    argparser.add_argument(
        "--list-displays", action="store_true", help="Show available displays"
    )
    argparser.add_argument(
        "-w", "--window", action="store_true", help="Display in a window"
    )
    argparser.add_argument(
        "--trails", action="store_true", help="Fade patterns out (trail mode)"
    )
    argparser.add_argument("--fps", action="store_true", help="Show FPS in console")
    argparser.add_argument(
        "-d", "--display", type=int, default=0, help="ID of the display to use"
    )

    # print some nice output after the pygame banner to separate our stuff from theirs
    print("")
    print("-" * (len(argparser.description) + 4))
    print(f"  {argparser.description}  ")
    print("-" * (len(argparser.description) + 4))
    print("")

    args = argparser.parse_args()

    if args.list_displays:
        print_displays()

    window, background = initialize(
        display_id=args.display, windowed=args.window, trails=args.trails
    )

    effects = [
        BouncingSpot(
            bounds=window.get_rect(),
            color=pg.Color(255, 255, 255),
            # color=color_wheel(),
            sizes=(window.get_height() / 8, window.get_height() / 8),
            velocity=(1, 1),
            x_factor=(1, 1),
            y_factor=(2.2, 2.2),
        ),
        BouncingSpot(bounds=window.get_rect(), color=color_wheel()),
        BouncingSpot(bounds=window.get_rect(), color=color_wheel(hue=180)),
    ]

    clock = pg.time.Clock()

    if args.render3d:
        loop = render_loop_3d(window, background, effects, clock)
    else:
        loop = render_loop_normal(window, background, effects, clock)

    while True:
        next(loop)
        pg.display.flip()
        clock.tick(60)
        if args.fps:
            print(clock.get_fps())


if __name__ == "__main__":
    main()