import random
import pygame

# pygame setup
pygame.init()
screen = pygame.display.set_mode((1280, 720))
clock = pygame.time.Clock()
running = True
dt = 0

force_coef = 20
player_applied_force = pygame.Vector2(0, 0)
player_mass = 1


# Helper to find all empty cells in the maze
def get_empty_cells(grid):
    empty = []
    for y, row in enumerate(grid):
        for x, cell in enumerate(row):
            if cell == 0:
                empty.append((x, y))
    return empty

# Set player_pos to a random empty cell
def get_cell_center(x, y, cell_size=80):
    return pygame.Vector2(x * cell_size + cell_size / 2, y * cell_size + cell_size / 2)

from mazelib import Maze
from mazelib.generate.Prims import Prims

maze_width = 7
maze_height = 4

m = Maze()
m.generator = Prims(maze_height, maze_width)
m.generate()
maze_grid = m.grid

empty_cells = get_empty_cells(maze_grid)
spawn_x, spawn_y = random.choice(empty_cells)
player_pos = get_cell_center(spawn_x, spawn_y)
player_speed = pygame.Vector2(0, 0)

raketa_original = pygame.image.load("raketa.png").convert_alpha()
raketa = pygame.transform.scale(raketa_original, (raketa_original.get_width() // 3, raketa_original.get_height() // 3))

# Function to get all meteor rectangles for collision detection
def get_meteor_rects(grid, cell_size=80):
    meteor_rects = []
    for y, row in enumerate(grid):
        for x, cell in enumerate(row):
            if cell == 1:
                rect = pygame.Rect(x * cell_size, y * cell_size, cell_size, cell_size)
                meteor_rects.append(rect)
    return meteor_rects

from mazelib import Maze
from mazelib.generate.Prims import Prims

maze_width = 7
maze_height = 4

m = Maze()
m.generator = Prims(maze_height, maze_width)
m.generate()
maze_grid = m.grid


def draw_maze(
    surface,
    grid,
    meteor_img=pygame.image.load("meteor.png").convert_alpha(),
    cell_size=80,
):
    for y, row in enumerate(grid):
        for x, cell in enumerate(row):
            if cell == 1:
                meteor_img = pygame.transform.scale(meteor_img, (cell_size, cell_size))
                rect = pygame.Rect(x * cell_size, y * cell_size, cell_size, cell_size)
                surface.blit(meteor_img, rect)


def draw_arrow(surface, color, start, end, width=3):
    pygame.draw.line(surface, color, start, end, width)
    direction = (end - start).normalize()
    arrow_size = 10
    perpendicular = pygame.Vector2(-direction.y, direction.x)
    point1 = end - direction * arrow_size + perpendicular * arrow_size / 2
    point2 = end - direction * arrow_size - perpendicular * arrow_size / 2
    pygame.draw.polygon(surface, color, [end, point1, point2])

def rescale_rocket(raketa, speed):
    scale_factor = 1 + player_speed.length() * 0.01
    raketa_w, raketa_h = raketa.get_size()
    raketa_w *= scale_factor
    raketa_h /= scale_factor
    scaled_raketa = pygame.transform.scale(raketa, (raketa_h, raketa_w))
    return scaled_raketa

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    screen.fill("purple")

    angle = player_speed.angle_to(pygame.Vector2(0, -1))

    scaled_raketa = rescale_rocket(raketa, player_speed)
    rotated_raketa = pygame.transform.rotate(scaled_raketa, angle)
    raketa_rect = rotated_raketa.get_rect(center=player_pos)

    player_applied_force = pygame.Vector2(0, 0)

    keys = pygame.key.get_pressed()
    if keys[pygame.K_w]:
        player_applied_force.y -= force_coef
    if keys[pygame.K_s]:
        player_applied_force.y += force_coef
    if keys[pygame.K_a]:
        player_applied_force.x -= force_coef
    if keys[pygame.K_d]:
        player_applied_force.x += force_coef

    if player_applied_force.length() > 0:
        player_applied_force.scale_to_length(force_coef)

    acceleration = player_applied_force / player_mass
    player_speed += acceleration * dt
    player_pos += player_speed * dt

    # Collision detection
    meteor_rects = get_meteor_rects(maze_grid)
    raketa_rect = rotated_raketa.get_rect(center=player_pos)
    for meteor_rect in meteor_rects:
        if raketa_rect.colliderect(meteor_rect):
            # Collision detected: reset position and speed
            player_speed = pygame.Vector2(0, 0)
            break

    # Draw rocket
    raketa_rect = rotated_raketa.get_rect(center=player_pos)
    screen.blit(rotated_raketa, raketa_rect)

    # Draw force arrow
    if player_applied_force.length() > 0:
        arrow_end = player_pos + player_applied_force * 2
        draw_arrow(screen, "yellow", player_pos, arrow_end)

    # Draw speed arrow
    if player_speed.length() > 0:
        arrow_end = player_pos + player_speed * 2
        draw_arrow(screen, "cyan", player_pos, arrow_end)

    draw_maze(screen, maze_grid)

    pygame.display.flip()
    dt = clock.tick(60) / 1000

pygame.quit()