# libraries
from enum import Enum
from functools import lru_cache
import pygame

# project modules
from constants import *

# NOTE_1
# Refactored using AI (Claude), did a great job, uses some obscure syntax
# still much better than the original which was becoming unusable
# I will use some of the stuff introduced by AI in this file to refactor all the other files

# NOTE_2
# The .get() function takes 2 values, the first is the field of the dictionary, the second
# is a default value to return if the key is not found

# -------------------------------------------------------------
# Enum and tables
# -------------------------------------------------------------

class ButtonTypes(Enum):
    # The value doubles as the asset prefix: assets/<value>_btn.png
    # and assets/<value>_dark_btn.png
    PLAY   = "play"
    RESUME = "resume"
    MENU   = "menu"
    DELETE = "delete"
    QUIT   = "quit"
    RETRY  = "retry"
    FW     = "fw"
    BW     = "bw"
    GO     = "go"


class MenuType(Enum):
    PAUSE     = 0
    COMPLETED = 1
    START     = 2
    SELECT    = 3


class EntryState(Enum):
    UNLOCKED = 0
    LOCKED   = 1


# Unscaled hitbox sizes (width, height); anything not listed uses the default
DEFAULT_HITBOX = (232, 82) # Defaults to a normal button
HITBOX_SIZES = {
    ButtonTypes.FW: (105, 82),
    ButtonTypes.BW: (105, 82),
    ButtonTypes.GO: (83, 63),
}

CARD_FILES = {
    MenuType.PAUSE:     "assets/paused_card.png",
    MenuType.COMPLETED: "assets/completed_card.png",
    MenuType.START:     "assets/menu_card.png",
    MenuType.SELECT:    "assets/select_card.png",
}

ICON_FILES = {
    EntryState.LOCKED:   "assets/locked.png",
    EntryState.UNLOCKED: "assets/unlocked.png",
}

# Layout constants for SelectEntry
ENTRY_WIDTH     = 500
ENTRY_PADDING   = 20
ENTRY_ICON_X    = 800


# --------------------------------------------------------------------------
# Helpers
# --------------------------------------------------------------------------

@lru_cache(maxsize=None)
def load_image(path, scale=1):
    image = pygame.image.load(path).convert_alpha()
    if scale != 1:
        image = pygame.transform.scale_by(image, scale)
    return image


# --------------------------------------------------------------------------
# Widgets
# --------------------------------------------------------------------------

class Widget:
    # Anything a MenuPage can stack vertically.

    # Note: this notation is called type hinting
    rect: pygame.Rect

    def place(self, centerx, y):
        self.rect.centerx = centerx
        self.rect.y = y

    # This are just the interface for the methods implemented in the various widgets
    def height(self, text_font):
        return self.rect.height

    def draw(self, screen, mouse, text_font):
        raise NotImplementedError


class Button(Widget):
    def __init__(self, value, kind=ButtonTypes.PLAY):
        # value is a free-form field returned when the button is clicked
        self.value = value
        self.kind = kind

        # f-strings from python 3.6 to embed variables
        prefix = f"assets/{kind.value}"
        self.image      = load_image(f"{prefix}_btn.png", SCALE_FACTOR)
        self.dark_image = load_image(f"{prefix}_dark_btn.png", SCALE_FACTOR)

        width, height = HITBOX_SIZES.get(kind, DEFAULT_HITBOX)
        self.rect = pygame.Rect(0, 0, int(width * SCALE_FACTOR), int(height * SCALE_FACTOR))

    # text font is none because buttons have text in the image but the interface is the same for all widgets
    def height(self, text_font=None):
        return self.image.get_height() / 2

    def clicked(self, mouse):
        return self.rect.collidepoint(mouse)

    def draw(self, screen, mouse, text_font=None):
        image = self.dark_image if self.clicked(mouse) else self.image
        # The image is always centred on the hitbox, so nothing to keep in sync
        screen.blit(image, image.get_rect(center=self.rect.center))


class TextArea(Widget):
    def __init__(self, text, color):
        self.text = text
        self.color = color
        self.rect = pygame.Rect(0, 0, 0, 0)

    def height(self, text_font):
        return text_font.size(self.text)[1]

    def draw(self, screen, mouse, text_font, anchor="midtop"):
        # The anchor can be any rect attribute like midtop, topleft etc
        surface = text_font.render(self.text, True, self.color)
        position = getattr(self.rect, anchor)
        screen.blit(surface, surface.get_rect(**{anchor: position}))


