"""
-----------------------------------------
Project: Escape Drake
Standard: 91896 (2.7)
School: Tauranga Boys' Collage
Auther: Hunter Moss
Date: May 2024
Python: 3.11.2
-----------------------------------------
"""
import pygame
import os
import pickle
from os import path

pygame.init()

# Screen Size
screen_width = 1000
screen_height = 1000
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption('Escape Drake')

# Load images
background_img = pygame.image.load('Images/Background.png')

# Colours
Red = (255, 0, 0)
Black = (0, 0, 0)

# Constants
tile_size = 50

# Game variables
level = 3

# Load player images function
def load_images(folder, prefix, count):
    return [pygame.transform.scale(pygame.image.load(os.path.join(folder, f"{prefix} ({i}).png")), (tile_size, tile_size)) for i in range(1, count + 1)]

# Player class
class Player(pygame.sprite.Sprite):
    def __init__(self):
        super().__init__()
        self.images = {
            'idle': load_images('Images/Player', 'idle', 16),
            'run': load_images('Images/Player', 'run', 20),
            'jump': load_images('Images/Player', 'jump', 30),
            'dead': load_images('Images/Player', 'dead', 30)
        }
        self.image = self.images['idle'][0]
        self.rect = self.image.get_rect(topleft=(50, screen_height - 50))
        self.vel = pygame.math.Vector2(0, 0)
        self.speed = 5
        self.gravity = 0.5
        self.jump_power = -10
        self.on_ground = False
        self.jumping = False
        self.facing_right = True
        self.frame = 0
        self.idle_animation_speed = 0.2
        self.animation_speed = 0.6
        self.alive = True
        self.keys_collected = 0  # Add key counter

    def reset(self):
        self.rect.topleft = (50, screen_height - 50)
        self.alive = True
        self.keys_collected = 0  # Reset key counter

    def update(self, tiles, door, drakes, keys):
        if self.alive:
            keys_pressed = pygame.key.get_pressed()

            # Horizontal movement
            if keys_pressed[pygame.K_LEFT]:
                self.vel.x = -self.speed
                self.facing_right = False
            elif keys_pressed[pygame.K_RIGHT]:
                self.vel.x = self.speed
                self.facing_right = True
            else:
                self.vel.x = 0

            # Jumping
            if (keys_pressed[pygame.K_SPACE] or keys_pressed[pygame.K_UP]) and self.on_ground:
                self.vel.y = self.jump_power
                self.on_ground = False
                self.jumping = True

            # Apply gravity
            self.vel.y += self.gravity

            # Horizontal collision
            self.rect.x += self.vel.x
            self.check_collision(tiles, 'x')

            # Vertical collision
            self.rect.y += self.vel.y
            self.check_collision(tiles, 'y')

            # Screen collision
            self.rect.clamp_ip(pygame.Rect(0, 0, screen_width, screen_height))

            # Check door collision
            if door and self.rect.colliderect(door.rect):
                door.open = True
                if level < 3:
                    increase_level()
                else:
                    end_game()

            # Check collision with DRAKES
            for drake in drakes:
                if self.rect.colliderect(drake.rect):
                    self.alive = False

            # Check collision with keys
            for key in keys:
                if self.rect.colliderect(key.rect):
                    keys.remove(key)
                    self.keys_collected += 1

            # Animation handling
            if self.jumping:
                self.image = self.images['jump'][0]
            elif self.vel.x != 0:
                self.frame = (self.frame + self.animation_speed) % len(self.images['run'])
                self.image = self.images['run'][int(self.frame)]
            else:
                self.frame = (self.frame + self.idle_animation_speed) % len(self.images['idle'])
                self.image = self.images['idle'][int(self.frame)]

            # Flip sprite
            if not self.facing_right:
                self.image = pygame.transform.flip(self.image, True, False)

            # Reset jumping state if on ground
            if self.rect.bottom >= screen_height:
                self.on_ground = True
                self.jumping = False
                self.vel.y = 0

        else:
            # Play death animation
            self.frame = (self.frame + self.animation_speed) % len(self.images['dead'])
            self.image = self.images['dead'][int(self.frame)]

            # Show retry screen after death animation finishes
            if self.frame >= len(self.images['dead']) - 1:
                show_retry_screen()

    def check_collision(self, tiles, direction):
        for tile in tiles:
            if self.rect.colliderect(tile.rect):
                if direction == 'x':
                    if self.vel.x > 0:  # Moving right
                        self.rect.right = tile.rect.left
                    elif self.vel.x < 0:  # Moving left
                        self.rect.left = tile.rect.right
                elif direction == 'y':
                    if self.vel.y > 0:  # Moving down
                        self.rect.bottom = tile.rect.top
                        self.vel.y = 0
                        self.on_ground = True
                        self.jumping = False
                    elif self.vel.y < 0:  # Moving up
                        self.rect.top = tile.rect.bottom
                        self.vel.y = 0


                # Check if collided with a platform
                if isinstance(tile, Platform) and direction == 'y':
                    if self.vel.y > 0:  # Moving down
                        self.rect.bottom = tile.rect.top
                        self.vel.y = 0
                        self.on_ground = True
                        self.jumping = False
                    elif self.vel.y < 0:  # Moving up
                        self.rect.top = tile.rect.bottom
                        self.vel.y = 0

        # Check collision with platform groups
        platform_groups = pygame.sprite.Group(world.platform_vert_group, world.platform_horz_group)
        platform_collisions = pygame.sprite.spritecollide(self, platform_groups, False)
        for platform in platform_collisions:
            if direction == 'y':
                if self.vel.y > 0:  # Moving down
                    self.rect.bottom = platform.rect.top
                    self.vel.y = 0
                    self.on_ground = True
                    self.jumping = False
                elif self.vel.y < 0:  # Moving up
                    self.rect.top = platform.rect.bottom
                    self.vel.y = 0




