import base64
import random
import logging
import time
from typing import Any, Dict, List, Union, Literal

from google import genai
from google.genai import types
from desktop_env import actions

logger = logging.getLogger("desktopenv.agent")
API_RETRY_TIMES = 500
API_RETRY_INTERVAL = 5

GEMINI_DI = """
* You are utilising an Ubuntu virtual machine using x86_64 architecture with internet access.
* You can feel free to install Ubuntu applications with bash. Use curl instead of wget.
* To open browser, please just click on the Chrome icon.  Note, Chrome is what is installed on your system.
* Using bash you can start GUI applications, but you need to set DISPLAY=:0.
* When using bash with commands that are expected to output very large quantities of text, redirect into a tmp file and use other bash commands like `grep -n -B <lines before> -A <lines after> <query> <filename>` to confirm output.
* DO NOT ask users for clarification during task execution. DO NOT stop to request more information from users. Always take action using available tools.
* TASK FEASIBILITY: You can declare a task infeasible at any point during execution - whether at the beginning after taking a screenshot, or later after attempting some actions and discovering barriers. Carefully evaluate whether the task is feasible given the current system state, available applications, and task requirements. If you determine that a task cannot be completed due to:
  - Missing required applications or dependencies that cannot be installed
  - Insufficient permissions or system limitations
  - Contradictory or impossible requirements
  - Any other fundamental barriers that make completion impossible
  Then you MUST use the infeasible tool call.
* Home directory of this Ubuntu system is '/home/user'.
* If you need a password for sudo, the password of the computer is 'password'.
* Before you make an initial plan, you MUST use the 'take_screenshot' tool to take a screenshot of the current screen. The screenshot returned will show important context needed to decide what to do next.
"""

_KEY_DESCRIPTION = f"The allowed keys are: {', '.join(actions.KEYBOARD_KEYS)}"


_KEY_PARAMETERS = {
    "key": types.Schema(
        type=types.Type.STRING,
        description=_KEY_DESCRIPTION,
    ),
}

_KEYS_LIST_PARAMETERS = {
    "keys": types.Schema(
        type=types.Type.ARRAY,
        description="The list of keys to press: " + _KEY_DESCRIPTION,
        items=types.Schema(
            type=types.Type.STRING,
        ),
    ),
}

_RESPONSE_SCHEMA = types.Schema(
    type=types.Type.OBJECT,
    properties={
        "status": types.Schema(
            type=types.Type.STRING,
            description="Optional human-readable status of the actuation.",
        ),
        "screenshot": types.Schema(
            type=types.Type.OBJECT,
            description="A screenshot taken after the action.",
        ),
    },
)

OSWORLD_INFEASIBLE_ACTION = types.FunctionDeclaration(
    name="computer_agent_api:infeasible",
    description="Signal that the task is infeasible.",
    response=_RESPONSE_SCHEMA,
)

# ==========================================
# V4 Function Declarations (GCL/Literal)
# ==========================================

CLICK_ACTION = types.FunctionDeclaration(
    name="computer_agent_api:click",
    description="Left clicks at the coordinate.",
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "x": types.Schema(
                type=types.Type.INTEGER,
                description="0 (left) to 999 (right) of the screen.",
            ),
            "y": types.Schema(
                type=types.Type.INTEGER,
                description="0 (top) to 999 (bottom) of the screen.",
            ),
        },
        required=["y", "x"],
    ),
    response=_RESPONSE_SCHEMA,
)

DOUBLE_CLICK_ACTION = types.FunctionDeclaration(
    name="computer_agent_api:double_click",
    description="Double clicks at the coordinate.",
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "x": types.Schema(
                type=types.Type.INTEGER,
                description="0 (left) to 999 (right) of the screen.",
            ),
            "y": types.Schema(
                type=types.Type.INTEGER,
                description="0 (top) to 999 (bottom) of the screen.",
            ),
        },
        required=["y", "x"],
    ),
    response=_RESPONSE_SCHEMA,
)

