"""
---------------------------------------------
    Project: Mario 
    Standard: NCEA Level 2
    School: Tauranga Boys College
    Author: Harry Burt
    Date: June 2025
    Python: 3.11.9
---------------------------------------------
"""
# Setup
import pygame
from pygame.locals import *
from pygame import mixer
import pickle
from os import path

# Pre-initialize mixer for better sound timing
pygame.mixer.pre_init(44100, -16, 2, 512)
mixer.init()
pygame.init()

clock = pygame.time.Clock()
fps = 60

# Screen setup
screen_width = 1000
screen_height = 1000
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption('Platformer')

# Fonts
font = pygame.font.SysFont('Bauhaus 93', 70)
font_score = pygame.font.SysFont('Bauhaus 93', 30)

# Game variables
tile_size = 50
game_over = 0
main_menu = True
level = 1
max_levels = 7
score = 0
lives = 3

# Colors
white = (255, 255, 255)
blue = (0, 0, 255)

# Load images
sun_img = pygame.image.load('img/sun.png')
bg_img = pygame.transform.scale(pygame.image.load('img/sky.jpg'), (screen_width, screen_height))
restart_img = pygame.image.load('img/restart_btn.png')
start_img = pygame.image.load('img/start_btn.png')
exit_img = pygame.image.load('img/exit_btn.png')
logo_img = pygame.transform.scale(pygame.image.load('img/logo.png'), (500, 250))

# Load sounds and music
pygame.mixer.music.load('img/music.mp3')
pygame.mixer.music.play(-1, 0.0, 5000)
coin_fx = pygame.mixer.Sound('img/coin.wav')
coin_fx.set_volume(0.5)
jump_fx = pygame.mixer.Sound('img/jump.wav')
jump_fx.set_volume(0.5)
game_over_fx = pygame.mixer.Sound('img/game_over.wav')
game_over_fx.set_volume(0.5)

# Utility Functions
def draw_text(text, font, text_col, x, y):
    img = font.render(text, True, text_col)
    screen.blit(img, (x, y))

