import pygame, time
from random import randint, shuffle
from copy import deepcopy
from uuid import uuid1 as ranuuid


scale = 10
width, height = 64, 64
scoreboard_width, scoreboard_border = 320, 6
pygame.init()
win = pygame.display.set_mode((width * 10 + scoreboard_width, height * 10))
pygame.display.set_caption("Snek")

class Snake:
    def __init__(self, x, y, color=None, controller=None, name=None):
        self.body = [(x, y)]
        self.direction = randint(0, 3)
        self.alive = True
        self.length = 3
        self.uuid = ranuuid()
        self.name = name

        if color is None:
            color = [(255, 0, 0)]
        self.color = color

        self.is_player = controller is None
        self.controller = controller

        self.ticks = 0
        self.kills = 0.0

    def tick(self):
        if len(self.body) == 0:
            return
        last_x, last_y = self.body[-1]
        # Dead
        if not self.alive:
            self.body.pop(0)
        # Right
        elif self.direction == 0:
            self.body.append((last_x + 1, last_y))
        # Up
        elif self.direction == 1:
            self.body.append((last_x, last_y - 1))
        # Left
        elif self.direction == 2:
            self.body.append((last_x - 1, last_y))
        # Down
        elif self.direction == 3:
            self.body.append((last_x, last_y + 1))
        else:
            print("Direction does not exist.")
        # Removes the oldest segments until it has reached the proper length
        while self.length < len(self.body):
            self.body.pop(0)
        self.ticks += 1

    def collision(self, snakes, apples):
        if len(self.body) == 0:
            return
        x, y = head = self.body[-1]
        # Collision with our own body
        if head in self.body[0:-1]:
            self.body.pop(len(self.body) - 1)
            self.alive = False
        # Collision with other snakes
        closest_snakes, closest_snake_dist = [], 0
        for snake in snakes:
            if snake != self:
                if head in snake.body:
                    self.body.pop(len(self.body) - 1)
                    self.alive = False
                    snake.kills += 1.0
                    snake.length += self.length // 2
                # Finds the closest snake by each segment to tell who killed it.
                for body in snake.body:
                    dist = abs(body[0] - head[0]) + abs(body[1] - head[1])
                    if closest_snake_dist == 0 or dist < closest_snake_dist:
                        closest_snakes = [snake]
                        closest_snake_dist = dist
                    elif dist == closest_snake_dist:
                        closest_snakes.append(snake)
        # Collisions with food
        for food in apples:
            if head == food:
                apples.remove(food)
                self.length += 1
        # Outside the map
        # TODO: Detect when snakes run other snakes off the map or trap them.
        if not (0 <= x < width and 0 <= y < height):
            print("Dist to closeset:", closest_snake_dist)
            if closest_snake_dist <= 2:
                for snake in closest_snakes:
                    snake.kills += 1 / len(closest_snakes)
                    snake.length += (self.length / 2) // len(closest_snakes)
            self.body.pop(len(self.body) - 1)
            self.alive = False

    def set_direction(self, direction):
        if direction % 2 != self.direction % 2:
            self.direction = direction % 4
            return True
        return False

    def get_head(self):
        return self.body[-1]

    def score(self):
        return int(self.kills / 3 * self.length)