TRIPLE_CLICK_ACTION = types.FunctionDeclaration(
    name="computer_agent_api:triple_click",
    description="Triple clicks at the coordinate.",
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "x": types.Schema(
                type=types.Type.INTEGER,
                description="0 (left) to 999 (right) of the screen.",
            ),
            "y": types.Schema(
                type=types.Type.INTEGER,
                description="0 (top) to 999 (bottom) of the screen.",
            ),
        },
        required=["y", "x"],
    ),
    response=_RESPONSE_SCHEMA,
)

MIDDLE_CLICK_ACTION = types.FunctionDeclaration(
    name="computer_agent_api:middle_click",
    description="Middle clicks at the coordinate.",
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "x": types.Schema(
                type=types.Type.INTEGER,
                description="0 (left) to 999 (right) of the screen.",
            ),
            "y": types.Schema(
                type=types.Type.INTEGER,
                description="0 (top) to 999 (bottom) of the screen.",
            ),
        },
        required=["y", "x"],
    ),
    response=_RESPONSE_SCHEMA,
)

RIGHT_CLICK_ACTION = types.FunctionDeclaration(
    name="computer_agent_api:right_click",
    description="Right clicks at the coordinate.",
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "x": types.Schema(
                type=types.Type.INTEGER,
                description="0 (left) to 999 (right) of the screen.",
            ),
            "y": types.Schema(
                type=types.Type.INTEGER,
                description="0 (top) to 999 (bottom) of the screen.",
            ),
        },
        required=["y", "x"],
    ),
    response=_RESPONSE_SCHEMA,
)

MOUSE_DOWN_ACTION = types.FunctionDeclaration(
    name="computer_agent_api:mouse_down",
    description="Presses and holds the mouse button at the coordinate.",
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "x": types.Schema(
                type=types.Type.INTEGER,
                description="0 (left) to 999 (right) of the screen.",
            ),
            "y": types.Schema(
                type=types.Type.INTEGER,
                description="0 (top) to 999 (bottom) of the screen.",
            ),
        },
        required=["y", "x"],
    ),
    response=_RESPONSE_SCHEMA,
)

MOUSE_UP_ACTION = types.FunctionDeclaration(
    name="computer_agent_api:mouse_up",
    description="Releases the mouse button at the coordinate.",
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "x": types.Schema(
                type=types.Type.INTEGER,
                description="0 (left) to 999 (right) of the screen.",
            ),
            "y": types.Schema(
                type=types.Type.INTEGER,
                description="0 (top) to 999 (bottom) of the screen.",
            ),
        },
        required=["y", "x"],
    ),
    response=_RESPONSE_SCHEMA,
)

MOVE_ACTION = types.FunctionDeclaration(
    name="computer_agent_api:move",
    description="Move the cursor to the specified position.",
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "x": types.Schema(
                type=types.Type.INTEGER,
                description="0 (left) to 999 (right) of the screen.",
            ),
            "y": types.Schema(
                type=types.Type.INTEGER,
                description="0 (top) to 999 (bottom) of the screen.",
            ),
        },
        required=["y", "x"],
    ),
    response=_RESPONSE_SCHEMA,
)

TYPE_ACTION = types.FunctionDeclaration(
    name="computer_agent_api:type",
    description="Types text.",
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "press_enter": types.Schema(
                type=types.Type.BOOLEAN,
                description="If true, presses enter after typing.",
                default=False,
            ),
            "text": types.Schema(
                type=types.Type.STRING,
                description="The text to type.",
            ),
        },
        required=["text"],
    ),
    response=_RESPONSE_SCHEMA,
)

DRAG_AND_DROP_ACTION = types.FunctionDeclaration(
    name="computer_agent_api:drag_and_drop",
    description="Drags item from start_y and start_x to end_y and end_x.",
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "end_x": types.Schema(
                type=types.Type.INTEGER,
                description="0 (left) to 999 (right) of the screen.",
            ),
            "end_y": types.Schema(
                type=types.Type.INTEGER,
                description="0 (top) to 999 (bottom) of the screen.",
            ),
            "start_x": types.Schema(
                type=types.Type.INTEGER,
                description="0 (left) to 999 (right) of the screen.",
            ),
            "start_y": types.Schema(
                type=types.Type.INTEGER,
                description="0 (top) to 999 (bottom) of the screen.",
            ),
        },
        required=["start_y", "start_x", "end_y", "end_x"],
    ),
    response=_RESPONSE_SCHEMA,
)