class SelectEntry(Widget):
    # One row of the level-select menu: [go button] [label] ........ [lock icon]

    def __init__(self, text, value, state):
        self.value = value
        self.btn   = Button(value, ButtonTypes.GO)
        self.label = TextArea(text, DARK)
        self.icon  = load_image(ICON_FILES[state])
        self.icon_rect = self.icon.get_rect()
        self.rect = pygame.Rect(0, 0, ENTRY_WIDTH, self.btn.rect.height)

    def place(self, centerx, y):
        super().place(centerx, y)
        x = self.rect.x + ENTRY_PADDING
        self.btn.rect.topleft   = (x, self.rect.y)
        self.label.rect.midleft = (x + self.btn.rect.width * 3, self.rect.centery)
        self.icon_rect.topleft  = (ENTRY_ICON_X, self.rect.y)

    def clicked(self, mouse):
        return self.btn.clicked(mouse)

    def draw(self, screen, mouse, text_font):
        self.btn.draw(screen, mouse)
        self.label.draw(screen, mouse, text_font, anchor="center")
        screen.blit(self.icon, self.icon_rect)


# --------------------------------------------------------------------------
# Menu page
# --------------------------------------------------------------------------

class MenuPage:
    """One card with optional top text, buttons, entries and bottom text,
    stacked vertically and centred on screen."""

    def __init__(self, kind):
        self.top_text_area = None
        self.bottom_text_area = None
        self.buttons = []
        self.entries = []
        self.total_height = 0  # set by draw(), read by Menu
        self.bg_image = load_image(CARD_FILES[kind], SCALE_FACTOR)

    @property
    def is_full(self):
        return len(self.buttons) + len(self.entries) >= MAX_BTN

    def add_button(self, value, kind):
        self.buttons.append(Button(value, kind))

    def add_entry(self, text, value, state):
        self.entries.append(SelectEntry(text, value, state))

    def add_top_text(self, text):
        self.top_text_area = TextArea(text, DARK)

    def add_bottom_text(self, text):
        self.bottom_text_area = TextArea(text, DARK)

    def _elements(self):
        elements = []
        if self.top_text_area:
            elements.append(self.top_text_area)
        elements.extend(self.buttons)
        elements.extend(self.entries)
        if self.bottom_text_area:
            elements.append(self.bottom_text_area)
        return elements

    def draw(self, screen, mouse, btn_font, text_font):
        # btn_font is unused now; kept so existing calls don't break
        center_x, center_y = screen.get_rect().center
        screen.blit(self.bg_image, self.bg_image.get_rect(center=(center_x, center_y)))

        elements = self._elements()
        spacing = BTN_HEIGHT * 0.5
        heights = [e.height(text_font) for e in elements]
        self.total_height = sum(heights) + spacing * max(len(elements) - 1, 0)

        y = center_y - self.total_height / 2
        for element, height in zip(elements, heights):
            element.place(center_x, y)
            element.draw(screen, mouse, text_font)
            y += height + spacing

    def handle_event(self, event, mouse):
        if event.type == pygame.QUIT:
            return "quit_game"

        if event.type == pygame.MOUSEBUTTONDOWN:
            for clickable in (*self.buttons, *self.entries):
                if clickable.clicked(mouse):
                    return clickable.value
        return None


# --------------------------------------------------------------------------
# Multi-page menu
# --------------------------------------------------------------------------

class Menu:
    """Paged menu: new pages are created automatically when one is full.
    For now pages hold only buttons or entries, no top/bottom text."""

    def __init__(self, menu_type):
        self.menu_type = menu_type
        self.pages = [MenuPage(menu_type)]
        self.current_page_i = 0

        self.prev_btn = Button(None, ButtonTypes.BW)
        self.menu_btn = Button("menu", ButtonTypes.MENU)
        self.next_btn = Button(None, ButtonTypes.FW)
        self.nav_buttons = (self.prev_btn, self.menu_btn, self.next_btn)

    @property
    def page(self):
        return self.pages[self.current_page_i]

    def _page_with_room(self):
        if self.pages[-1].is_full:
            self.pages.append(MenuPage(self.menu_type))
        return self.pages[-1]

    def add_button(self, value, kind):
        self._page_with_room().add_button(value, kind)

    def add_entry(self, text, value, state):
        self._page_with_room().add_entry(text, value, state)

    def draw(self, screen, mouse, btn_font, text_font):
        center_x, center_y = screen.get_rect().center
        self.page.draw(screen, mouse, btn_font, text_font)

        spacing = BTN_HEIGHT * 0.5
        bottom_y = center_y + self.page.total_height / 2 + spacing

        group_width = sum(b.rect.width for b in self.nav_buttons)
        step = group_width / 2.5

        # Left / middle / right, spaced `step` apart around the centre
        for i, button in enumerate(self.nav_buttons):
            button.place(center_x + (i - 1) * step, bottom_y)
            button.draw(screen, mouse)

    def handle_event(self, event, mouse):
        if event.type == pygame.MOUSEBUTTONDOWN:
            if self.prev_btn.clicked(mouse):
                self.current_page_i = max(self.current_page_i - 1, 0)
                return None
            if self.next_btn.clicked(mouse):
                self.current_page_i = min(self.current_page_i + 1, len(self.pages) - 1)
                return None
            if self.menu_btn.clicked(mouse):
                return self.menu_btn.value

        # Also handles pygame.QUIT and clicks on the page's own buttons
        return self.page.handle_event(event, mouse)