from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException
import time
import logging

from athenaidx.config import Config

class UserMenu:
    USER_MENU_BUTTON = (By.ID, 'user_menu_btn-button')
    
    # Map of menu option names to button IDs
    MENU_OPTION_IDS = {
        'HOG Screen': 'tools_HOG_1',
        'Help': 'user_showHelp',
        'Keyboard Help': 'user_keyboardHelp',
        'Sign Out': 'user_logout'
    }
    
    def __init__(self, config:Config):
        self.driver = config.get_driver()
    
    def _open_menu(self):
        user_menu_button = self.driver.find_element(*self.USER_MENU_BUTTON)
        user_menu_button.click()
        time.sleep(0.5)
        
    def _get_menu_options(self):
        self._open_menu()
        time.sleep(0.5)
        menu_items = self.driver.find_elements(By.CSS_SELECTOR, 'div.menu-item button')
        options = [item.text for item in menu_items]
        return options
    
    def select_menu_option(self, option_name):
        self._open_menu()
        
        # Try using the ID mapping first
        if option_name in self.MENU_OPTION_IDS:
            button_id = self.MENU_OPTION_IDS[option_name]
            try:
                option_element = WebDriverWait(self.driver, 5).until(
                    EC.element_to_be_clickable((By.ID, button_id))
                )
                option_element.click()
                time.sleep(2)
                logging.info(f"Successfully selected '{option_name}'.")
                return
            except TimeoutException:
                logging.error(f"Failed to find button with ID '{button_id}'.")
            
    def open_HOG_screen(self):
        self.select_menu_option('HOG Screen')
    
    def sign_out(self):
        self.select_menu_option('Sign Out')
