import pygame
from pocket_friends.elements import sprites
from pocket_friends.elements import surface


class Surface(surface.GameSurface):
    def __init__(self, game_res, resources_dir, game_fps, **kwargs):
        super().__init__(game_res, resources_dir, game_fps)

        preselected_color = None
        for key in kwargs.keys():
            if key == 'selected_color':
                preselected_color = kwargs.get(key)

        egg_list = [
            'dev_egg',
            'blue',
            'rainbow'
        ]
        self.eggs = []
        for egg in egg_list:
            self.eggs.append(sprites.SelectionEgg(egg, self.resource_dir))

        self.eggs_per_row = 3
        distance_between_eggs = 36 / self.eggs_per_row

        # Count the total rows.
        self.total_rows = -(-len(self.eggs) // self.eggs_per_row)
        distance_between_rows = 32 / self.eggs_per_row

        # Determine the location of each egg.
        for egg in self.eggs:
            current_row = self.eggs.index(egg) // self.eggs_per_row
            rows_after = self.total_rows - (current_row + 1)
            egg_in_row = self.eggs.index(egg) % self.eggs_per_row
            eggs_after = min(len(self.eggs) - (current_row * self.eggs_per_row), self.eggs_per_row) - (egg_in_row + 1)

            x_offset = 32
            y_offset = 30

            # The x coordinate of an egg is determined by which egg in the row it is, and how many eggs
            # are in that row. If there is only 1 egg in a row, it is in the middle of the screen. If
            # there are two, they're on equal halves and so on.
            x = x_offset - (eggs_after * distance_between_eggs) + (egg_in_row * distance_between_eggs)
            y = y_offset - (rows_after * distance_between_rows) + (current_row * distance_between_rows)

            egg.rect.x = x
            egg.rect.y = y

            # Add the egg to the sprite list.
            self._sprites.add(egg)

        self.selected_egg = 0
        self.selected_color = ''
        if preselected_color is not None:
            self.selected_color = self.eggs[self.selected_egg].egg_color
            for i in range(len(self.eggs)):
                if self.eggs[i].egg_color == preselected_color:
                    self.selected_egg = i

    def get_cursor_coords(self):
        """
        Gets the coordinates of an egg on the selection screen by index and returns it as a tuple
        :return: tuple of the coordinates of the selected egg
        """
        cursor_x_offset = -2
        cursor_y_offset = -2

        return (self.eggs[self.selected_egg].rect.x + cursor_x_offset,
                self.eggs[self.selected_egg].rect.y + cursor_y_offset)

    def sel_left(self):
        """
        Select the egg to the left with constraints.
        """

        if self.selected_egg % self.eggs_per_row != 0:
            self.selected_egg -= 1

    def sel_right(self):
        """
        Select the egg to the right with constraints.
        """

        row = self.selected_egg // self.eggs_per_row
        eggs_in_row = min(len(self.eggs) - (row * self.eggs_per_row), self.eggs_per_row)

        if self.selected_egg % self.eggs_per_row != eggs_in_row - 1:
            self.selected_egg += 1

    def sel_up(self):
        """
        Select the egg above with constraints.
        """

        if self.selected_egg // self.eggs_per_row != 0:
            self.selected_egg -= self.eggs_per_row

    def sel_down(self):
        """
        Select the egg below with constraints.
        """

        if self.selected_egg // self.eggs_per_row != self.total_rows - 1:
            self.selected_egg += self.eggs_per_row

    def update(self):
        self.preprocess()

        cursor = pygame.image.load(
            self.resource_dir + '/images/gui/egg_selector.png').convert_alpha()
        self.blit(cursor, self.get_cursor_coords())

        self.selected_color = self.eggs[self.selected_egg].egg_color

        for event in pygame.event.get():
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_RIGHT:
                    self.sel_right()
                if event.key == pygame.K_LEFT:
                    self.sel_left()
                if event.key == pygame.K_DOWN:
                    self.sel_down()
                if event.key == pygame.K_UP:
                    self.sel_up()
                if event.key == pygame.K_a:
                    self.additional_args = {'selected_egg': self.selected_color}
                    self.next_surface = 'selection_info'
                    self.running = False