import pygame
import sys

pygame.init()

# =====================================================
# НАСТРОЙКИ
# =====================================================

WIDTH, HEIGHT = 1000, 800
SCREEN = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Goblin & Knight")

CLOCK = pygame.time.Clock()
BLACK = (0, 0, 0)
GRAVITY = 1

# =====================================================
# ЗАГРУЗКА ФОНОВ
# =====================================================

BACKGROUND = pygame.image.load("images/back.png")
BACK_MAIN = pygame.image.load("images/back_main.jpg")
BACK_MAIN = pygame.transform.scale(BACK_MAIN, (WIDTH, HEIGHT))

# =====================================================
# ФУНКЦИЯ ВЫРЕЗКИ СПРАЙТОВ
# =====================================================

def get_image(sheet, frame_x, frame_y, width, height, scale, color):
    image = pygame.Surface((width, height)).convert_alpha()
    image.blit(sheet, (0, 0), (frame_x * width, frame_y * height, width, height))
    image = pygame.transform.scale(image, (width * scale, height * scale))
    image.set_colorkey(color)
    return image


# =====================================================
# КЛАСС ИГРОКА
# =====================================================

class Player:
    def __init__(self, x, y, sprite_sheet, frame_size, scale, controls):
        self.start_x = x
        self.start_y = y

        self.x = x
        self.y = y

        self.y_velocity = 0
        self.jump_power = 20
        self.jumping = False
        self.facing_right = True
        self.moving = False

        self.controls = controls

        self.anim_counter = 0
        self.ANIM_SPEED = 8

        w, h = frame_size

        # Анимации
        self.stand = get_image(sprite_sheet, 0, 0, w, h, scale, BLACK)
        self.walk1 = get_image(sprite_sheet, 1, 0, w, h, scale, BLACK)
        self.walk2 = get_image(sprite_sheet, 2, 0, w, h, scale, BLACK)
        self.jump_img = get_image(sprite_sheet, 3, 0, w, h, scale, BLACK)

        self.surface = self.stand
        self.rect = self.surface.get_rect(center=(self.x, self.y))

    # -------------------------
    # СБРОС ПЕРСОНАЖА
    # -------------------------
    def reset(self):
        self.x = self.start_x
        self.y = self.start_y
        self.y_velocity = 0
        self.jumping = False
        self.moving = False
        self.anim_counter = 0
        self.facing_right = True
        self.rect = self.surface.get_rect(center=(self.x, self.y))

    # -------------------------
    # УПРАВЛЕНИЕ
    # -------------------------
    def handle_input(self, keys):
        self.moving = False

        if keys[self.controls["right"]]:
            self.x += 5
            self.facing_right = True
            self.moving = True

        if keys[self.controls["left"]]:
            self.x -= 5
            self.facing_right = False
            self.moving = True

        if keys[self.controls["jump"]] and not self.jumping:
            self.jumping = True
            self.y_velocity = self.jump_power

        # Ограничение по экрану
        if self.x < 0:
            self.x = 0
        if self.x > WIDTH:
            self.x = WIDTH

    # -------------------------
    # ГРАВИТАЦИЯ
    # -------------------------
    def apply_gravity(self):
        self.y_velocity -= GRAVITY
        self.y -= self.y_velocity

        if self.y >= HEIGHT:
            self.y = HEIGHT
            self.y_velocity = 0
            self.jumping = False

    # -------------------------
    # АНИМАЦИЯ
    # -------------------------
    def update_animation(self):
        if self.jumping:
            self.surface = self.jump_img
        elif self.moving:
            self.anim_counter += 1
            frame = (self.anim_counter // self.ANIM_SPEED) % 2
            self.surface = self.walk1 if frame == 0 else self.walk2
        else:
            self.surface = self.stand

        if not self.facing_right:
            self.surface = pygame.transform.flip(self.surface, True, False)

        self.rect = self.surface.get_rect(center=(self.x, self.y))

    # -------------------------
    # ОТРИСОВКА
    # -------------------------
    def draw(self, screen):
        screen.blit(self.surface, self.rect)


# =====================================================
# ЗАГРУЗКА СПРАЙТОВ
# =====================================================

goblin_sheet = pygame.image.load("images/GoblinWorker/GoblinWorker.png").convert_alpha()
knight_sheet = pygame.image.load("images/sprites/knight.png").convert_alpha()

goblin = Player(
    200, 800,
    goblin_sheet,
    (48, 48),
    3,
    {
        "left": pygame.K_a,
        "right": pygame.K_d,
        "jump": pygame.K_SPACE
    }
)

knight = Player(
    500, 800,
    knight_sheet,
    (32, 32),
    3,
    {
        "left": pygame.K_KP4,
        "right": pygame.K_KP6,
        "jump": pygame.K_KP8
    }
)

# =====================================================
# ОБЪЕКТЫ УРОВНЯ
# =====================================================

door_image = pygame.image.load("images/door.png")
door_rect = door_image.get_rect(center=(800, 700))

button_play = pygame.image.load("images/play_button.png")
button_rect = button_play.get_rect(center=(500, 400))

# =====================================================
# ФУНКЦИЯ РЕСТАРТА
# =====================================================

def restart_level():
    goblin.reset()
    knight.reset()

# =====================================================
# GAME LOOP
# =====================================================

level = 0

while True:

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()

    keys = pygame.key.get_pressed()
    mouse_pos = pygame.mouse.get_pos()

    SCREEN.fill((0, 0, 0))

    # ================= MENU =================
    if level == 0:
        SCREEN.blit(BACK_MAIN, (0, 0))
        SCREEN.blit(button_play, button_rect)

        if button_rect.collidepoint(mouse_pos) and pygame.mouse.get_pressed()[0]:
            restart_level()
            level = 1

    # ================= GAME =================
    elif level == 1:

        SCREEN.blit(BACKGROUND, (0, 0))
        SCREEN.blit(door_image, door_rect)

        # --- Goblin ---
        goblin.handle_input(keys)
        goblin.apply_gravity()
        goblin.update_animation()
        goblin.draw(SCREEN)

        # --- Knight ---
        knight.handle_input(keys)
        knight.apply_gravity()
        knight.update_animation()
        knight.draw(SCREEN)

        # Победа
        if goblin.rect.colliderect(door_rect) and keys[pygame.K_e]:
            print("Goblin wins!")
            restart_level()
            level = 0

        if knight.rect.colliderect(door_rect) and keys[pygame.K_KP5]:
            print("Knight wins!")
            restart_level()
            level = 0

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