import sys
from itertools import product

import pygame
from pygame.locals import *
from maze import Maze, Cell
from agent import Agent

pygame.init()
vec = pygame.math.Vector2  # 2 for two dimensional

HEIGHT = 450
WIDTH = 400

FPS = 60

FramePerSec = pygame.time.Clock()

displaysurface = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Game")

SCREEN_BUFFER = 10
MAZE_SIZE = 10
GRID_SIZE = 25


class VisitedSpot(pygame.sprite.Sprite):
    def __init__(self, cell: Cell):
        super().__init__()
        self.surf = pygame.Surface((GRID_SIZE, GRID_SIZE))
        self.surf.fill((0, 122, 000))
        self.rect = self.surf.get_rect(topleft=(cell.x * GRID_SIZE + SCREEN_BUFFER, cell.y * GRID_SIZE + SCREEN_BUFFER))

class DeadEndSpot(pygame.sprite.Sprite):
    def __init__(self, cell: Cell):
        super().__init__()
        self.surf = pygame.Surface((GRID_SIZE, GRID_SIZE))
        self.surf.fill((122, 0, 000))
        self.rect = self.surf.get_rect(topleft=(cell.x * GRID_SIZE + SCREEN_BUFFER, cell.y * GRID_SIZE + SCREEN_BUFFER))

class MazeWall(pygame.sprite.Sprite):
    def __init__(self, wall):
        super().__init__()
        wall_width = 2
        start_cell, end_cell = wall
        x, y = start_cell
        vertical = start_cell[1] == end_cell[1]
        if vertical:
            self.surf = pygame.Surface((wall_width, GRID_SIZE))
        else:
            self.surf = pygame.Surface((GRID_SIZE, wall_width))
        self.surf.fill((0, 0, 255))
        if vertical:
            self.rect = self.surf.get_rect(midtop=((x+1) * GRID_SIZE + SCREEN_BUFFER, y * GRID_SIZE + SCREEN_BUFFER))
        else:
            self.rect = self.surf.get_rect(midleft=(x * GRID_SIZE + SCREEN_BUFFER, (y + 1) * GRID_SIZE + SCREEN_BUFFER))

agents = list()
agents.append(Agent(SCREEN_BUFFER, GRID_SIZE))
# agents.append(Agent(SCREEN_BUFFER, GRID_SIZE))
# agents.append(Agent(SCREEN_BUFFER, GRID_SIZE))
# agents.append(Agent(SCREEN_BUFFER, GRID_SIZE))


all_sprites = pygame.sprite.Group()
visited_cells_sprites = pygame.sprite.Group()
dead_end_sprites = pygame.sprite.Group()
maze_sprites = pygame.sprite.Group()

MOVEEVENT = pygame.USEREVENT+1
RESETMAZE = pygame.USEREVENT+2
reset_maze_event = pygame.event.Event(RESETMAZE, message="Reset Maze")
pygame.event.post(reset_maze_event)
pygame.time.set_timer(MOVEEVENT, 1000)


while True:
    for event in pygame.event.get():
        if event.type == QUIT:
            pygame.quit()
            sys.exit()
        if event.type == MOVEEVENT:
            for agent in agents:
                agent.move(maze)
                if agent.location not in maze.visited:
                    maze.visited.add(agent.location)
                    visited_cells_sprites.add(VisitedSpot(agent.location))
                if agent.location == max(maze.cells):
                    pygame.event.post(reset_maze_event)
                if maze.check_dead_end(agent.location):
                    dead_end_sprites.add(DeadEndSpot(agent.location))
        if event.type == RESETMAZE:
            for agent in agents:
                agent.reset()
            maze = Maze(MAZE_SIZE, MAZE_SIZE)
            maze_sprites.empty()
            visited_cells_sprites.empty()
            dead_end_sprites.empty()
            for wall in maze.walls:
                wall_sprite = MazeWall(wall)
                maze_sprites.add(wall_sprite)
    # wipe buffer
    displaysurface.fill((0, 0, 0))

    from maze import Cell
    location = Cell(0, 5)

    for entity in visited_cells_sprites:
        displaysurface.blit(entity.surf, entity.rect)

    for entity in dead_end_sprites:
        displaysurface.blit(entity.surf, entity.rect)

    for agent in agents:
        pygame.draw.circle(displaysurface, (0, 255, 0), (agent.get_x_position(), agent.get_y_position()), GRID_SIZE / 3)

    for entity in maze_sprites:
        displaysurface.blit(entity.surf, entity.rect)

    # draw border
    pygame.draw.rect(displaysurface, (255, 0, 0),
                     pygame.Rect((SCREEN_BUFFER, SCREEN_BUFFER, MAZE_SIZE * GRID_SIZE, MAZE_SIZE * GRID_SIZE)), 2)

    pygame.display.update()
    FramePerSec.tick(FPS)