import pygame
import sys
import random

pygame.init()

width, height = 800, 800
screen = pygame.display.set_mode((width, height))
pygame.display.set_caption('Labirints')

BLACK = (0, 0, 0)
BLUE = (0, 0, 255)
RED = (255, 0, 0)

maze = [
    "#########################",
    "#......#..........#.....#",
    "######.########.#.#.#####",
    "#...##.#........#.#.....#",
    "#.#.##...########.#####.#",
    "#.#.#########.....#...#.#",
    "#.#.......#.#.#####.#.#.#",
    "#.##..###.#.#.#.....#...#",
    "#.....###.#.#.#.#######.#",
    "#.##....#...#...#.#.....#",
    "#.##.##.###.#####.#.##.##",
    "#.....#...#...#...#.#...#",
    "#.#.#.###.###.#.###.#.#.#",
    "#.#.#...........#...#.#.#",
    "#.#.##.###.####.#.###.#.#",
    "#.#.#...##.##...#.#...#.#",
    "#.#.#.#.##....###.#.###.#",
    "#.#.#.#....##.#...#.#...#",
    "#.#.#.#.##..###.###.#.###",
    "#.......................#",
    "#########################"
]

class Player(pygame.sprite.Sprite):
    def __init__(self, x, y):
        super().__init__()
        self.images = {
            "up": [pygame.image.load('up.png').convert_alpha(),
                   pygame.image.load('up1.png').convert_alpha()],
            "down": [pygame.image.load('down.png').convert_alpha(),
                     pygame.image.load('down1.png').convert_alpha()],
            "left": [pygame.image.load('left.png').convert_alpha(),
                     pygame.image.load('left1.png').convert_alpha()],
            "right": [pygame.image.load('right.png').convert_alpha(),
                      pygame.image.load('right1.png').convert_alpha()]
        }
        self.direction = "down"  # Initial direction
        self.image_index = 0  # Index of the current image in the animation
        self.image = self.images[self.direction][self.image_index]
        self.rect = self.image.get_rect()
        self.rect.topleft = (x, y)
        self.speed = 10
        self.cheese_count = 0  # Cheese count

    def move(self, dx, dy, walls):
        old_rect = self.rect.copy()
        self.rect.x += dx
        self.rect.y += dy
        for wall in walls:
            if self.rect.colliderect(wall):
                if dx > 0:
                    self.rect.right = wall.left
                elif dx < 0:
                    self.rect.left = wall.right
                elif dy > 0:
                    self.rect.bottom = wall.top
                elif dy < 0:
                    self.rect.top = wall.bottom
        if self.rect.collidelist(walls) != -1:
            self.rect = old_rect

    def update_image(self, dx, dy):
        if dx > 0:
            self.direction = "right"
        elif dx < 0:
            self.direction = "left"
        elif dy > 0:
            self.direction = "down"
        elif dy < 0:
            self.direction = "up"

        # Update the image index for animation
        self.image_index = (self.image_index + 1) % len(self.images[self.direction])
        self.image = self.images[self.direction][self.image_index]

    def draw(self, surface):
        surface.blit(self.image, self.rect)

class Cheese(pygame.sprite.Sprite):
    def __init__(self, x, y):
        super().__init__()
        self.image = pygame.image.load('cheese.png').convert_alpha()
        self.rect = self.image.get_rect()
        self.rect.topleft = (x, y)

class Villain(pygame.sprite.Sprite):
    def __init__(self, x, y, player):
        super().__init__()
        self.image = pygame.image.load('villain.png').convert_alpha()
        self.rect = self.image.get_rect()
        self.rect.topleft = (x, y)
        self.speed = 8
        self.player = player

    def move_towards_player(self, walls):
        dx = self.player.rect.x - self.rect.x
        dy = self.player.rect.y - self.rect.y
        distance = max(abs(dx), abs(dy))
        dx = dx / distance if distance != 0 else 0
        dy = dy / distance if distance != 0 else 0
        dx *= self.speed
        dy *= self.speed
        self.move(dx, dy, walls)

    def move(self, dx, dy, walls):
        old_rect = self.rect.copy()
        self.rect.x += dx
        self.rect.y += dy
        for wall in walls:
            if self.rect.colliderect(wall):
                if dx > 0:
                    self.rect.right = wall.left
                elif dx < 0:
                    self.rect.left = wall.right
                elif dy > 0:
                    self.rect.bottom = wall.top
                elif dy < 0:
                    self.rect.top = wall.bottom
        if self.rect.collidelist(walls) != -1:
            self.rect = old_rect

    def draw(self, surface):
        surface.blit(self.image, self.rect)

