import pygame
from constants import *
from enum import Enum

class ButtonTypes(Enum):
    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


# Button class that can store a value
class Button:
    def __init__(self, x, y, value, type=ButtonTypes.PLAY):
        # value is a utility field that can store something useful for the use of the button
        self.x = x
        self.y = y
        self.value = value

        self.original_image, self.dark_image = self.load_image(type)
        self.image = self.original_image.convert_alpha()

        # Collide rect measures based on the type
        if type == ButtonTypes.FW or type == ButtonTypes.BW:
            width = 105
            height = 82
        elif type == ButtonTypes.GO:
            width = 83
            height = 63
        else:
            width = 232
            height = 82
        
        self.rect = pygame.Rect(
            x,
            y,
            int(width * SCALE_FACTOR),
            int(height * SCALE_FACTOR)
        )

        self.image_rect = self.image.get_rect(
            center=self.rect.center
        )

    def load_image(self, type):
        if type == ButtonTypes.MENU:
            image = pygame.transform.scale_by(pygame.image.load('assets/menu_btn.png'), SCALE_FACTOR)
            dark_image =  pygame.transform.scale_by(pygame.image.load('assets/menu_dark_btn.png').convert_alpha(), SCALE_FACTOR)                   
        elif type == ButtonTypes.QUIT:
            image = pygame.transform.scale_by(pygame.image.load('assets/quit_btn.png'), SCALE_FACTOR)
            dark_image =  pygame.transform.scale_by(pygame.image.load('assets/quit_dark_btn.png').convert_alpha(), SCALE_FACTOR)
        elif type == ButtonTypes.DELETE:
            image = pygame.transform.scale_by(pygame.image.load('assets/delete_btn.png'), SCALE_FACTOR)
            dark_image =  pygame.transform.scale_by(pygame.image.load('assets/delete_dark_btn.png').convert_alpha(), SCALE_FACTOR)
        elif type == ButtonTypes.RESUME:
            image = pygame.transform.scale_by(pygame.image.load('assets/resume_btn.png'), SCALE_FACTOR)
            dark_image =  pygame.transform.scale_by(pygame.image.load('assets/resume_dark_btn.png').convert_alpha(), SCALE_FACTOR)
        elif type == ButtonTypes.PLAY:
            image = pygame.transform.scale_by(pygame.image.load('assets/play_btn.png'), SCALE_FACTOR)
            dark_image =  pygame.transform.scale_by(pygame.image.load('assets/play_dark_btn.png').convert_alpha(), SCALE_FACTOR)
        elif type == ButtonTypes.RETRY:
            image = pygame.transform.scale_by(pygame.image.load('assets/retry_btn.png'), SCALE_FACTOR)
            dark_image =  pygame.transform.scale_by(pygame.image.load('assets/retry_dark_btn.png').convert_alpha(), SCALE_FACTOR)
        elif type == ButtonTypes.FW:
            image = pygame.transform.scale_by(pygame.image.load('assets/fw_btn.png'), SCALE_FACTOR)
            dark_image =  pygame.transform.scale_by(pygame.image.load('assets/fw_dark_btn.png').convert_alpha(), SCALE_FACTOR)
        elif type == ButtonTypes.BW:
            image = pygame.transform.scale_by(pygame.image.load('assets/bw_btn.png'), SCALE_FACTOR)
            dark_image =  pygame.transform.scale_by(pygame.image.load('assets/bw_dark_btn.png').convert_alpha(), SCALE_FACTOR)
        elif type == ButtonTypes.GO:
            image = pygame.transform.scale_by(pygame.image.load('assets/go_btn.png'), SCALE_FACTOR)
            dark_image =  pygame.transform.scale_by(pygame.image.load('assets/go_dark_btn.png').convert_alpha(), SCALE_FACTOR)

        return image, dark_image

    def draw(self, screen, mouse, font):
        #text_surface = text_outline_perfect(font, self.text, WHITE, DARK)
        #text_rect = text_surface.get_rect(center=self.rect.center)
        # shadow_rect = sahdow_surface.get_rect(center=self.rect.center)

        if self.rect.collidepoint(mouse):
            screen.blit(self.dark_image, self.image_rect)
        else:
            screen.blit(self.image, self.image_rect)

        #pygame.draw.rect(screen, (255, 0, 255), self.rect, 2)
        #screen.blit(text_surface, text_rect)

    def clicked(self, mouse):
        return self.rect.collidepoint(mouse)
    
# Easy, text area
class TextArea:
    def __init__(self, x, y, text, color):
        self.text = text
        self.rect = pygame.Rect(x, y, 0, 0)
        self.color = color
    def draw(self, screen, font):
        text_surface = font.render(self.text, True, self.color)
        self.rect = text_surface.get_rect(center=self.rect.center)
        screen.blit(text_surface, self.rect)

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

class SelectEntry:
    def __init__(self, x, y, text, value, state):
        self.x = x
        self.y = y
        self.text   = text
        self.btn    = Button(0, 0, value, ButtonTypes.GO)
        self.area   = TextArea(0, 0, text, DARK)
        self.icon   = self.load_icon(state)
        self.icon_rect = self.icon.get_rect()
        self.value = value

        # Calculate full width
        #self.width = self.btn.rect.width + self.area.rect.width + self.icon_rect.width
        self.width = 500
        self.height = self.btn.rect.height
        self.rect = pygame.Rect(x, y, self.width, self.height)

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

    def load_icon(self, state):
        if state == EntryState.LOCKED:
            return pygame.image.load("assets/locked.png").convert_alpha()
        elif state == EntryState.UNLOCKED:
            return pygame.image.load("assets/unlocked.png").convert_alpha()
        
    def draw(self, screen, mouse, font):

        self.btn.rect.y = self.rect.y
        self.btn.image_rect.y = self.rect.y
        self.area.rect.y = self.rect.y
        self.icon_rect.y = self.rect.y

        x = self.rect.x + 20
        self.btn.rect.x = x
        self.btn.image_rect.x = x
        x += self.btn.rect.width*2

        self.area.rect.x = x

        self.icon_rect.x = 800

        self.btn.draw(screen, mouse, font)
        self.area.draw(screen, font)
        screen.blit(self.icon, self.icon_rect)