# Block Classes
class Block(pygame.sprite.Sprite):
    def __init__(self, x, y, image, scale_to_tile=True):
        super().__init__()
        if scale_to_tile:
            self.image = pygame.transform.scale(image, (tile_size, tile_size))
        else:
            self.image = image
        self.rect = self.image.get_rect(topleft=(x, y))

class Brick0(Block):
    def __init__(self, x, y):
        image = pygame.image.load('Images/Level/brick_dark0.png')
        super().__init__(x, y, image)

class Brick1(Block):
    def __init__(self, x, y):
        image = pygame.image.load('Images/Level/brick_dark1.png')
        super().__init__(x, y, image)

class Brick2(Block):
    def __init__(self, x, y):
        image = pygame.image.load('Images/Level/brick_dark2.png')
        super().__init__(x, y, image)

class Brick3(Block):
    def __init__(self, x, y):
        image = pygame.image.load('Images/Level/brick_dark3.png')
        super().__init__(x, y, image)

class Cobbweb(Block):
    def __init__(self, x, y):
        image = pygame.image.load('Images/Level/Cobbwebs.png')
        super().__init__(x, y, image)

class Grill(Block):
    def __init__(self, x, y):
        image = pygame.image.load('Images/Level/Grill.png')
        super().__init__(x, y, image)

class Skulls(Block):
    def __init__(self, x, y):
        image = pygame.image.load('Images/Level/skulls.png')
        super().__init__(x, y, image)

class Platform(Block):
    def __init__(self, x, y):
        image = pygame.image.load('Images/Level/platform.png')
        original_width, original_height = image.get_size()
        resized_image = pygame.transform.scale(image, (tile_size, original_height + 5))
        adjusted_y = y + tile_size - original_height - 5  # Position at the bottom of the tile
        super().__init__(x, adjusted_y, resized_image, scale_to_tile=False)

