import os
import pygame


class Sprite(object):

    def __init__(self, path=None, sprite=None):
        self.sprite = None
        self.rect = None
        self.mask = None
        self.mask_image = None

        if path is not None:
            self.load(path)

        if sprite is not None:
            self.set_sprite(sprite)

    def rebuild(self):
        self.rect = self.sprite.get_rect()
        self.mask = pygame.mask.from_surface(self.sprite)
        self.mask_image = self.mask.to_surface()

    def load(self, path):
        self.sprite = pygame.image.load(path).convert_alpha()
        self.rebuild()

    def set_sprite(self, sprite):
        self.sprite = sprite
        self.rebuild()

    def reflect(self, horizontal=False, vertical=True):
        self.sprite = pygame.transform.flip(self.sprite, horizontal, vertical)
        self.rebuild()

    def rotate(self, angle=0):
        angle = 360 - (angle - 360 * int(angle // 360))
        self.sprite = pygame.transform.rotate(self.sprite, angle)
        self.rebuild()

    def scale2x(self):
        self.sprite = pygame.transform.scale2x(self.sprite)
        self.rebuild()

    def set_x(self, x):
        self.rect.x = x

    def set_y(self, y):
        self.rect.y = y

    def set_width(self, width):
        self.rect.size[0] = width

    def set_height(self, height):
        self.rect.size[1] = height

    @property
    def get_x(self):
        return self.rect.x

    @property
    def get_y(self):
        return self.rect.y

    @property
    def get_width(self):
        return self.rect.size[0]

    @property
    def get_height(self):
        return self.rect.size[1]

    @property
    def get_sprite(self):
        return self.sprite

    @property
    def get_rect(self):
        return self.rect

    @property
    def get_mask(self):
        return self.mask

    @property
    def get_mask_image(self):
        return self.mask_image

    def draw_circle(self, radius):
        self.sprite = pygame.Surface((radius * 2, radius * 2), pygame.SRCALPHA)
        self.rect = self.sprite.get_rect()
        x = self.rect.size[0] / 2
        y = self.rect.size[1] / 2
        pygame.draw.circle(self.sprite, 'White', (x, y), radius)
        self.rebuild()

    def render(self, screen, position):
        x = position[0] - self.get_rect.size[0] / 2
        y = position[1] - self.get_rect.size[1] / 2

        screen.blit(self.get_sprite, (x, y))

    def render_mask(self, screen):
        screen.blit(self.get_mask_image, self.get_rect)


class Animation(object):

    def __init__(self, path):
        self.sprites = []
        self.rects = []
        self.masks = []
        self.masks_image = []
        self.step = 0

        if path is not None:
            self.load(path)

    def rebuild(self):
        self.rects = []
        self.masks = []
        for index, sprite in enumerate(self.sprites):
            self.rects.append(self.sprites[index].get_rect())
            self.masks.append(pygame.mask.from_surface(self.sprites[index]))
            self.masks_image.append(self.masks[len(self.masks) - 1].to_surface())

    def load(self, path):
        for dir_path, dir_names, file_names in os.walk(f'./{path}'):
            for file_name in file_names:
                if file_name[-3:] == 'png':
                    self.sprites.append(pygame.image.load(f'./{path}/{file_name}').convert_alpha())
        self.rebuild()

    def reflect(self, horizontal=False, vertical=True):
        for index, sprite in enumerate(self.sprites):
            self.sprites[index] = pygame.transform.flip(sprite, horizontal, vertical)
        self.rebuild()

    def rotate(self, angle=0):
        angle = 360 - (angle - 360 * int(angle // 360))
        for index, sprite in enumerate(self.sprites):
            self.sprites[index] = pygame.transform.rotate(sprite, angle)
        self.rebuild()

    def scale2x(self):
        for index, sprite in enumerate(self.sprites):
            self.sprites[index] = pygame.transform.scale2x(sprite)
        self.rebuild()

    def get_sprite(self, index):
        return self.sprites[index]

    def get_sprites(self):
        return self.sprites

    def get_count(self):
        return len(self.sprites)

    def get_rect(self, index):
        return self.rects[index]

    def get_mask(self, index):
        return self.masks[index]

    def generator(self):
        return (sprite for sprite in self.get_sprites())

    def reset_step(self):
        self.step = 0

    def render(self, screen, position, tick):
        sprite_index = int(self.step / (tick / self.get_count()))
        x = position[0] - self.get_rect(sprite_index).size[0] / 2
        y = position[1] - self.get_rect(sprite_index).size[1] / 2
        screen.blit(self.get_sprite(sprite_index), (x, y))

        self.step += 1
        if self.step == tick:
            self.step = 0
            return True
        else:
            return False
