# isort: skip_file

import argparse
import os
import time
from typing import Any

import numpy as np

# Qwen2.5VL imports
# E2B imports

# SmolaAgents imports
from smolagents import HfApiModel, Model, Tool, tool
from smolagents.monitoring import LogLevel


from android_world import registry, suite_utils  # type: ignore
from android_world.env import json_action  # type: ignore
from screensuite.agents.client.android_env_client import AndroidEnvClient
from screensuite.agents.vision_agents.base_vision_agent import VisionAgent


def get_agent_summary_erase_images(agent):
    for memory_step in agent.memory.steps:
        if hasattr(memory_step, "observations_images"):
            memory_step.observations_images = None
        if hasattr(memory_step, "task_images"):
            memory_step.task_images = None
    return agent.write_memory_to_messages()


def _find_adb_directory() -> str:
    """Returns the directory where adb is located."""
    potential_paths = [
        os.path.expanduser("~/Library/Android/sdk/platform-tools/adb"),
        os.path.expanduser("~/Android/Sdk/platform-tools/adb"),
    ]
    for path in potential_paths:
        if os.path.isfile(path):
            return path
    raise EnvironmentError(
        "adb not found in the common Android SDK paths. Please install Android"
        " SDK and ensure adb is in one of the expected directories. If it's"
        " already installed, point to the installed location."
    )


GUIDANCE = (
    "Here are some useful guidelines you need to follow:\n"
    "General\n"
    "- Usually there will be multiple ways to complete a task, pick the"
    " easiest one. Also when something does not work as expected (due"
    " to various reasons), sometimes a simple retry can solve the problem,"
    " but if it doesn't (you can see that from the history), try to"
    " switch to other solutions.\n"
    "- Sometimes you may need to navigate the phone to gather information"
    " needed to complete the task, for example if user asks"
    ' "what is my schedule tomorrow", then you may want to open the calendar'
    " app (using the `open_app` action), look up information there, answer"
    " user's question (using the `answer` action) and finish (using"
    " the `status` action with complete as goal_status).\n"
    "- For requests that are questions (or chat messages), remember to use"
    " the `answer` action to reply to user explicitly before finish!"
    " Merely displaying the answer on the screen is NOT sufficient (unless"
    ' the goal is something like "show me ...").\n'
    "- If the desired state is already achieved (e.g., enabling Wi-Fi when"
    " it's already on), you can just complete the task.\n"
    "Action Related\n"
    "- Use the `open_app` action whenever you want to open an app"
    " (nothing will happen if the app is not installed), do not use the"
    " app drawer to open an app unless all other ways have failed.\n"
    "- Use the `input_text` action whenever you want to type"
    " something (including password) instead of clicking characters on the"
    " keyboard one by one. Sometimes there is some default text in the text"
    " field you want to type in, remember to delete them before typing.\n"
    "- For `click`, `long_press` and `input_text`, the index parameter you"
    " pick must be VISIBLE in the screenshot and also in the UI element"
    " list given to you (some elements in the list may NOT be visible on"
    " the screen so you can not interact with them).\n"
    "- Consider exploring the screen by using the `scroll`"
    " action with different directions to reveal additional content.\n"
    "- The direction parameter for the `scroll` action can be confusing"
    " sometimes as it's opposite to swipe, for example, to view content at the"
    ' bottom, the `scroll` direction should be set to "down". It has been'
    " observed that you have difficulties in choosing the correct direction, so"
    " if one does not work, try the opposite as well.\n"
    "Text Related Operations\n"
    "- Normally to select some text on the screen: <i> Enter text selection"
    " mode by long pressing the area where the text is, then some of the words"
    " near the long press point will be selected (highlighted with two pointers"
    " indicating the range) and usually a text selection bar will also appear"
    " with options like `copy`, `paste`, `select all`, etc."
    " <ii> Select the exact text you need. Usually the text selected from the"
    " previous step is NOT the one you want, you need to adjust the"
    " range by dragging the two pointers. If you want to select all text in"
    " the text field, simply click the `select all` button in the bar.\n"
    "- At this point, you don't have the ability to drag something around the"
    " screen, so in general you can not select arbitrary text.\n"
    "- To delete some text: the most traditional way is to place the cursor"
    " at the right place and use the backspace button in the keyboard to"
    " delete the characters one by one (can long press the backspace to"
    " accelerate if there are many to delete). Another approach is to first"
    " select the text you want to delete, then click the backspace button"
    " in the keyboard.\n"
    "- To copy some text: first select the exact text you want to copy, which"
    " usually also brings up the text selection bar, then click the `copy`"
    " button in bar.\n"
    "- To paste text into a text box, first long press the"
    " text box, then usually the text selection bar will appear with a"
    " `paste` button in it.\n"
    "- When typing into a text field, sometimes an auto-complete dropdown"
    " list will appear. This usually indicating this is a enum field and you"
    " should try to select the best match by clicking the corresponding one"
    " in the list.\n"
)