class Platform_horz(Block):
    def __init__(self, x, y):
        image = pygame.image.load('Images/Level/platform_horz.png')
        original_width, original_height = image.get_size()
        resized_image = pygame.transform.scale(image, (tile_size, original_height + 5))
        adjusted_y = y + tile_size - original_height - 5  # Position at the bottom of the tile
        super().__init__(x, adjusted_y, resized_image, scale_to_tile=False)
        self.move_speed = 2
        self.range = 150  # Adjust this value to change the platform's movement range
        self.start_x = x  # Remember the starting x position
        self.direction = 1

    def update(self):
        self.rect.x += self.move_speed * self.direction
        if self.direction == 1:
            if self.rect.x - self.start_x >= self.range:
                self.direction = -1
        else:
            if self.start_x - self.rect.x >= self.range:
                self.direction = 1
                


class Platform_vert(Block):
    def __init__(self, x, y):
        image = pygame.image.load('Images/Level/platform_vert.png')
        original_width, original_height = image.get_size()
        resized_image = pygame.transform.scale(image, (tile_size, original_height + 5))
        adjusted_y = y + tile_size - original_height - 5  # Position at the bottom of the tile
        super().__init__(x, adjusted_y, resized_image, scale_to_tile=False)
        self.move_speed = 2
        self.range = 150  # Adjust this value to change the platform's movement range
        self.start_y = y  # Remember the starting y position
        self.direction = 1

    def update(self):
        self.rect.y += self.move_speed * self.direction
        if self.direction == 1:
            if self.rect.y - self.start_y >= self.range:
                self.direction = -1
        else:
            if self.start_y - self.rect.y >= self.range:
                self.direction = 1
                


class Trog(Block):
    def __init__(self, x, y):
        image = pygame.image.load('Images/dngn_altar_trog.png')
        super().__init__(x, y, image)

class Key(Block):
    def __init__(self, x, y):
        image = pygame.image.load('Images/Key.png')
        resized_image = pygame.transform.scale(image, (50, 20))  # Resize key image to 50x20 pixels
        # Calculate adjusted_y to center the key vertically within the grid
        adjusted_y = y + (tile_size - resized_image.get_height()) // 2
        super().__init__(x, adjusted_y, resized_image, scale_to_tile=False)

class Drake(Block):
    def __init__(self, x, y):
        image = pygame.image.load('Images/DRAKE.png')
        image = pygame.transform.scale(image, (tile_size, tile_size))
        super().__init__(x, y, image)

class Door(Block):
    def __init__(self, x, y):
        self.closed_image = pygame.image.load('Images/dngn_closed_door.png')
        self.open_image = pygame.image.load('Images/dngn_open_door.png')
        # Resize images to match the tile size
        self.closed_image = pygame.transform.scale(self.closed_image, (tile_size, tile_size))
        self.open_image = pygame.transform.scale(self.open_image, (tile_size, tile_size))
        super().__init__(x, y, self.closed_image)
        self.open = False

    def update(self):
        if self.open:
            self.image = self.open_image
        else:
            self.image = self.closed_image

# World class
class World:
    def __init__(self, data):
        self.tile_list = []
        self.floor_group = pygame.sprite.Group()  # Group for floor tiles
        self.drake_group = pygame.sprite.Group()  # Group for drakes
        self.key_group = pygame.sprite.Group()  # Group for keys
        self.platform_vert_group = pygame.sprite.Group()  # Group for vertical platforms
        self.platform_horz_group = pygame.sprite.Group()  # Group for horizontal platforms
        self.door_group = pygame.sprite.GroupSingle()  # Group for the door

        self.load_world(data)

    def load_world(self, data):
        block_types = {
            'brick0': Brick0,
            'brick1': Brick1,
            'brick2': Brick2,
            'brick3': Brick3,
            'cobbweb': Cobbweb,
            'grill': Grill,
            'skulls': Skulls,
            'platform': Platform,
            'platform_vert': Platform_vert,
            'platform_horz': Platform_horz,
            'trog': Trog,
            'key': Key,
            'drake': Drake,
            'door': Door
        }
        
        for block in data:
            x, y, block_type = block
            if block_type in block_types:
                block_class = block_types[block_type]
                block_instance = block_class(x, y)
                if block_type == 'door':
                    self.door_group.add(block_instance)
                elif block_type == 'drake':
                    self.drake_group.add(block_instance)
                elif block_type == 'key':
                    self.key_group.add(block_instance)
                elif block_type == 'platform_vert':
                    self.platform_vert_group.add(block_instance)
                elif block_type == 'platform_horz':
                    self.platform_horz_group.add(block_instance)
                else:
                    self.tile_list.append(block_instance)
                    self.floor_group.add(block_instance)

    def draw(self):
        for tile in self.tile_list:
            screen.blit(tile.image, tile.rect)
        self.door_group.draw(screen)
        self.drake_group.draw(screen)
        self.key_group.draw(screen)
        self.platform_vert_group.draw(screen)
        self.platform_horz_group.draw(screen)

    def get_tiles(self):
        return self.floor_group

    def get_keys(self):
        return self.key_group

