from browsergym.core.action.parsers import highlevel_action_parser
from pydantic import BaseModel, ValidationError
from typing import Union, Literal, List, Optional, Dict, Any
import pyparsing as pp

# Enable packrat parsing for efficiency

# Define the base Action class
class Action(BaseModel):
    action_type: str

# Action subclasses for each specific action type

class SendMsgToUserAction(Action):
    action_type: Literal['send_msg_to_user'] = 'send_msg_to_user'
    text: str

class ReportInfeasibleAction(Action):
    action_type: Literal['report_infeasible'] = 'report_infeasible'
    reason: str

class NoopAction(Action):
    action_type: Literal['noop'] = 'noop'
    wait_ms: Optional[float] = 1000

class FillAction(Action):
    action_type: Literal['fill'] = 'fill'
    bid: str
    value: str

class CheckAction(Action):
    action_type: Literal['check'] = 'check'
    bid: str

class UncheckAction(Action):
    action_type: Literal['uncheck'] = 'uncheck'
    bid: str

class SelectOptionAction(Action):
    action_type: Literal['select_option'] = 'select_option'
    bid: str
    options: Union[str, List[str]]

class ClickAction(Action):
    action_type: Literal['click'] = 'click'
    bid: str
    button: Literal["left", "middle", "right"] = "left"
    modifiers: List[Literal["Alt", "Control", "Meta", "Shift"]] = []

class DblClickAction(Action):
    action_type: Literal['dblclick'] = 'dblclick'
    bid: str
    button: Literal["left", "middle", "right"] = "left"
    modifiers: List[Literal["Alt", "Control", "Meta", "Shift"]] = []

class HoverAction(Action):
    action_type: Literal['hover'] = 'hover'
    bid: str

class PressAction(Action):
    action_type: Literal['press'] = 'press'
    bid: str
    key_comb: str

class FocusAction(Action):
    action_type: Literal['focus'] = 'focus'
    bid: str

class ClearAction(Action):
    action_type: Literal['clear'] = 'clear'
    bid: str

class DragAndDropAction(Action):
    action_type: Literal['drag_and_drop'] = 'drag_and_drop'
    from_bid: str
    to_bid: str

class ScrollAction(Action):
    action_type: Literal['scroll'] = 'scroll'
    delta_x: float
    delta_y: float

class MouseMoveAction(Action):
    action_type: Literal['mouse_move'] = 'mouse_move'
    x: float
    y: float

class MouseUpAction(Action):
    action_type: Literal['mouse_up'] = 'mouse_up'
    x: float
    y: float
    button: Literal["left", "middle", "right"] = "left"

class MouseDownAction(Action):
    action_type: Literal['mouse_down'] = 'mouse_down'
    x: float
    y: float
    button: Literal["left", "middle", "right"] = "left"

class MouseClickAction(Action):
    action_type: Literal['mouse_click'] = 'mouse_click'
    x: float
    y: float
    button: Literal["left", "middle", "right"] = "left"

class MouseDblClickAction(Action):
    action_type: Literal['mouse_dblclick'] = 'mouse_dblclick'
    x: float
    y: float
    button: Literal["left", "middle", "right"] = "left"

class MouseDragAndDropAction(Action):
    action_type: Literal['mouse_drag_and_drop'] = 'mouse_drag_and_drop'
    from_x: float
    from_y: float
    to_x: float
    to_y: float

class KeyboardPressAction(Action):
    action_type: Literal['keyboard_press'] = 'keyboard_press'
    key: str

class KeyboardUpAction(Action):
    action_type: Literal['keyboard_up'] = 'keyboard_up'
    key: str

class KeyboardDownAction(Action):
    action_type: Literal['keyboard_down'] = 'keyboard_down'
    key: str

class KeyboardTypeAction(Action):
    action_type: Literal['keyboard_type'] = 'keyboard_type'
    text: str

class KeyboardInsertTextAction(Action):
    action_type: Literal['keyboard_insert_text'] = 'keyboard_insert_text'
    text: str

class GotoAction(Action):
    action_type: Literal['goto'] = 'goto'
    url: str

class GoBackAction(Action):
    action_type: Literal['go_back'] = 'go_back'

class GoForwardAction(Action):
    action_type: Literal['go_forward'] = 'go_forward'

class NewTabAction(Action):
    action_type: Literal['new_tab'] = 'new_tab'

class TabCloseAction(Action):
    action_type: Literal['tab_close'] = 'tab_close'

class TabFocusAction(Action):
    action_type: Literal['tab_focus'] = 'tab_focus'
    index: int

