import pygame
from pygame.locals import *


class InputManager:
    def __init__(self):
        self.keys_pressed = {}
        self.keys_just_pressed = {}
        self.keys_released = {}
        self.mouse_position = (0, 0)
        self.mouse_delta = (0, 0)
        self.mouse_buttons = [False, False, False]  # Left, Middle, Right
        self.mouse_buttons_just_pressed = [False, False, False]
        self.mouse_wheel = 0
        
    def update(self):
        # Reset one-frame states
        self.keys_just_pressed.clear()
        self.keys_released.clear()
        self.mouse_buttons_just_pressed = [False, False, False]
        self.mouse_wheel = 0
        
        # Get current key states
        keys = pygame.key.get_pressed()
        for key_code in range(len(keys)):
            was_pressed = self.keys_pressed.get(key_code, False)
            is_pressed = keys[key_code]
            
            if is_pressed and not was_pressed:
                self.keys_just_pressed[key_code] = True
            elif not is_pressed and was_pressed:
                self.keys_released[key_code] = True
            
            self.keys_pressed[key_code] = is_pressed
        
        # Get mouse state
        prev_mouse_pos = self.mouse_position
        self.mouse_position = pygame.mouse.get_pos()
        self.mouse_delta = (self.mouse_position[0] - prev_mouse_pos[0], 
                           self.mouse_position[1] - prev_mouse_pos[1])
        
        prev_buttons = self.mouse_buttons[:]
        self.mouse_buttons = list(pygame.mouse.get_pressed(3))
        
        for i in range(3):
            if self.mouse_buttons[i] and not prev_buttons[i]:
                self.mouse_buttons_just_pressed[i] = True
        
        # Handle events
        for event in pygame.event.get():
            if event.type == pygame.MOUSEWHEEL:
                self.mouse_wheel = event.y


# Global input manager instance
_input_manager = InputManager()


def update_input():
    """Call this once per frame to update input states"""
    _input_manager.update()


def listen(key):
    """
    Check if a keyboard key is currently pressed.
    Usage: if listen("W"): or if listen("SPACE"): or if listen("LEFT"):
    """
    key_map = {
        # Letters
        "A": K_a, "B": K_b, "C": K_c, "D": K_d, "E": K_e, "F": K_f,
        "G": K_g, "H": K_h, "I": K_i, "J": K_j, "K": K_k, "L": K_l,
        "M": K_m, "N": K_n, "O": K_o, "P": K_p, "Q": K_q, "R": K_r,
        "S": K_s, "T": K_t, "U": K_u, "V": K_v, "W": K_w, "X": K_x,
        "Y": K_y, "Z": K_z,
        # Numbers
        "0": K_0, "1": K_1, "2": K_2, "3": K_3, "4": K_4,
        "5": K_5, "6": K_6, "7": K_7, "8": K_8, "9": K_9,
        # Special keys
        "SPACE": K_SPACE, "ENTER": K_RETURN, "RETURN": K_RETURN,
        "ESCAPE": K_ESCAPE, "ESC": K_ESCAPE, "TAB": K_TAB,
        "BACKSPACE": K_BACKSPACE, "DELETE": K_DELETE,
        "UP": K_UP, "DOWN": K_DOWN, "LEFT": K_LEFT, "RIGHT": K_RIGHT,
        "SHIFT": K_LSHIFT, "CTRL": K_LCTRL, "ALT": K_LALT,
        "LSHIFT": K_LSHIFT, "RSHIFT": K_RSHIFT,
        "LCTRL": K_LCTRL, "RCTRL": K_RCTRL,
        "LALT": K_LALT, "RALT": K_RALT,
        "F1": K_F1, "F2": K_F2, "F3": K_F3, "F4": K_F4, "F5": K_F5,
        "F6": K_F6, "F7": K_F7, "F8": K_F8, "F9": K_F9, "F10": K_F10,
        "F11": K_F11, "F12": K_F12,
    }
    
    key_code = key_map.get(key.upper(), None)
    if key_code is not None:
        return _input_manager.keys_pressed.get(key_code, False)
    return False


def listen_pressed(key):
    """
    Check if a key was just pressed this frame (one-time detection).
    Usage: if listen_pressed("SPACE"): print("Jump!")
    """
    key_map = {
        "A": K_a, "B": K_b, "C": K_c, "D": K_d, "E": K_e, "F": K_f,
        "G": K_g, "H": K_h, "I": K_i, "J": K_j, "K": K_k, "L": K_l,
        "M": K_m, "N": K_n, "O": K_o, "P": K_p, "Q": K_q, "R": K_r,
        "S": K_s, "T": K_t, "U": K_u, "V": K_v, "W": K_w, "X": K_x,
        "Y": K_y, "Z": K_z,
        "0": K_0, "1": K_1, "2": K_2, "3": K_3, "4": K_4,
        "5": K_5, "6": K_6, "7": K_7, "8": K_8, "9": K_9,
        "SPACE": K_SPACE, "ENTER": K_RETURN, "RETURN": K_RETURN,
        "ESCAPE": K_ESCAPE, "ESC": K_ESCAPE, "TAB": K_TAB,
        "BACKSPACE": K_BACKSPACE, "DELETE": K_DELETE,
        "UP": K_UP, "DOWN": K_DOWN, "LEFT": K_LEFT, "RIGHT": K_RIGHT,
        "SHIFT": K_LSHIFT, "CTRL": K_LCTRL, "ALT": K_LALT,
        "F1": K_F1, "F2": K_F2, "F3": K_F3, "F4": K_F4, "F5": K_F5,
        "F6": K_F6, "F7": K_F7, "F8": K_F8, "F9": K_F9, "F10": K_F10,
        "F11": K_F11, "F12": K_F12,
    }
    
    key_code = key_map.get(key.upper(), None)
    if key_code is not None:
        return _input_manager.keys_just_pressed.get(key_code, False)
    return False


def listen_mouse(button="LEFT"):
    """
    Check if a mouse button is pressed.
    Usage: if listen_mouse("LEFT"): or if listen_mouse("RIGHT"):
    """
    button_map = {"LEFT": 0, "MIDDLE": 1, "RIGHT": 2}
    btn_index = button_map.get(button.upper(), 0)
    return _input_manager.mouse_buttons[btn_index]


def listen_mouse_pressed(button="LEFT"):
    """
    Check if a mouse button was just pressed this frame.
    """
    button_map = {"LEFT": 0, "MIDDLE": 1, "RIGHT": 2}
    btn_index = button_map.get(button.upper(), 0)
    return _input_manager.mouse_buttons_just_pressed[btn_index]


def get_mouse_position():
    """
    Get current mouse position as (x, y).
    """
    return _input_manager.mouse_position


def get_mouse_delta():
    """
    Get mouse movement since last frame as (dx, dy).
    Positive dx = moved right, Negative dx = moved left
    Positive dy = moved down, Negative dy = moved up
    """
    return _input_manager.mouse_delta


def get_mouse_wheel():
    """
    Get mouse wheel scroll amount.
    Positive = scrolled up, Negative = scrolled down
    """
    return _input_manager.mouse_wheel


def get_mouse_x():
    """Get mouse X position"""
    return _input_manager.mouse_position[0]


def get_mouse_y():
    """Get mouse Y position"""
    return _input_manager.mouse_position[1]


def get_mouse_dx():
    """Get mouse horizontal movement (right positive, left negative)"""
    return _input_manager.mouse_delta[0]


def get_mouse_dy():
    """Get mouse vertical movement (down positive, up negative)"""
    return _input_manager.mouse_delta[1]