def increase_level():
    global world
    global player
    global level
    level += 1
    print(f"Loading level {level}")
    if path.exists(f'level_data/level{level}_data'):
        with open(f'level_data/level{level}_data', 'rb') as pickle_in:
            world_data = pickle.load(pickle_in)
        player.reset()  # Reset player state
        world = World(world_data)  # Re-initialize the world
    else:
        end_game()

def reset_level():
    global world
    global player
    global level
    print(f"Resetting level {level}")
    if path.exists(f'level_data/level{level}_data'):
        with open(f'level_data/level{level}_data', 'rb') as pickle_in:
            world_data = pickle.load(pickle_in)
        world = World(world_data)
    else:
        world = World([])  # Default to an empty world if level data is not found
    player.reset()  # Reset player state

def end_game():
    print("Game Over")

def show_retry_screen():
    global run, retry_screen

    font = pygame.font.Font(None, 75)
    text = font.render('You failed to Escape Drake!', True, (255, 255, 255))
    text_rect = text.get_rect(center=(screen_width // 2, screen_height // 2))

    retry_button_rect = pygame.Rect(0, 0, 400, 50)
    retry_button_rect.center = (screen_width // 2, screen_height // 2 + 100)

    retry_screen = True
    while retry_screen:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                quit()
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_SPACE:  # Press 'Space' to retry
                    reset_level()
                    retry_screen = False

        screen.fill((0, 0, 0))
        screen.blit(text, text_rect)
        pygame.draw.rect(screen, (255, 0, 0), retry_button_rect)
        
        font_small = pygame.font.Font(None, 50)
        retry_text = font_small.render('Press Space to Retry', True, (255, 255, 255))
        retry_text_rect = retry_text.get_rect(center=retry_button_rect.center)
        screen.blit(retry_text, retry_text_rect)

        pygame.display.flip()

def main():
    global world
    global player
    global level
    global run

    run = True
    clock = pygame.time.Clock()
    player = Player()
    player_group = pygame.sprite.Group()
    player_group.add(player)

    # Load the first level
    reset_level()

    while run:
        screen.blit(background_img, (0, 0))

        # Event handling
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                run = False

        player.update(world.get_tiles(), world.door_group.sprite, world.drake_group, world.get_keys())
        if world.door_group.sprite:
            world.door_group.update()

        # Update moving platforms
        for platform in world.platform_vert_group:
            platform.update()
        for platform in world.platform_horz_group:
            platform.update()

        # Update drakes
        world.drake_group.update()

        world.draw()
        player_group.draw(screen)

        # Display current level
        font = pygame.font.Font(None, 30)
        level_text = font.render(f'Level: {level}', True, (255, 255, 255))
        screen.blit(level_text, (10, 10))

        # Display key counter
        key_text = font.render(f'Keys: {player.keys_collected}', True, (255, 255, 255))
        screen.blit(key_text, (10, 40))

        pygame.display.update()
        clock.tick(60)

    pygame.quit()

if __name__ == "__main__":
    main()