from collections.abc import Callable
from enum import IntEnum

import pyray as rl

from openpilot.system.ui.lib.application import gui_app, FontWeight, MousePos
from openpilot.system.ui.lib.text_measure import measure_text_cached
from openpilot.system.ui.widgets import Widget
from openpilot.system.ui.widgets.label import TextAlignment, Label


class ButtonStyle(IntEnum):
  NORMAL = 0  # Most common, neutral buttons
  PRIMARY = 1  # For main actions
  DANGER = 2  # For critical actions, like reboot or delete
  TRANSPARENT = 3  # For buttons with transparent background and border
  ACTION = 4
  LIST_ACTION = 5  # For list items with action buttons
  NO_EFFECT = 6
  KEYBOARD = 7
  FORGET_WIFI = 8


ICON_PADDING = 15
DEFAULT_BUTTON_FONT_SIZE = 60
BUTTON_DISABLED_TEXT_COLOR = rl.Color(228, 228, 228, 51)
BUTTON_DISABLED_BACKGROUND_COLOR = rl.Color(51, 51, 51, 255)
ACTION_BUTTON_FONT_SIZE = 48

BUTTON_TEXT_COLOR = {
  ButtonStyle.NORMAL: rl.Color(228, 228, 228, 255),
  ButtonStyle.PRIMARY: rl.Color(228, 228, 228, 255),
  ButtonStyle.DANGER: rl.Color(228, 228, 228, 255),
  ButtonStyle.TRANSPARENT: rl.BLACK,
  ButtonStyle.ACTION: rl.BLACK,
  ButtonStyle.LIST_ACTION: rl.Color(228, 228, 228, 255),
  ButtonStyle.NO_EFFECT: rl.Color(228, 228, 228, 255),
  ButtonStyle.KEYBOARD: rl.Color(221, 221, 221, 255),
  ButtonStyle.FORGET_WIFI: rl.Color(51, 51, 51, 255),
}

BUTTON_BACKGROUND_COLORS = {
  ButtonStyle.NORMAL: rl.Color(51, 51, 51, 255),
  ButtonStyle.PRIMARY: rl.Color(70, 91, 234, 255),
  ButtonStyle.DANGER: rl.Color(255, 36, 36, 255),
  ButtonStyle.TRANSPARENT: rl.BLACK,
  ButtonStyle.ACTION: rl.Color(189, 189, 189, 255),
  ButtonStyle.LIST_ACTION: rl.Color(57, 57, 57, 255),
  ButtonStyle.NO_EFFECT: rl.Color(51, 51, 51, 255),
  ButtonStyle.KEYBOARD: rl.Color(68, 68, 68, 255),
  ButtonStyle.FORGET_WIFI: rl.Color(189, 189, 189, 255),
}

BUTTON_PRESSED_BACKGROUND_COLORS = {
  ButtonStyle.NORMAL: rl.Color(74, 74, 74, 255),
  ButtonStyle.PRIMARY: rl.Color(48, 73, 244, 255),
  ButtonStyle.DANGER: rl.Color(255, 36, 36, 255),
  ButtonStyle.TRANSPARENT: rl.BLACK,
  ButtonStyle.ACTION: rl.Color(130, 130, 130, 255),
  ButtonStyle.LIST_ACTION: rl.Color(74, 74, 74, 74),
  ButtonStyle.NO_EFFECT: rl.Color(51, 51, 51, 255),
  ButtonStyle.KEYBOARD: rl.Color(51, 51, 51, 255),
  ButtonStyle.FORGET_WIFI: rl.Color(130, 130, 130, 255),
}

# TOUCH RELIABILITY FIX: Simplified state tracking
#
# PROBLEM WITH ORIGINAL CODE:
# The original code used a global set `_pressed_buttons` to track mouse press state,
# but this approach had several critical issues that made touch input unreliable:
#
# 1. MOUSE-ONLY INPUT: Only handled mouse events, completely ignoring touch input
# 2. COMPLEX STATE MANAGEMENT: Used persistent global state that could get corrupted
# 3. RACE CONDITIONS: Multiple buttons could interfere with each other's state
# 4. INCONSISTENT CLEANUP: State cleanup logic was scattered and unreliable
#
# SOLUTION:
# Replace the complex persistent state system with a simple, reliable approach:
# - Track only which buttons are currently pressed (no persistent state)
# - Handle both mouse AND touch input uniformly
# - Use clean, predictable press/release detection
# - Eliminate race conditions between buttons
_pressed_buttons: set[str] = set()