class AndroidAgent(VisionAgent):
    """Agent for Android automation with Qwen2.5VL vision capabilities"""

    def __init__(
        self,
        model: Model,
        data_dir: str,
        tools: list[Tool] | None = None,
        max_steps: int = 200,
        verbosity_level: LogLevel = LogLevel.INFO,
        planning_interval: int | None = None,
        transition_pause: float | None = None,
        **kwargs,
    ):
        self.env = AndroidEnvClient()
        self.env.reset(go_home=True)
        if transition_pause is not None and transition_pause < 0:
            raise ValueError(f"transition_pause must be non-negative, got {transition_pause}")
        self._transition_pause = transition_pause
        self.env.reinitialize_suite()

        # Initialize base agent
        super().__init__(
            model=model,
            data_dir=data_dir,
            tools=tools,
            max_steps=max_steps,
            verbosity_level=verbosity_level,
            planning_interval=planning_interval,
            **kwargs,
        )

    def reset(self, go_home: bool = False, respawn: bool = False) -> None:
        """Resets the agent."""
        if respawn:
            self.close()
            self.env = AndroidEnvClient()
        self.env.reset(go_home=go_home)

    def tear_down_task(self, task_type: str, task_idx: int) -> None:
        """Tear down the task."""
        self.env.tear_down_task(task_type, task_idx)

    def get_screenshot(self) -> np.ndarray[Any, Any]:
        """Get a screenshot from the Android environment."""

        if self._transition_pause is None:
            print("Waiting for screen to stabilize before grabbing state...")
            start = time.time()
            screenshot = self.env.get_screenshot(wait_to_stabilize=True)
            print("Fetched after %.1f seconds.", time.time() - start)
            return screenshot
        else:
            time.sleep(self._transition_pause)
            print("Pausing {:2.1f} seconds before grabbing state.".format(self._transition_pause))
            return self.env.get_screenshot(wait_to_stabilize=False)

    @property
    def transition_pause(self) -> float | None:
        return self._transition_pause

    @transition_pause.setter
    def transition_pause(self, transition_pause: float | None) -> None:
        self._transition_pause = transition_pause

    def _setup_desktop_tools(self):
        """Register all desktop tools"""

        @tool
        def click(x: int, y: int) -> str:
            """
            Click/tap at the specified coordinates on the screen
            Args:
                x: The x coordinate (horizontal position)
                y: The y coordinate (vertical position)
            """
            action = json_action.JSONAction(action_type="click", x=x, y=y)
            self.env.execute_action(action)
            self.logger.log(f"Clicked at coordinates ({x}, {y})")
            return f"Clicked at coordinates ({x}, {y})"

        @tool
        def long_press(x: int, y: int) -> str:
            """
            Long press at the specified coordinates on the screen
            Args:
                x: The x coordinate (horizontal position)
                y: The y coordinate (vertical position)
            """
            action = json_action.JSONAction(action_type="long_press", x=x, y=y)
            self.env.execute_action(action)
            self.logger.log(f"Long pressed at coordinates ({x}, {y})")
            return f"Long pressed at coordinates ({x}, {y})"

        @tool
        def input_text(text: str, x: int, y: int) -> str:
            """
            Type text into an editable text field at the specified coordinates
            Args:
                text: The text to input
                x: The x coordinate (horizontal position) of the text field
                y: The y coordinate (vertical position) of the text field
            """
            action = json_action.JSONAction(action_type="input_text", text=text, x=x, y=y)
            self.env.execute_action(action)
            self.logger.log(f"Input text '{text}' at coordinates ({x}, {y})")
            return f"Input text '{text}' at coordinates ({x}, {y})"

        @tool
        def keyboard_enter() -> str:
            """
            Press the Enter key
            """
            action = json_action.JSONAction(action_type="keyboard_enter")
            self.env.execute_action(action)
            self.logger.log("Pressed Enter key")
            return "Pressed Enter key"

        @tool
        def navigate_home() -> str:
            """
            Navigate to the home screen
            """
            action = json_action.JSONAction(action_type="navigate_home")
            self.env.execute_action(action)
            self.logger.log("Navigated to home screen")
            return "Navigated to home screen"

        @tool
        def navigate_back() -> str:
            """
            Navigate back to the previous screen
            """
            action = json_action.JSONAction(action_type="navigate_back")
            self.env.execute_action(action)
            self.logger.log("Navigated back")
            return "Navigated back"

        @tool
        def scroll(direction: str, x: int | None = None, y: int | None = None) -> str:
            """
            Scroll the screen or a specific area in one of the four directions
            Args:
                direction: The direction to scroll ("up", "down", "left", "right")
                x: Optional x coordinate to start the scroll from
                y: Optional y coordinate to start the scroll from
            """
            action = json_action.JSONAction(action_type="scroll", direction=direction, x=x, y=y)
            self.env.execute_action(action)
            if x is not None and y is not None:
                self.logger.log(f"Scrolled from coordinates ({x}, {y}) {direction}")
                return f"Scrolled from coordinates ({x}, {y}) {direction}"
            else:
                self.logger.log(f"Scrolled screen {direction}")
                return f"Scrolled screen {direction}"

        @tool
        def open_app(app_name: str) -> str:
            """
            Open an app
            Args:
                app_name: The name of the app to open
            """
            action = json_action.JSONAction(action_type="open_app", app_name=app_name)
            self.env.execute_action(action)
            self.logger.log(f"Opened app: {app_name}")
            return f"Opened app: {app_name}"

        @tool
        def wait() -> str:
            """
            Wait for the screen to update
            """
            action = json_action.JSONAction(action_type="wait")
            self.env.execute_action(action)
            self.logger.log("Waited for screen to update")
            return "Waited for screen to update"

        @tool
        def status(goal_status: str) -> str:
            """
            Set the status of the current goal
            Args:
                goal_status: The status of the goal ("complete" or "infeasible")
            """
            if goal_status == "complete":
                self.logger.log("Task completed successfully")
                return "Task completed successfully"
            elif goal_status == "infeasible":
                self.logger.log("Task is infeasible")
                return "Task is infeasible"
            else:
                self.logger.log(f"Unknown goal status: {goal_status}")
                return f"Unknown goal status: {goal_status}"

        @tool
        def answer(text: str) -> str:
            """
            Answer user's question
            Args:
                text: The answer text
            """
            action = json_action.JSONAction(action_type="answer", text=text)
            self.env.execute_action(action)
            self.logger.log(f"Answer: {text}")
            return f"Answer: {text}"

        # Register the tools
        self.tools["click"] = click
        self.tools["long_press"] = long_press
        self.tools["input_text"] = input_text
        self.tools["keyboard_enter"] = keyboard_enter
        self.tools["navigate_home"] = navigate_home
        self.tools["navigate_back"] = navigate_back
        self.tools["scroll"] = scroll
        self.tools["open_app"] = open_app
        self.tools["wait"] = wait
        self.tools["status"] = status
        self.tools["answer"] = answer

    def close(self) -> None:
        """Clean up resources"""
        if self.env:
            try:
                self.env.close()
            except Exception:
                pass


def main():
    """Run the Android Agent"""
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_dir", type=str, default="android_agent_data")
    parser.add_argument("--max_steps", type=int, default=200)
    parser.add_argument("--planning_interval", type=int, default=None)
    parser.add_argument("--use_v1_prompt", action="store_true")
    parser.add_argument("--transition_pause", type=float, default=None)
    args = parser.parse_args()

    # Initialize environment
    env = AndroidEnvClient()

    # Initialize agent
    agent = AndroidAgent(
        model=HfApiModel("Qwen/Qwen2.5-VL-7B"),
        data_dir=args.data_dir,
        env=env,
        max_steps=args.max_steps,
        planning_interval=args.planning_interval,
        use_v1_prompt=args.use_v1_prompt,
        transition_pause=args.transition_pause,
    )

    # Get task registry and create suite
    task_registry = registry.TaskRegistry()
    suite = suite_utils.create_suite(
        task_registry.get_registry(registry.TaskRegistry.ANDROID_WORLD_FAMILY),
        n_task_combinations=1,
        tasks=None,  # Run all tasks
    )
    agent.run(list(suite.values())[0].goal)


if __name__ == "__main__":
    main()