class UploadFileAction(Action):
    action_type: Literal['upload_file'] = 'upload_file'
    bid: str
    file: Union[str, List[str]]

class MouseUploadFileAction(Action):
    action_type: Literal['mouse_upload_file'] = 'mouse_upload_file'
    x: float
    y: float
    file: Union[str, List[str]]

class ReadPageAction(Action):
    action_type: Literal['read_page'] = 'read_page'
    pages: Optional[List[int]] = None

# Union of all action types
ActionType = Union[
    SendMsgToUserAction,
    ReportInfeasibleAction,
    NoopAction,
    FillAction,
    CheckAction,
    UncheckAction,
    SelectOptionAction,
    ClickAction,
    DblClickAction,
    HoverAction,
    PressAction,
    FocusAction,
    ClearAction,
    DragAndDropAction,
    ScrollAction,
    MouseMoveAction,
    MouseUpAction,
    MouseDownAction,
    MouseClickAction,
    MouseDblClickAction,
    MouseDragAndDropAction,
    KeyboardPressAction,
    KeyboardUpAction,
    KeyboardDownAction,
    KeyboardTypeAction,
    KeyboardInsertTextAction,
    GotoAction,
    GoBackAction,
    GoForwardAction,
    NewTabAction,
    TabCloseAction,
    TabFocusAction,
    UploadFileAction,
    MouseUploadFileAction,
    ReadPageAction
]

# Mapping from action type strings to Action classes
action_classes = {
    'send_msg_to_user': SendMsgToUserAction,
    'report_infeasible': ReportInfeasibleAction,
    'noop': NoopAction,
    'fill': FillAction,
    'check': CheckAction,
    'uncheck': UncheckAction,
    'select_option': SelectOptionAction,
    'click': ClickAction,
    'dblclick': DblClickAction,
    'hover': HoverAction,
    'press': PressAction,
    'focus': FocusAction,
    'clear': ClearAction,
    'drag_and_drop': DragAndDropAction,
    'scroll': ScrollAction,
    'mouse_move': MouseMoveAction,
    'mouse_up': MouseUpAction,
    'mouse_down': MouseDownAction,
    'mouse_click': MouseClickAction,
    'mouse_dblclick': MouseDblClickAction,
    'mouse_drag_and_drop': MouseDragAndDropAction,
    'keyboard_press': KeyboardPressAction,
    'keyboard_up': KeyboardUpAction,
    'keyboard_down': KeyboardDownAction,
    'keyboard_type': KeyboardTypeAction,
    'keyboard_insert_text': KeyboardInsertTextAction,
    'goto': GotoAction,
    'go_back': GoBackAction,
    'go_forward': GoForwardAction,
    'new_tab': NewTabAction,
    'tab_close': TabCloseAction,
    'tab_focus': TabFocusAction,
    'upload_file': UploadFileAction,
    'mouse_upload_file': MouseUploadFileAction,
    'read_page': ReadPageAction,
}


def parse_actions_from_string(action_string: str) -> List[ActionType]:
    """
    Parse a string of actions into a list of Action objects using the provided parser.
    """
    try:
        parsed_result = highlevel_action_parser.parse_string(action_string,parseAll=True)
    except pp.ParseException as pe:
        raise ValueError(f"Invalid action string: {pe}")

    actions = []

    for func_call in parsed_result:
        function_name = func_call[0]
        function_args = func_call[1]
        # function_kwargs = func_call[2]

        # Convert ParseResults to regular lists and dicts
        args = []
        for arg in function_args:
            args.append(arg)

        kwargs = {}
        # for kwarg in function_kwargs:
        #     key = kwarg[0]
        #     value = kwarg[1]
        #     kwargs[key] = value

        # Reconstruct the action dictionary
        action_dict = {'action_type': function_name}

        action_class = action_classes.get(function_name)
        if not action_class:
            raise ValueError(f"Unknown action type: {function_name}")

        # Get the parameter names of the action class, excluding 'action_type'
        param_names = list(action_class.__fields__.keys())
        param_names.remove('action_type')

        # Map positional arguments to parameter names
        for name, value in zip(param_names, args):
            action_dict[name] = value

        # Add keyword arguments
        action_dict.update(kwargs)

        # Create Action instance
        try:
            action = action_class(**action_dict)
            actions.append(action)
        except ValidationError as ve:
            raise ValueError(f"Validation error for action '{function_name}': {ve}")

    return actions

# Example usage:

action_string = 'click("b22", button = "right")'

parsed_actions = parse_actions_from_string(action_string)

if __name__ == '__main__':

    for action in parsed_actions:
        print(action)