def draw_maze(maze):
    walls = []
    roads = []  # 存储道路的矩形列表
    for y, row in enumerate(maze):
        for x, char in enumerate(row):
            if char == '#':
                pygame.draw.rect(screen, BLACK, (x * 32, y * 32, 32, 32))
                walls.append(pygame.Rect(x * 32, y * 32, 32, 32))
            elif char == '.':
                roads.append(pygame.Rect(x * 32, y * 32, 32, 32))
            elif char == 'C':  # Cheese
                roads.append(pygame.Rect(x * 32, y * 32, 32, 32))
                screen.blit(pygame.image.load('cheese.png'), (x * 32, y * 32))
            elif char == 'V':  # Villain
                roads.append(pygame.Rect(x * 32, y * 32, 32, 32))
                screen.blit(pygame.image.load('villain.png'), (x * 32, y * 32))
    return walls, roads

def draw_menu():
    font = pygame.font.Font(None, 36)
    title_text = font.render("Labirints", True, RED)
    start_text = font.render("Press 1 to Control Player", True, BLACK)
    start_text2 = font.render("Press 2 to Control Villain", True, BLACK)
    quit_text = font.render("Press Q to Quit", True, BLACK)

    screen.blit(title_text, (width // 2 - title_text.get_width() // 2, height // 2 - 100))
    screen.blit(start_text, (width // 2 - start_text.get_width() // 2, height // 2 - 50))
    screen.blit(start_text2, (width // 2 - start_text2.get_width() // 2, height // 2))
    screen.blit(quit_text, (width // 2 - quit_text.get_width() // 2, height // 2 + 50))

    pygame.display.flip()

player = Player(32, 32)
cheese = Cheese(416, 416)  # Cheese position
villain = Villain(704, 704, player)  # Villain's starting position
villain2 = Villain(32, 32, player)   # Second villain's starting position

clock = pygame.time.Clock()
running = True
show_menu = True
control_player = True  # 初始设置玩家控制主角
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_q:
                running = False
            if event.key == pygame.K_1:
                control_player = True
                show_menu = False
            if event.key == pygame.K_2:
                control_player = False
                show_menu = False

    if show_menu:
        screen.fill(BLUE)
        draw_menu()
    else:
        if control_player:
            keys = pygame.key.get_pressed()
            dx, dy = 0, 0

            if keys[pygame.K_LEFT]:
                dx = -player.speed
            elif keys[pygame.K_RIGHT]:
                dx = player.speed

            if keys[pygame.K_UP]:
                dy = -player.speed
            elif keys[pygame.K_DOWN]:
                dy = player.speed

            if dx != 0 and dy != 0:
                if keys[pygame.K_LEFT]:
                    dy = 0
                elif keys[pygame.K_RIGHT]:
                    dy = 0
                elif keys[pygame.K_UP]:
                    dx = 0
                elif keys[pygame.K_DOWN]:
                    dx = 0
            
            walls, roads = draw_maze(maze)
            player.move(dx, dy, walls)
            player.update_image(dx, dy)

            villain.move_towards_player(walls)

            if player.rect.colliderect(villain.rect):
                print("Game Over! You were caught by the villain.")
                running = False

            if player.rect.colliderect(cheese.rect):
                player.cheese_count += 1
                print("You collected a cheese! Total cheese count:", player.cheese_count)
                cheese.rect.topleft = random.choice(roads).topleft  # Move cheese to a random position on a road
                if player.cheese_count >= 10:
                    print("Congratulations! You collected 10 cheeses!")
                    running = False

            screen.fill(BLUE)
            draw_maze(maze)
            player.draw(screen)
            villain.draw(screen)
            pygame.display.flip()
        else:
            keys = pygame.key.get_pressed()
            dx, dy = 0, 0

            if keys[pygame.K_a]:
                dx = -villain2.speed
            elif keys[pygame.K_d]:
                dx = villain2.speed

            if keys[pygame.K_w]:
                dy = -villain2.speed
            elif keys[pygame.K_s]:
                dy = villain2.speed

            if dx != 0 and dy != 0:
                if keys[pygame.K_a]:
                    dy = 0
                elif keys[pygame.K_d]:
                    dy = 0
                elif keys[pygame.K_w]:
                    dx = 0
                elif keys[pygame.K_s]:
                    dx = 0
            
            walls, roads = draw_maze(maze)
            villain2.move(dx, dy, walls)

            if villain2.rect.colliderect(player.rect):
                print("Game Over! The second villain caught you.")
                running = False

            screen.fill(BLUE)
            draw_maze(maze)
            player.draw(screen)
            villain.draw(screen)
            villain2.draw(screen)
            pygame.display.flip()

    clock.tick(10)

pygame.quit()
sys.exit()