WAIT_ACTION = types.FunctionDeclaration(
    name="computer_agent_api:wait",
    description="Wait for the specified number of seconds.",
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "seconds": types.Schema(
                type=types.Type.NUMBER,
                description="Duration in seconds to wait. Accepts fractional values (e.g. 0.5 for half a second).",
                default=1.0,
            ),
        },
    ),
    response=_RESPONSE_SCHEMA,
)

PRESS_KEY_ACTION = types.FunctionDeclaration(
    name="computer_agent_api:press_key",
    description="Press the specified key and release it.",
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "key": types.Schema(
                type=types.Type.STRING,
                description="The key to press. Specify the key as an X KeySym (e.g. 'Return', 'BackSpace', 'Tab', 'space', 'a').",
            ),
        },
        required=["key"],
    ),
    response=_RESPONSE_SCHEMA,
)

KEY_DOWN_ACTION = types.FunctionDeclaration(
    name="computer_agent_api:key_down",
    description="Presses and holds the specified key.",
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "key": types.Schema(
                type=types.Type.STRING,
                description="The key to press. Specify the key as an X KeySym (e.g. 'Return', 'BackSpace', 'Tab', 'space', 'a').",
            ),
        },
        required=["key"],
    ),
    response=_RESPONSE_SCHEMA,
)

KEY_UP_ACTION = types.FunctionDeclaration(
    name="computer_agent_api:key_up",
    description="Releases the specified key.",
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "key": types.Schema(
                type=types.Type.STRING,
                description="The key to releaseSpecify the key as an X KeySym (e.g. 'Return', 'BackSpace', 'Tab', 'space', 'a').",
            ),
        },
        required=["key"],
    ),
    response=_RESPONSE_SCHEMA,
)

HOTKEY_ACTION = types.FunctionDeclaration(
    name="computer_agent_api:hotkey",
    description="Press the specified key combination.",
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "keys": types.Schema(
                type=types.Type.ARRAY,
                items=types.Schema(type=types.Type.STRING),
                description="The list of keys to press simultaneously. Specify the key as an X KeySym (e.g. 'Return', 'BackSpace', 'Tab', 'space', 'a').",
            ),
        },
        required=["keys"],
    ),
    response=_RESPONSE_SCHEMA,
)

SCROLL_ACTION = types.FunctionDeclaration(
    name="computer_agent_api:scroll",
    description="Scrolls up, down, left, or right at a coordinate.",
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "direction": types.Schema(
                type=types.Type.STRING,
                description='The direction to scroll ("up", "down", "left", "right").',
            ),
            "magnitude_in_wheel_clicks": types.Schema(
                type=types.Type.INTEGER,
                description="Number of mouse-wheel clicks to scroll (1 = one wheel click). Use larger values to scroll further.",
                default=2,
            ),
            "x": types.Schema(
                type=types.Type.INTEGER,
                description="0 (left) to 999 (right) of the screen.",
            ),
            "y": types.Schema(
                type=types.Type.INTEGER,
                description="0 (top) to 999 (bottom) of the screen.",
            ),
        },
        required=["y", "x", "direction"],
    ),
    response=_RESPONSE_SCHEMA,
)

TAKE_SCREENSHOT_ACTION = types.FunctionDeclaration(
    name="computer_agent_api:take_screenshot",
    description="Returns a screenshot of the current screen.",
    response=_RESPONSE_SCHEMA,
)

# ==========================================
# Lists grouping the FDs
# ==========================================

COMPLETE_OSWORLD_ACTIONS = [
    CLICK_ACTION,
    DOUBLE_CLICK_ACTION,
    TRIPLE_CLICK_ACTION,
    MIDDLE_CLICK_ACTION,
    RIGHT_CLICK_ACTION,
    MOUSE_DOWN_ACTION,
    MOUSE_UP_ACTION,
    MOVE_ACTION,
    TYPE_ACTION,
    DRAG_AND_DROP_ACTION,
    WAIT_ACTION,
    PRESS_KEY_ACTION,
    KEY_DOWN_ACTION,
    KEY_UP_ACTION,
    HOTKEY_ACTION,
    TAKE_SCREENSHOT_ACTION,
    SCROLL_ACTION,
    OSWORLD_INFEASIBLE_ACTION,
]