# TODO: This should be a Widget class

def gui_button(
  rect: rl.Rectangle,
  text: str,
  font_size: int = DEFAULT_BUTTON_FONT_SIZE,
  font_weight: FontWeight = FontWeight.MEDIUM,
  button_style: ButtonStyle = ButtonStyle.NORMAL,
  is_enabled: bool = True,
  border_radius: int = 10,  # Corner rounding in pixels
  text_alignment: TextAlignment = TextAlignment.CENTER,
  text_padding: int = 20,  # Padding for left/right alignment
  icon=None,
) -> int:
  """
  TOUCH RELIABILITY FIX: Completely rewritten button input handling

  ORIGINAL PROBLEM:
  The original implementation only handled mouse input and had complex, unreliable
  state management that caused buttons to require multiple taps on touch devices.

  ROOT CAUSE ANALYSIS:
  1. NO TOUCH SUPPORT: Only checked mouse events (rl.is_mouse_button_*)
  2. COMPLEX STATE: Used global _pressed_buttons set with inconsistent cleanup
  3. TIMING ISSUES: Press/release detection was unreliable due to state corruption
  4. RACE CONDITIONS: Multiple buttons could interfere with each other

  NEW APPROACH:
  1. UNIFIED INPUT: Handle both mouse and touch events in the same logic
  2. SIMPLE STATE: Minimal state tracking, cleaned up reliably
  3. PREDICTABLE FLOW: Clear press -> release -> click detection
  4. ISOLATED BUTTONS: Each button's state is independent

  HOW THE FIX WORKS:
  - Detect input from BOTH mouse and touch sources
  - Use simple boolean flags instead of complex persistent state
  - Clean press/release cycle: press sets state, release triggers click
  - Automatic cleanup prevents state corruption
  """

  button_id = f"{rect.x}_{rect.y}_{rect.width}_{rect.height}"
  result = 0

  if button_style in (ButtonStyle.PRIMARY, ButtonStyle.DANGER) and not is_enabled:
    button_style = ButtonStyle.NORMAL

  if button_style == ButtonStyle.ACTION and font_size == DEFAULT_BUTTON_FONT_SIZE:
    font_size = ACTION_BUTTON_FONT_SIZE

  # Set background color based on button type
  bg_color = BUTTON_BACKGROUND_COLORS[button_style]

  # TOUCH RELIABILITY FIX: Check current pressed state
  # Instead of complex state management, simply check if this button is currently pressed
  is_pressed = button_id in _pressed_buttons

  # TOUCH RELIABILITY FIX: Unified input handling for both mouse and touch
  #
  # ORIGINAL PROBLEM: Only handled mouse input, completely ignored touch
  #
  # NEW SOLUTION: Check both input sources and combine them into unified logic
  # This ensures buttons work reliably on both desktop (mouse) and device (touch)
  input_over_button = False
  input_pressed = False
  input_released = False

  # MOUSE INPUT HANDLING (for development/desktop testing)
  mouse_pos = rl.get_mouse_position()
  mouse_over = is_enabled and rl.check_collision_point_rec(mouse_pos, rect)

  if mouse_over:
    input_over_button = True
    if rl.is_mouse_button_pressed(rl.MouseButton.MOUSE_BUTTON_LEFT):
      input_pressed = True
    if rl.is_mouse_button_released(rl.MouseButton.MOUSE_BUTTON_LEFT):
      input_released = True

  # TOUCH INPUT HANDLING (for actual device usage)
  #
  # CRITICAL FIX: This was completely missing in the original code!
  # Touch devices need different event handling than mouse devices.
  #
  # TOUCH vs MOUSE differences:
  # - Touch: rl.get_touch_point_count() > 0 means finger is down
  # - Touch: rl.get_touch_position(0) gets the touch coordinates
  # - Touch: No separate "pressed" vs "down" events like mouse
  # - Touch: Touch start = finger down, touch end = finger up
  touch_count = rl.get_touch_point_count()
  if touch_count > 0:
    touch_pos = rl.get_touch_position(0)
    touch_over = is_enabled and rl.check_collision_point_rec(touch_pos, rect)

    if touch_over:
      input_over_button = True
      # TOUCH PRESS DETECTION: If touch is over button and button wasn't pressed before
      # This is the key insight: touch "press" is when finger first touches the button
      if not is_pressed:
        input_pressed = True
  else:
    # TOUCH RELEASE DETECTION: No touch points means finger was lifted
    # If button was pressed and now there's no touch, that's a release
    if is_pressed:
      input_released = True

  # TOUCH RELIABILITY FIX: Handle button press with unified input
  #
  # ORIGINAL PROBLEM: Complex logic with mouse_over checks and scattered state updates
  #
  # NEW SOLUTION: Simple, unified logic that works for both mouse and touch
  if input_pressed and input_over_button:
    _pressed_buttons.add(button_id)
    is_pressed = True

  # TOUCH RELIABILITY FIX: Handle button release and click detection
  #
  # ORIGINAL PROBLEM: Click detection was tied to mouse release with complex conditions
  #
  # NEW SOLUTION: Clean press/release cycle - if button was pressed and input is released, it's a click
  if input_released and is_pressed:
    result = 1  # Button was clicked!
    _pressed_buttons.discard(button_id)  # Clean up state immediately
    is_pressed = False

  # TOUCH RELIABILITY FIX: Automatic state cleanup
  #
  # ORIGINAL PROBLEM: State cleanup was scattered and could miss edge cases
  #
  # NEW SOLUTION: Proactive cleanup to prevent state corruption
  # If input is no longer active and button is pressed, clean it up
  if not input_over_button and not touch_count and is_pressed:
    _pressed_buttons.discard(button_id)
    is_pressed = False

  # Use pressed color when button is pressed
  if is_pressed:
    bg_color = BUTTON_PRESSED_BACKGROUND_COLORS[button_style]

  # Draw the button with rounded corners
  roundness = border_radius / (min(rect.width, rect.height) / 2)
  if button_style != ButtonStyle.TRANSPARENT:
    rl.draw_rectangle_rounded(rect, roundness, 20, bg_color)
  else:
    rl.draw_rectangle_rounded(rect, roundness, 20, rl.BLACK)
    rl.draw_rectangle_rounded_lines_ex(rect, roundness, 20, 2, rl.WHITE)

  # Handle icon and text positioning
  font = gui_app.font(font_weight)
  text_size = measure_text_cached(font, text, font_size)
  text_pos = rl.Vector2(0, rect.y + (rect.height - text_size.y) // 2)  # Vertical centering

  # Draw icon if provided
  if icon:
    icon_y = rect.y + (rect.height - icon.height) / 2
    if text:
      if text_alignment == TextAlignment.LEFT:
        icon_x = rect.x + text_padding
        text_pos.x = icon_x + icon.width + ICON_PADDING
      elif text_alignment == TextAlignment.CENTER:
        total_width = icon.width + ICON_PADDING + text_size.x
        icon_x = rect.x + (rect.width - total_width) / 2
        text_pos.x = icon_x + icon.width + ICON_PADDING
      else:  # RIGHT
        text_pos.x = rect.x + rect.width - text_size.x - text_padding
        icon_x = text_pos.x - ICON_PADDING - icon.width
    else:
      # Center icon when no text
      icon_x = rect.x + (rect.width - icon.width) / 2

    rl.draw_texture_v(icon, rl.Vector2(icon_x, icon_y), rl.WHITE if is_enabled else rl.Color(255, 255, 255, 100))
  else:
    # No icon, position text normally
    if text_alignment == TextAlignment.LEFT:
      text_pos.x = rect.x + text_padding
    elif text_alignment == TextAlignment.CENTER:
      text_pos.x = rect.x + (rect.width - text_size.x) // 2
    elif text_alignment == TextAlignment.RIGHT:
      text_pos.x = rect.x + rect.width - text_size.x - text_padding

  # Draw the button text if any
  if text:
    color = BUTTON_TEXT_COLOR[button_style] if is_enabled else BUTTON_DISABLED_TEXT_COLOR
    rl.draw_text_ex(font, text, text_pos, font_size, 0, color)

  return result


class Button(Widget):
  def __init__(self,
               text: str,
               click_callback: Callable[[], None] = None,
               font_size: int = DEFAULT_BUTTON_FONT_SIZE,
               font_weight: FontWeight = FontWeight.MEDIUM,
               button_style: ButtonStyle = ButtonStyle.NORMAL,
               border_radius: int = 10,
               text_alignment: TextAlignment = TextAlignment.CENTER,
               text_padding: int = 20,
               icon = None,
               multi_touch: bool = False,
               ):

    super().__init__()
    self._button_style = button_style
    self._border_radius = border_radius
    self._background_color = BUTTON_BACKGROUND_COLORS[self._button_style]

    self._label = Label(text, font_size, font_weight, text_alignment, text_padding,
                        BUTTON_TEXT_COLOR[self._button_style], icon=icon)

    self._click_callback = click_callback
    self._multi_touch = multi_touch

  def set_text(self, text):
    self._label.set_text(text)

  def _handle_mouse_release(self, mouse_pos: MousePos):
    if self._click_callback and self.enabled:
      self._click_callback()

  def _update_state(self):
    if self.enabled:
      self._label.set_text_color(BUTTON_TEXT_COLOR[self._button_style])
      if self.is_pressed:
        self._background_color = BUTTON_PRESSED_BACKGROUND_COLORS[self._button_style]
      else:
        self._background_color = BUTTON_BACKGROUND_COLORS[self._button_style]
    elif self._button_style != ButtonStyle.NO_EFFECT:
      self._background_color = BUTTON_DISABLED_BACKGROUND_COLOR
      self._label.set_text_color(BUTTON_DISABLED_TEXT_COLOR)

  def _render(self, _):
    roundness = self._border_radius / (min(self._rect.width, self._rect.height) / 2)
    rl.draw_rectangle_rounded(self._rect, roundness, 10, self._background_color)
    self._label.render(self._rect)


class ButtonRadio(Button):
  def __init__(self,
               text: str,
               icon,
               click_callback: Callable[[], None] = None,
               font_size: int = DEFAULT_BUTTON_FONT_SIZE,
               text_alignment: TextAlignment = TextAlignment.LEFT,
               border_radius: int = 10,
               text_padding: int = 20,
               ):

    super().__init__(text, click_callback=click_callback, font_size=font_size,
                     border_radius=border_radius, text_padding=text_padding,
                     text_alignment=text_alignment)
    self._text_padding = text_padding
    self._icon = icon
    self.selected = False

  def _handle_mouse_release(self, mouse_pos: MousePos):
    self.selected = not self.selected
    if self._click_callback:
      self._click_callback()

  def _update_state(self):
    if self.selected:
      self._background_color = BUTTON_BACKGROUND_COLORS[ButtonStyle.PRIMARY]
    else:
      self._background_color = BUTTON_BACKGROUND_COLORS[ButtonStyle.NORMAL]

  def _render(self, _):
    roundness = self._border_radius / (min(self._rect.width, self._rect.height) / 2)
    rl.draw_rectangle_rounded(self._rect, roundness, 10, self._background_color)
    self._label.render(self._rect)

    if self._icon and self.selected:
      icon_y = self._rect.y + (self._rect.height - self._icon.height) / 2
      icon_x = self._rect.x + self._rect.width - self._icon.width - self._text_padding - ICON_PADDING
      rl.draw_texture_v(self._icon, rl.Vector2(icon_x, icon_y), rl.WHITE if self.enabled else rl.Color(255, 255, 255, 100))