from random import randint, choice
from games.snake.apple import Apple
from games.snake.snake import SnakeBodySegment
from retro.game import Game


class SnakeHead:
    RIGHT = (1, 0)
    UP = (0, -1)
    LEFT = (-1, 0)
    DOWN = (0, 1)
    name = "Snake head"
    position = (0, 0)
    direction = DOWN
    character = 'v'
    next_segment = None
    growing = False

    def play_turn(self, game):
        x, y = self.position
        dx, dy = self.direction
        next_pos = (x + dx, y + dy)
        if self.can_move(next_pos, game):
            apple = game.get_agent_by_name("Apple")
            ax, ay = apple.position
            old_dist = abs(x - ax) + abs(y - ay)
            new_dist = abs(next_pos[0] - ax) + abs(next_pos[1] - ay)
            game.state['reward'] += old_dist - new_dist
            game.state['energy'] -= 1
            self.position = next_pos
            if self.is_on_apple(self.position, game):
                apple.relocate(game)
                self.growing = True
                game.state['score'] += 50
                game.state['reward'] += 50
                game.state['energy'] = 150
            if self.next_segment:
                self.next_segment.move((x, y), game, growing=self.growing)
            elif self.growing:
                self.next_segment = SnakeBodySegment(1, (x, y))
                game.add_agent(self.next_segment)
            self.growing = False
            if game.state['energy'] <= 0:
                game.state['reward'] -= 10
                game.end()
        else:
            game.state['reward'] -= 10
            game.end()

    def handle_keystroke(self, keystroke, game):
        if keystroke.name == "KEY_RIGHT":
            self.direction = self.RIGHT
            self.character = '>'
        elif keystroke.name == "KEY_UP":
            self.direction = self.UP
            self.character = '^'
        elif keystroke.name == "KEY_LEFT":
            self.direction = self.LEFT
            self.character = '<'
        elif keystroke.name == "KEY_DOWN":
            self.direction = self.DOWN
            self.character = 'v'

    def can_move(self, position, game):
        return game.on_board(position) and (game.is_empty(position) or self.is_on_apple(position, game))

    def is_on_apple(self, position, game):
        return game.get_agent_by_name("Apple").position == position


def create_game():
    head = SnakeHead()
    apple = Apple()
    game = Game(
        [head, apple],
        {'score': 0, 'reward': 0, 'energy': 150},
        board_size=(32, 16),
        framerate=12,
        show_state=['score'],
    )
    bw, bh = game.board_size
    head.position = (randint(1, bw - 2), randint(1, bh - 2))
    direction, character = choice([
        (SnakeHead.RIGHT, '>'),
        (SnakeHead.UP,    '^'),
        (SnakeHead.LEFT,  '<'),
        (SnakeHead.DOWN,  'v'),
    ])
    head.direction = direction
    head.character = character
    apple.relocate(game)
    return game