class GeminiFCAgent:
    """Agent that predicts next actions using Gemini with function calling."""

    client: genai.Client
    action_functions: list[types.FunctionDeclaration]
    messages: list[types.Content]

    def __init__(
        self,
        client: genai.Client,
        model_id: str,
        action_space: Literal["gemini_computer_use", "pyautogui"] = "gemini_computer_use",
        client_password="password",
        media_res="high",
        thinking_level: Literal["high", "medium", "low"] = "high",
    ):
        """Initialize the agent.

        Args:
          client: Gemini client for API calls
          model_id: Model ID to use for generation
          action_space: Action space type
          client_password: Password for client authentication
          thinking_level: Thinking depth
        """
        self.client = client
        self.model_id = model_id
        self.media_res = media_res
        self.thinking_level = thinking_level

        if action_space == "gemini_computer_use":
            self.tools = [
                types.Tool(
                    computer_use=types.ComputerUse(environment="ENVIRONMENT_DESKTOP")
                ),
                types.Tool(function_declarations=[OSWORLD_INFEASIBLE_ACTION]),
            ]
        elif action_space == "pyautogui":
            self.tools = [
                types.Tool(function_declarations=COMPLETE_OSWORLD_ACTIONS),
            ]
        else:
            raise RuntimeError(f"Unknown action space: {action_space}")

        self.action_space = action_space
        self.client_password = client_password

        self.messages = []  # Store conversation history
        self.thoughts = []
        self.actions = []
        self.observations = []


    def predict(
        self,
        instruction: str,
        obs: Dict,
        function_responses: list[dict[str, Any]] | None = None,
    ) -> List:
        """Predict the next action given instruction and current observation.

        Args:
          instruction: The task instruction
          obs: Current observation from the environment

        Returns:
          Tuple of (response, actions)
        """
        # Check if this is the first step (no messages yet)
        is_first_step = len(self.messages) == 0

        if is_first_step:
            # Add user message
            self.messages.append(
                types.Content(role="user", parts=[types.Part(text=instruction)])
            )

            # Store observation
            self.observations.append(obs)

        else:
            # Subsequent steps: Add function response from previous action if provided
            all_fr_parts = []
            assert function_responses
            for function_response in function_responses:
                if function_response and "func_name" in function_response:
                    # Extract function name and output from the response
                    func_name = function_response["func_name"]
                    output = function_response.get("output", {})

                    # Handle special case where output contains bytes for screenshot functions
                    action_call = function_response.get("action", {})
                    has_safety = "safety_decision" in action_call.get("parameters", {})

                    if isinstance(output, dict) and "result" in output:
                        # Error message
                        response_dict = dict(output)
                        if has_safety:
                            response_dict["safety_acknowledgement"] = {"decision": "acknowledged"}
                        part_to_append = types.Part.from_function_response(
                            name=func_name, response=response_dict
                        )
                    else:
                        # Assume it's always a screenshot
                        logger.info("Appending screenshot")
                        if not isinstance(output, bytes):
                            raise TypeError(f"Expected bytes for screenshot, got {type(output)}. This indicates an environment failure.")

                        response_dict = {"screenshot": {"$ref": "screenshot"}}
                        if has_safety:
                            response_dict["safety_acknowledgement"] = {"decision": "acknowledged"}

                        part_to_append = types.Part(
                            function_response=types.FunctionResponse(
                                name=func_name,
                                response=response_dict,
                                parts=[
                                    types.FunctionResponsePart(
                                        inline_data=types.FunctionResponseBlob(
                                            mime_type="image/png",
                                            data=output,
                                            display_name="screenshot",
                                        )
                                    )
                                ],
                            )
                        )
                    all_fr_parts.append(part_to_append)
            if all_fr_parts:
                self.messages.append(types.Content(role="user", parts=all_fr_parts))
        try:
            # Call Gemini with message history
            generation_config = types.GenerateContentConfig()
            generation_config.tools = self.tools
            generation_config.thinking_config = types.ThinkingConfig(
                include_thoughts=True,
                thinking_level=self.thinking_level,
            )
            generation_config.system_instruction = GEMINI_DI

            if self.media_res:
                _MEDIA_RES_DICT = {
                    "ultra_high": "MEDIA_RESOLUTION_ULTRA_HIGH",
                    "high": types.MediaResolution.MEDIA_RESOLUTION_HIGH,
                    "medium": types.MediaResolution.MEDIA_RESOLUTION_MEDIUM,
                    "low": types.MediaResolution.MEDIA_RESOLUTION_LOW,
                }
                generation_config.media_resolution = _MEDIA_RES_DICT[self.media_res]

            for attempt in range(API_RETRY_TIMES):
                try:
                    logger.info("Starting LM call")
                    response = self.client.models.generate_content(
                        model=self.model_id,
                        contents=self.messages,
                        config=generation_config,
                    )
                    logger.info("LLM response: %s", response)
                    logger.info("Finished LM call")
                    break
                except Exception as e:
                    error_msg = str(e)
                    logger.warning(
                        f"API error (attempt {attempt+1}/{API_RETRY_TIMES}): {error_msg}"
                    )
                    if attempt < API_RETRY_TIMES - 1:
                        time.sleep(
                            API_RETRY_INTERVAL
                            * min(2**attempt, 64)
                            * random.uniform(1, 1.5)
                        )
                    else:
                        raise
            # Extract function call from response
            if response.prompt_feedback and response.prompt_feedback.block_reason:
                raise ValueError(
                    f"Gemini response blocked due to safety: {response.prompt_feedback.block_reason}"
                )

            if response.prompt_feedback and response.prompt_feedback.block_reason:
                raise ValueError(
                    f"Gemini response blocked due to safety: {response.prompt_feedback.block_reason}"
                )

            if not response or not response.candidates or not response.candidates[0].content.parts:
                raise ValueError("No valid response from Gemini")

            parts_generated = response.candidates[0].content.parts

            self.messages.append(response.candidates[0].content)
            all_out = []
            for part in parts_generated:
                if hasattr(part, "function_call") and part.function_call is not None:
                    function_call = part.function_call
                    action = {
                        "action_type": function_call.name,
                        "is_function_call": True,
                        "parameters": (
                            dict(function_call.args) if function_call.args else {}
                        ),
                    }
                    all_out.append(action)
                    self.actions.append(action)
                elif part.text and part.thought:
                    response_text = part.text if hasattr(part, "text") else ""
                    action = {"action_type": "thought", "is_function_call": False, "parameters": response_text}
                    all_out.append(action)
                    self.thoughts.append(response_text)
                elif part.text and not part.thought:
                    # Plain user-visible text. the task is done.
                    response_text = part.text
                    action = {"action_type": "done", "is_function_call": False, "parameters": response_text}
                    all_out.append(action)
                    self.thoughts.append(response_text)
                else:
                    raise ValueError(f"Cannot handle part: {part}")
            return [], all_out

        except Exception as e:
            logger.error(f"Error in predict: {e}")
            # Return wait action as fallback
            # Handle text response - convert to n/a action
            error_msg = str(e)
            if "safety" in error_msg.lower():
                na_action = {
                    "action_type": "safety_error",
                    "parameters": {"reason": error_msg},
                }
            else:
                error_msg = str(e)
            if "safety" in error_msg.lower():
                na_action = {
                    "action_type": "safety_error",
                    "parameters": {"reason": error_msg},
                }
            else:
                na_action = {"action_type": "parsing_error", "parameters": {}}
            response_text = "FAILED"
            self.actions.append(na_action)
            self.thoughts.append(response_text)
            return [], [na_action]

    def reset(self, _logger=None):
        """Reset agent state to match PromptAgent specification."""
        self.thoughts = []
        self.actions = []
        self.observations = []
        self.messages = []

    def reset_history(self):
        """Clear conversation history."""
        self.messages.clear()