# Single page menu class
class MenuPage:
    def __init__(self, type):
        self.top_text_area = None
        self.bottom_text_area = None
        self.buttons = []
        self.entries = []
        self.bg_image = self.load_card(type)

    def load_card(self, type):
        if type == MenuType.PAUSE:
            return pygame.transform.scale_by(pygame.image.load('assets/paused_card.png').convert_alpha(), SCALE_FACTOR)
        elif type == MenuType.COMPLETED:
            return pygame.transform.scale_by(pygame.image.load('assets/completed_card.png').convert_alpha(), SCALE_FACTOR)
        elif type == MenuType.START:
            return pygame.transform.scale_by(pygame.image.load('assets/menu_card.png').convert_alpha(), SCALE_FACTOR)
        elif type == MenuType.SELECT:
            return pygame.transform.scale_by(pygame.image.load('assets/select_card.png').convert_alpha(), SCALE_FACTOR)
        
    # Button methods
    def add_button(self, value, type):
        self.buttons.append(Button(0, 0, value, type))
    def add_entry(self, text, value, state):
        self.entries.append(SelectEntry(0, 0, text, value, state))

    def add_top_text(self, text):
        self.top_text_area = TextArea(0, 0, text, DARK)
    def add_bottom_text(self, text):
        self.bottom_text_area = TextArea(0, 0, text, DARK)
    def draw(self, screen, mouse, btn_font, text_font):
        center_x, center_y = screen.get_rect().center

        # Draw the card image
        screen.blit(self.bg_image, (center_x-self.bg_image.get_width()/2, center_y-self.bg_image.get_height()/2))

        spacing = BTN_HEIGHT * 0.5
        elements = []

        if self.top_text_area is not None:
            elements.append(self.top_text_area)

        elements.extend(self.buttons)
        elements.extend(self.entries)

        if self.bottom_text_area is not None:
            elements.append(self.bottom_text_area)

        # Calculate the dimensions
        heights = []
        for element in elements:
            if isinstance(element, TextArea):
                text_surface = text_font.render(element.text, True, WHITE)
                heights.append(text_surface.get_height())
            elif isinstance(element, Button):
                heights.append(element.image.get_height()/2)
            else:
                heights.append(element.height)

        self.total_height = sum(heights) + spacing * (len(elements) - 1)
        y = center_y - self.total_height / 2

        for element, height in zip(elements, heights):
            element.rect.centerx = center_x
            element.rect.y = y

            if isinstance(element, Button):
                element.image_rect.center = element.rect.center
                element.draw(screen, mouse, btn_font)
            elif isinstance(element, TextArea):
                element.draw(screen, text_font)
            else:
                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 button in self.buttons:
                if button.clicked(mouse):
                    return button.value
            for entry in self.entries:
                if entry.clicked(mouse):
                    return entry.value

        return None

# Multi-page menu, automatically creates new page 3 buttons per page
class Menu:
    def __init__(self, menu_type):
        self.bottom_buttons = [Button(0, 0, None, ButtonTypes.BW), Button(0, 0, "menu", ButtonTypes.MENU), Button(0, 0, None, ButtonTypes.FW)]

        self.pages = []
        self.menu_type = menu_type
        # Creating the first page
        self.pages.append(MenuPage(self.menu_type))
        self.current_page_i = 0
        

    # For now the paged menu can only have pages with just buttons, no top or bottom text
    def add_button(self, value, type):
        current_page = self.pages[-1]
        if len(current_page.buttons) >= MAX_BTN:
            current_page = MenuPage(self.menu_type)
            self.pages.append(current_page)
        
        current_page.add_button(value, type)

    def add_entry(self, text, value, state):
        current_page = self.pages[-1]
        if len(current_page.entries) >= MAX_BTN:
            current_page = MenuPage(self.menu_type)
            self.pages.append(current_page)
        
        current_page.add_entry(text, value, state)

    def draw(self, screen, mouse, btn_font, text_font):
        center_x, center_y = screen.get_rect().center
        spacing = BTN_HEIGHT * 0.5

        self.pages[self.current_page_i].draw(screen, mouse, btn_font, text_font)

        # bottom y, OK
        bottom_y = center_y + self.pages[self.current_page_i].total_height / 2 + spacing

        # Calculate total width
        group_width = 0
        for button in self.bottom_buttons:
            group_width += button.rect.width
        
        x = center_x - group_width / 2.5
        for button in self.bottom_buttons:
            button.rect.centerx = x
            button.rect.y = bottom_y
            button.image_rect.center = button.rect.center
            x += group_width / 2.5

        for button in self.bottom_buttons:
            button.draw(screen, mouse, btn_font)

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

        # Maybe in the future the index can wrap around
        if event.type == pygame.MOUSEBUTTONDOWN:
            if self.bottom_buttons[0].clicked(mouse):
                if self.current_page_i > 0:
                    self.current_page_i -= 1
                return None

            if self.bottom_buttons[1].clicked(mouse):
                return self.bottom_buttons[1].value

            if self.bottom_buttons[2].clicked(mouse):
                if self.current_page_i < len(self.pages) - 1:
                    self.current_page_i += 1
                return None

            return self.pages[self.current_page_i].handle_event(event, mouse)
        
        return None