#!/usr/bin/python3

import os
from random import randrange
from settings import *


class Cache:
    def __init__(self):
        self.stacked_sprite_cache = {}
        self.viewing_angle = 360 // NUM_ANGLES
        self.outline_thickness = 2
        self.get_stacked_sprite_cache()

    def get_stacked_sprite_cache(self):
        for object_name in SPRITE_ATTRS:
            self.stacked_sprite_cache[object_name] = {
                "rotated_sprites": {}
            }

            attrs = SPRITE_ATTRS[object_name]
            layer_array = self.get_layer_array(attrs)
            self.run_prerender(object_name, layer_array, attrs)

        for object_name in ANIMATED_SPRITE_ATTRS:
            self.stacked_sprite_cache[object_name] = {
                "frames": {}
            }

            attrs = ANIMATED_SPRITE_ATTRS[object_name]
            frames = self.get_frame_dict(attrs)
            self.run_animation_prerender(object_name, frames, attrs)

    def run_prerender(self, object_name, layer_array, attrs):
        outline = attrs.get("outline", OUTLINE)

        for angle in range(NUM_ANGLES):
            surface = pygame.Surface(layer_array[0].get_size())
            surface = pygame.transform.rotate(surface, angle * self.viewing_angle)
            sprite_surface = pygame.Surface([
                surface.get_width(),
                surface.get_height() + attrs["layers"] * attrs["scale"]
            ])

            sprite_surface.fill(TRANSP_COLOR)
            sprite_surface.set_colorkey(TRANSP_COLOR)

            for i, layer in enumerate(layer_array):
                layer = pygame.transform.rotate(layer, angle * self.viewing_angle)
                sprite_surface.blit(layer, (0, i * attrs["scale"]))

            # get outline
            if outline:
                outline_coords = pygame.mask.from_surface(sprite_surface).outline()
                pygame.draw.polygon(sprite_surface, "black", outline_coords, self.outline_thickness)

            image = pygame.transform.flip(sprite_surface, True, True)
            self.stacked_sprite_cache[object_name]["rotated_sprites"][angle] = image

    def run_animation_prerender(self, object_name, frames, attrs):
        outline = attrs.get("outline", OUTLINE)

        for angle in range(NUM_ANGLES):
            surface = pygame.Surface(list(frames.items())[0][1][0].get_size())
            surface = pygame.transform.rotate(surface, angle * self.viewing_angle)
            surface_default = pygame.Surface([
                    surface.get_width(),
                    surface.get_height() + attrs["layers"] * attrs["scale"]
                ])

            for frame_name, layer_array in frames.items():
                sprite_surface = surface_default

                sprite_surface.fill(TRANSP_COLOR)
                sprite_surface.set_colorkey(TRANSP_COLOR)

                for i, layer in enumerate(layer_array):
                    layer = pygame.transform.rotate(layer, angle * self.viewing_angle)
                    sprite_surface.blit(layer, (0, i * attrs["scale"]))

                # get outline
                if outline:
                    outline_coords = pygame.mask.from_surface(sprite_surface).outline()
                    pygame.draw.polygon(sprite_surface, "black", outline_coords, self.outline_thickness)

                image = pygame.transform.flip(sprite_surface, True, True)

                if not frame_name in self.stacked_sprite_cache[object_name]["frames"]:
                    self.stacked_sprite_cache[object_name]["frames"][frame_name] = {"rotated_sprites": {}}

                self.stacked_sprite_cache[object_name]["frames"][frame_name]["rotated_sprites"][angle] = image

    def get_layer_array(self, attrs, path=None):
        if not path is None:
            attrs["path"] = path

        # load sprite sheet
        sprite_sheet = pygame.image.load(attrs["path"]).convert_alpha()
        # scaling
        sprite_sheet = pygame.transform.scale(sprite_sheet, vec2(sprite_sheet.get_size()) * attrs["scale"])
        sheet_width, sheet_height = sprite_sheet.get_size()
        sprite_width = sheet_width // attrs["layers"]
        # new width to prevent error
        sheet_width = sprite_width * attrs["layers"]

        # get sprites
        layer_array = []
        for x in range(0, sheet_width, sprite_width):
            sprite = sprite_sheet.subsurface((x, 0, sprite_width, sheet_height))
            layer_array.append(sprite)

        return layer_array

    def get_frame_dict(self, attrs):
        frames = {}

        for path, folders, files in os.walk(attrs["path"]):  # loop through all frames and add them to the sheet list
            for file in files:
                # use the name of the file without the file format specifier (the part before the dot) as the frame name
                # and store the layer array
                frames[file.split(".")[0]] = self.get_layer_array(attrs, f"{path}/{file}")

        return frames