# Reset the level and reload all entities
def reset_level(level):
    player.reset(100, screen_height - 130)
    blob_group.empty()
    platform_group.empty()
    coin_group.empty()
    lava_group.empty()
    exit_group.empty()

    if path.exists(f'level{level}_data'):
        pickle_in = open(f'level{level}_data', 'rb')
        world_data = pickle.load(pickle_in)
    world = World(world_data, level)
    score_coin = Coin(tile_size // 2, tile_size // 2)
    coin_group.add(score_coin)
    return world

# UI Button Class
class Button():
    def __init__(self, x, y, image):
        self.image = image
        self.rect = self.image.get_rect()
        self.rect.x = x
        self.rect.y = y
        self.clicked = False

    # Returns True if the button is clicked
    def draw(self):
        action = False
        pos = pygame.mouse.get_pos()
        if self.rect.collidepoint(pos):
            if pygame.mouse.get_pressed()[0] == 1 and not self.clicked:
                action = True
                self.clicked = True
        if pygame.mouse.get_pressed()[0] == 0:
            self.clicked = False
        screen.blit(self.image, self.rect)
        return action

# Player Class

class Player():
    def __init__(self, x, y):
        self.reset(x, y)

    def update(self, game_over):
        dx = 0
        dy = 0
        walk_cooldown = 5
        col_thresh = 20

        if game_over == 0:
            key = pygame.key.get_pressed()

            # Handle jumping (and double jumping)
            if key[pygame.K_SPACE] and not self.jumped:
                if not self.in_air:
                    jump_fx.play()
                    self.vel_y = -15
                    self.jumped = True
                    self.can_double_jump = self.double_jump_unlocked
                elif self.double_jump_unlocked and self.can_double_jump:
                    jump_fx.play()
                    self.vel_y = -15
                    self.jumped = True
                    self.can_double_jump = False

            if not key[pygame.K_SPACE]:
                self.jumped = False

            # Horizontal movement
            if key[pygame.K_a]:
                dx -= 5
                self.counter += 1
                self.direction = -1
            if key[pygame.K_d]:
                dx += 5
                self.counter += 1
                self.direction = 1
            if not key[pygame.K_a] and not key[pygame.K_d]:
                self.counter = 0
                self.index = 0
                self.image = self.images_right[self.index] if self.direction == 1 else self.images_left[self.index]

            # Walking animation
            if self.counter > walk_cooldown:
                self.counter = 0
                self.index = (self.index + 1) % len(self.images_right)
                self.image = self.images_right[self.index] if self.direction == 1 else self.images_left[self.index]

            # Apply gravity
            self.vel_y += 1
            if self.vel_y > 10:
                self.vel_y = 10
            dy += self.vel_y

            # Collision detection
            self.in_air = True
            for tile in world.tile_list:
                if tile[1].colliderect(self.rect.x + dx, self.rect.y, self.width, self.height):
                    dx = 0
                if tile[1].colliderect(self.rect.x, self.rect.y + dy, self.width, self.height):
                    if self.vel_y < 0:
                        dy = tile[1].bottom - self.rect.top
                        self.vel_y = 0
                    elif self.vel_y >= 0:
                        dy = tile[1].top - self.rect.bottom
                        self.vel_y = 0
                        self.in_air = False
                        self.can_double_jump = False

            # Check collision with hazards
            if pygame.sprite.spritecollide(self, blob_group, False) or pygame.sprite.spritecollide(self, lava_group, False):
                game_over = -1
                game_over_fx.play()

            # Level completion check
            if pygame.sprite.spritecollide(self, exit_group, False):
                game_over = 1

            # Platform collision logic
            for platform in platform_group:
                if platform.rect.colliderect(self.rect.x + dx, self.rect.y, self.width, self.height):
                    dx = 0
                if platform.rect.colliderect(self.rect.x, self.rect.y + dy, self.width, self.height):
                    if abs((self.rect.top + dy) - platform.rect.bottom) < col_thresh:
                        self.vel_y = 0
                        dy = platform.rect.bottom - self.rect.top
                    elif abs((self.rect.bottom + dy) - platform.rect.top) < col_thresh:
                        self.rect.bottom = platform.rect.top - 1
                        self.in_air = False
                        dy = 0
                        self.can_double_jump = False
                    if platform.move_x != 0:
                        self.rect.x += platform.move_direction

            # Apply movement
            self.rect.x += dx
            self.rect.y += dy

        elif game_over == -1:
            self.image = self.dead_image
            draw_text('YOU DIED', font, blue, (screen_width // 2) - 140, screen_height // 2)
            if self.rect.y > 200:
                self.rect.y -= 5

        screen.blit(self.image, self.rect)
        return game_over

    def reset(self, x, y):
        self.images_right = []
        self.images_left = []
        self.index = 0
        self.counter = 0
        for num in range(1, 5):
            img_right = pygame.image.load(f'img/guy{num}.png')
            img_right = pygame.transform.scale(img_right, (40, 80))
            img_left = pygame.transform.flip(img_right, True, False)
            self.images_right.append(img_right)
            self.images_left.append(img_left)
        self.dead_image = pygame.image.load('img/ghost.png')
        self.image = self.images_right[self.index]
        self.rect = self.image.get_rect()
        self.rect.x = x
        self.rect.y = y
        self.width = self.image.get_width()
        self.height = self.image.get_height()
        self.vel_y = 0
        self.jumped = False
        self.direction = 0
        self.in_air = True
        self.double_jump_unlocked = False
        self.can_double_jump = False

# World and Entities
class World():
    def __init__(self, data, level):
        self.tile_list = []
        dirt_img = pygame.image.load('img/dirt.png')
        grass_img = pygame.image.load('img/grass.png')
        for row_index, row in enumerate(data):
            for col_index, tile in enumerate(row):
                x = col_index * tile_size
                y = row_index * tile_size
                if tile == 1:
                    img = pygame.transform.scale(dirt_img, (tile_size, tile_size))
                    self.tile_list.append((img, img.get_rect(topleft=(x, y))))
                elif tile == 2:
                    img = pygame.transform.scale(grass_img, (tile_size, tile_size))
                    self.tile_list.append((img, img.get_rect(topleft=(x, y))))
                elif tile == 3:
                    blob_group.add(Enemy(x, y, level))
                elif tile == 4:
                    platform_group.add(Platform(x, y, 1, 0))
                elif tile == 5:
                    platform_group.add(Platform(x, y, 0, 1))
                elif tile == 6:
                    lava_group.add(Lava(x, y + (tile_size // 2)))
                elif tile == 7:
                    coin_group.add(Coin(x + tile_size // 2, y + tile_size // 2))
                elif tile == 8:
                    exit_group.add(Exit(x, y - tile_size // 2))

    def draw(self):
        for tile in self.tile_list:
            screen.blit(tile[0], tile[1])

# Other Entities (Enemies, Platforms, Lava
class Enemy(pygame.sprite.Sprite):
    def __init__(self, x, y, level):
        super().__init__()
        self.is_boss = level == 4
        if self.is_boss:
            bowser_img = pygame.image.load('img/bowser.png')
            bowser_img = pygame.transform.scale(bowser_img, (85, 85))
            self.images = [bowser_img]
            self.index = 0
            self.image = self.images[self.index]
            self.rect = self.image.get_rect(topleft=(x, y - 35))
        else:
            self.images = [
                pygame.transform.scale(pygame.image.load(f'img/blob{i}.png'), (50, 50))
                for i in range(1, 3)
            ]
            self.index = 0
            self.image = self.images[self.index]
            self.rect = self.image.get_rect(topleft=(x, y))
        self.move_direction = 1
        self.move_counter = 0
        self.animation_counter = 0

    def update(self):
        self.rect.x += self.move_direction
        self.move_counter += 1
        if abs(self.move_counter) > 50:
            self.move_direction *= -1
            self.move_counter *= -1
        if not self.is_boss:
            self.animation_counter += 1
            if self.animation_counter >= 15:
                self.animation_counter = 0
                self.index = (self.index + 1) % len(self.images)
                self.image = self.images[self.index]

class Platform(pygame.sprite.Sprite):
    def __init__(self, x, y, move_x, move_y):
        super().__init__()
        self.image = pygame.transform.scale(pygame.image.load('img/platform.png'), (tile_size, tile_size // 2))
        self.rect = self.image.get_rect(topleft=(x, y))
        self.move_x = move_x
        self.move_y = move_y
        self.move_direction = 1
        self.move_counter = 0

    def update(self):
        self.rect.x += self.move_direction * self.move_x
        self.rect.y += self.move_direction * self.move_y
        self.move_counter += 1
        if abs(self.move_counter) > 50:
            self.move_direction *= -1
            self.move_counter *= -1

class Lava(pygame.sprite.Sprite):
    def __init__(self, x, y):
        super().__init__()
        self.image = pygame.transform.scale(pygame.image.load('img/lava.png'), (tile_size, tile_size // 2))
        self.rect = self.image.get_rect(topleft=(x, y))

class Coin(pygame.sprite.Sprite):
    def __init__(self, x, y):
        super().__init__()
        self.image = pygame.transform.scale(pygame.image.load('img/coin.png'), (tile_size // 2, tile_size // 2))
        self.rect = self.image.get_rect(center=(x, y))

class Exit(pygame.sprite.Sprite):
    def __init__(self, x, y):
        super().__init__()
        self.image = pygame.transform.scale(pygame.image.load('img/exit.png'), (tile_size, int(tile_size * 1.5)))
        self.rect = self.image.get_rect(topleft=(x, y))

# Game Setup
player = Player(100, screen_height - 130)

# Sprite groups
blob_group = pygame.sprite.Group()
platform_group = pygame.sprite.Group()
lava_group = pygame.sprite.Group()
coin_group = pygame.sprite.Group()
exit_group = pygame.sprite.Group()

# Persistent UI coin
score_coin = Coin(tile_size // 2, tile_size // 2)
coin_group.add(score_coin)

# Load level data
if path.exists(f'level{level}_data'):
    pickle_in = open(f'level{level}_data', 'rb')
    world_data = pickle.load(pickle_in)
world = World(world_data, level)

# UI buttons
restart_button = Button(screen_width // 2 - 50, screen_height // 2 + 100, restart_img)
start_button = Button(screen_width // 2 - 350, screen_height // 2, start_img)
exit_button = Button(screen_width // 2 + 150, screen_height // 2, exit_img)

# Main Game Loop
run = True
while run:
    clock.tick(fps)
    screen.blit(bg_img, (0, 0))
    screen.blit(sun_img, (100, 100))

    if main_menu:
        screen.blit(logo_img, (screen_width // 2 - logo_img.get_width() // 2, screen_height // 2 - 300))
        if exit_button.draw():
            run = False
        if start_button.draw():
            main_menu = False
    else:
        world.draw()

        if game_over == 0:
            blob_group.update()
            platform_group.update()
            if pygame.sprite.spritecollide(player, coin_group, True):
                score += 1
                coin_fx.play()
                if score >= 30:
                    player.double_jump_unlocked = True
            draw_text('X ' + str(score), font_score, white, tile_size - 10, 10)
            draw_text('Lives: ' + str(lives), font_score, white, screen_width - 150, 10)
            if player.double_jump_unlocked:
                draw_text('Double Jump Unlocked!', font_score, white, screen_width // 2 - 120, 10)

        blob_group.draw(screen)
        platform_group.draw(screen)
        lava_group.draw(screen)
        coin_group.draw(screen)
        exit_group.draw(screen)

        game_over = player.update(game_over)

        if game_over == -1:
            if restart_button.draw():
                lives -= 1
                if lives > 0:
                    world = reset_level(level)
                    game_over = 0
                else:
                    level = 1
                    lives = 3
                    score = 0
                    world = reset_level(level)
                    game_over = 0

        if game_over == 1:
            level += 1
            if level <= max_levels:
                world = reset_level(level)
                game_over = 0
            else:
                draw_text('YOU WON MARIO!', font, blue, screen_width // 2 - 140, screen_height // 2)
                if restart_button.draw():
                    level = 1
                    score = 0
                    lives = 3
                    world = reset_level(level)
                    game_over = 0

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

    pygame.display.update()

pygame.quit()