import json
from mm_agents.uipath.types_utils import (
    ComputerUseAction,
    ComputerUseStep,
    SupportedActions,
    PlanActionType,
    PlanAction,
    key_maps,
    ExecutionState,
    State,
)
import mm_agents.uipath.utils as utils
from mm_agents.uipath.action_planner import ActionPlanner, PlannerOutput
from mm_agents.uipath.grounder_client import GrounderClient


class UiPathComputerUseV1(object):
    def __init__(self):
        self.planner = ActionPlanner()
        self.executor = GrounderClient()

    async def predict_request(
        self, request_body: dict, model_name: str
    ) -> tuple[dict, dict]:
        state = State(
            task=request_body["userTask"],
            image_base64=request_body["image"],
            previous_steps=request_body.get("previousSteps", []),
        )

        execution_state = ExecutionState(model_name=model_name, execution_info={})
        output = await self.predict(state, execution_state)
        return output

    def process_grounding(
        self,
        plan_action: PlanAction,
        grounding_result: utils.GroundingOutput,
        x: int,
        y: int,
    ):
        match plan_action.action_type:
            case PlanActionType.Scroll:
                # guess the scroll direction if missing in the plan output
                if "direction" not in plan_action.parameters:
                    if "scroll up" in plan_action.description.lower():
                        scroll_direction = "up"
                    else:
                        scroll_direction = "down"
                else:
                    scroll_direction = plan_action.parameters["direction"]

                action = ComputerUseAction(
                    name=SupportedActions.Scroll,
                    description=plan_action.description,
                    parameters={"position": [x, y], "direction": scroll_direction},
                )

                if "distance" in plan_action.parameters:
                    match scroll_direction:
                        case "up":
                            action.parameters["offset"] = [
                                0,
                                plan_action.parameters["distance"],
                            ]
                        case "down":
                            action.parameters["offset"] = [
                                0,
                                -plan_action.parameters["distance"],
                            ]
                        case "left":
                            action.parameters["offset"] = [
                                plan_action.parameters["distance"],
                                0,
                            ]
                        case "right":
                            action.parameters["offset"] = [
                                -plan_action.parameters["distance"],
                                0,
                            ]
            case PlanActionType.Drag:
                assert grounding_result.end_position is not None, (
                    "End position must be provided for drag action"
                )
                x_end, y_end = grounding_result.end_position
                action = ComputerUseAction(
                    name=SupportedActions.Drag,
                    description=plan_action.description,
                    parameters={
                        "path": [
                            {"x": x, "y": y},
                            {"x": x_end, "y": y_end},
                        ]
                    },
                )
            case _:
                action_name = plan_action.action_type
                parameters = {"position": [x, y]}

                if plan_action.action_type == PlanActionType.DoubleClick:
                    action_name = SupportedActions.Click
                    parameters["click_type"] = "double"
                elif plan_action.action_type == PlanActionType.RightClick:
                    action_name = SupportedActions.Click
                    parameters["button"] = "right"
                elif plan_action.action_type == PlanActionType.MouseMove:
                    action_name = SupportedActions.MouseMove  # different names

                assert action_name in [
                    SupportedActions.Click,
                    SupportedActions.MouseMove,
                ]
                action = ComputerUseAction(
                    name=action_name,
                    description=plan_action.description,
                    parameters=parameters,
                )
        return action

    async def predict(
        self, state: State, execution_state: ExecutionState
    ) -> tuple[dict, dict]:
        planer_output: PlannerOutput = self.planner.predict(state, execution_state)
        plan_action = planer_output.plan_action

        action: ComputerUseAction | None = None
        step: ComputerUseStep | None = None

        match plan_action.action_type:
            case PlanActionType.KeyPress:
                keys = plan_action.parameters["key"].split(" ")
                keys = [key.strip() for key in keys]
                keys = [key_maps.get(key, key) for key in keys]
                action = ComputerUseAction(
                    name=SupportedActions.KeyPress,
                    description=plan_action.description,
                    parameters={"keys": keys},
                )
            case PlanActionType.Wait:
                action = ComputerUseAction(
                    name=SupportedActions.Wait,
                    description=plan_action.description,
                    parameters={},
                )
            case PlanActionType.ExtractData:
                # return a step with no action, just to store the extracted data
                step = ComputerUseStep(
                    description=plan_action.description,
                    actions=[],
                    additional_parameters={
                        "extracted_data": plan_action.parameters,
                    },
                    thought=planer_output.thought,
                )
            case PlanActionType.Finish:
                action = ComputerUseAction(
                    name=SupportedActions.Finish,
                    description=plan_action.description,
                    parameters=plan_action.parameters,
                )
            case (
                PlanActionType.Click
                | PlanActionType.MouseMove
                | PlanActionType.Scroll
                | PlanActionType.Drag
                | PlanActionType.DoubleClick
                | PlanActionType.RightClick
            ):
                if plan_action.action_type != PlanActionType.Drag:
                    grounding_result = await self.executor.predict(
                        state.image_base64,
                        plan_action.description,
                        action=plan_action.action_type,
                    )
                else:
                    grounding_result = await self.executor.predict(
                        state.image_base64,
                        plan_action.parameters["start_description"],
                        action=plan_action.action_type,
                    )
                    grounding_result_end = await self.executor.predict(
                        state.image_base64,
                        plan_action.parameters["end_description"],
                        action=plan_action.action_type,
                    )
                    grounding_result.end_position = grounding_result_end.position
                x, y = grounding_result.position
                action = self.process_grounding(plan_action, grounding_result, x, y)
            case PlanActionType.Type:
                action = ComputerUseAction(
                    name=SupportedActions.TypeInto,
                    description=plan_action.description,
                    parameters={"value": plan_action.parameters["text"]},
                )

        if step is None:
            assert action is not None
            step = ComputerUseStep(
                description=plan_action.description,
                actions=[action],
                additional_parameters={},
                thought=planer_output.thought,
            )

        # save additional data for history
        assert step.additional_parameters is not None
        step.additional_parameters["thought"] = planer_output.thought
        step.additional_parameters["review"] = planer_output.review
        step.additional_parameters.update(planer_output.additional_sections)
        step.additional_parameters["plan_action"] = json.dumps(plan_action.to_dict())

        history_image = state.image_base64
        previous_steps_parameters = {
            "max_chat_history_messages": 1000,
            "max_chat_history_images": self.planner.number_history_steps_with_images,
            "image": history_image,
        }
        agent_response = {
            "step": step.to_response_dict(),
            "previous_steps_parameters": previous_steps_parameters,
        }

        return agent_response