def default_think(snake: Snake, snakes: list, apples: list):
    if not snake.alive:
        return snake.direction
    x, y = head = snake.get_head()
    dir = snake.direction

    filled_tiles = [] + snake.body
    for snek in snakes:
        if snek != snake and snek.uuid != snake.uuid:
            filled_tiles += snek.body
            if snek.alive:
                head_x, head_y = snek.get_head()
                if snek.direction == 0:
                    filled_tiles.append((head_x + 1, head_y))
                elif snek.direction == 1:
                    filled_tiles.append((head_x, head_y - 1))
                elif snek.direction == 2:
                    filled_tiles.append((head_x - 1, head_y))
                elif snek.direction == 3:
                    filled_tiles.append((head_x, head_y + 1))

    for num in range(-1, 65):
        filled_tiles += [(-1, num), (64, num), (num, -1), (num, 64)]

    # for spot in filled_tiles:
    #     pygame.draw.rect(win, (255, 0, 0), (spot[0] * scale, spot[1] * scale, scale, scale))

    # The four possible directions
    head_up, head_down = (x, y - 1), (x, y + 1)
    head_left, head_right = (x - 1, y), (x + 1, y)

    if dir == 0:
        x += 1
    if dir == 1:
        y -= 1
    if dir == 2:
        x -= 1
    if dir == 3:
        y += 1

    if (x, y) in filled_tiles or ((x == width // 2 or y == height // 2) and randint(1, 2) == 1):
        possible_directions = []
        if head_up not in filled_tiles:
            possible_directions.append(1)
        if head_down not in filled_tiles:
            possible_directions.append(3)
        if head_left not in filled_tiles:
            possible_directions.append(2)
        if head_right not in filled_tiles:
            possible_directions.append(0)

        shuffle(possible_directions)
        if len(possible_directions):
            return possible_directions[0]

    return dir

def ben_think(snake: Snake, snakes: list, apples: list):
    # Hahahahaa you THOUGHT you would get to see my code!
    return snake.direction

def draw_tile(x, y, color, translate=True):
    if translate:
        pygame.draw.rect(win, color, (x * scale + 1 + scoreboard_width, y * scale + 1, scale - 2, scale - 2))
    else:
        pygame.draw.rect(win, color, (x * scale + 1, y * scale + 1, scale - 2, scale - 2))

def ran_color(min_avg=0):
    r, g, b = randint(0, 255), randint(0, 255), randint(0, 255)
    while r + g + b / 3 < min_avg:
        r, g, b = randint(0, 255), randint(0, 255), randint(0, 255)
    return r, g, b

def draw_text(win, text, c_x, c_y, size=10, color=(0, 0, 0), x_left=False, y_up=False):
    text = str(text)
    basic_font = pygame.font.SysFont(None, size)
    text = basic_font.render(text, True, color)
    text_rect = text.get_rect()
    text_rect.centerx = c_x
    text_rect.centery = c_y
    if x_left:
        text_rect.x = c_x
    if y_up:
        text_rect.y = c_y
    win.blit(text, text_rect)

def main():
    ticks = 0
    snakes, dead_snakes, foods = [], [], []
    def random_open_spot():
        filled = [] + [bod for bod in [snake.body for snake in snakes]] + foods
        ran = (randint(0, 64), randint(0, 64))
        while ran in filled:
            ran = (randint(0, width - 1), randint(0, width - 1))
        return ran

    # snakes.append(Snake(32, 32, [(100, 0, 200), (0, 0, 200)]))
    for i in range(25):
        ran = random_open_spot()
        snakes.append(Snake(ran[0], ran[1], [ran_color(150) for _ in range(0, randint(0, 2) + 1)], default_think))

    max_foods = 1
    movement_events = []
    clock = pygame.time.Clock()
    running, paused = True, False

    tps = 10
    # Game loop
    while running:
        keys = pygame.key.get_pressed()
        # Events
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    running = False
                if event.key == pygame.K_r:
                    return True
                if event.key == pygame.K_c:
                    foods.clear()
                if event.key == pygame.K_p:
                    paused = not paused
                if not paused and event.key in [pygame.K_w, pygame.K_a, pygame.K_s, pygame.K_d,
                                 pygame.K_LEFT, pygame.K_RIGHT, pygame.K_UP, pygame.K_DOWN]:
                    movement_events.append(event)

        # Movement
        if not paused:
            if len(movement_events) > 0:
                event = movement_events.pop(0)
                key = event.key
                player_controlled = list(filter(lambda x: x.is_player, snakes))
                for snake in player_controlled:
                    if key == pygame.K_d or key == pygame.K_RIGHT:
                        snake.set_direction(0)
                    elif key == pygame.K_w or key == pygame.K_UP:
                        snake.set_direction(1)
                    elif key == pygame.K_a or key == pygame.K_LEFT:
                        snake.set_direction(2)
                    elif key == pygame.K_s or key == pygame.K_DOWN:
                        snake.set_direction(3)
            last_snakes = deepcopy(snakes)
            for snake in snakes:
                if snake.controller is not None:
                    snake.set_direction(snake.controller(snake, last_snakes, foods))
                snake.tick()

            for snake in snakes:
                if not snake.alive and len(snake.body) == 0:
                    snakes.remove(snake)
                    dead_snakes.insert(0, snake)
                snake.collision(snakes, foods)

            while len(foods) < max_foods:
                foods.append(random_open_spot())

        win.fill((50, 50, 50))
        # Game Drawing
        for snake in snakes:
            body, color = snake.body, snake.color
            for body_index, segment in enumerate(list(reversed(body))):
                draw_tile(segment[0], segment[1], color[body_index % len(color)])
        for food in foods:
            food_x, food_y = food
            draw_tile(food_x, food_y, (222, 0, 0))
        # Scoreboard Drawing
        # TODO: Make it so if the mouse is hovering over a snake on the scoreboard, the snake gets highlighted
        # Do this by filling the squares behind it with a rainbow color (shifting through the rainbow)
        pygame.draw.rect(win, (200, 200, 200), (0, 0, scoreboard_width, height * scale))
        pygame.draw.rect(win, (50, 50, 50), (scoreboard_border // 2, scoreboard_border // 2,
                                             scoreboard_width - scoreboard_border, height * scale - scoreboard_border))

        def draw_scoreboard(sneks):
            index, ret = 0, 0
            for snake in sneks:
                x, y = 0, (2 * index) + 9
                for i in range(4):
                    x = 1 + i
                    draw_tile(x, y, snake.color[i % len(snake.color)], translate=False)
                text_color = (255, 255, 255)
                if not snake.alive:
                    text_color = (255, 0, 0)
                draw_text(win, "Kills: " + str(snake.kills) + " Length: " + str(snake.length), x * scale + 15, (y + .5) * scale, color=text_color, size=30, x_left=True)
                ret = (y + .5) * scale
                index += 1
            return ret
        sorted_snakes = list(sorted(snakes, key=lambda s: s.score(), reverse=True))
        draw_scoreboard(sorted_snakes + dead_snakes)


        # Updates the window
        pygame.display.flip()

        # Food amount management
        if not paused:
            ticks += 1
            if ticks % 100 == 0:
                max_foods += 1
                # Makes the snakes longer
                for snake in snakes:
                    if snake.alive:
                        snake.length += 1
        # Speed up the game
        if keys[pygame.K_SPACE]:
            clock.tick(tps * 5)
        else:
            clock.tick(tps)


if __name__ == '__main__':
    while True:
        if not main():
            break