# isort: skip_file

import json
import logging
import os
import time
from typing import Any, Literal
import requests
from playwright.sync_api import sync_playwright, Browser, BrowserContext, Playwright

from screensuite.agents.client.osworld_vm_downloader import get_osworld_vm_path
from screensuite.agents.remote_env import create_remote_env_provider
from screensuite.agents.remote_env.docker.provider import DockerProviderConfig
from pydantic import BaseModel

logger = logging.getLogger()
logger.setLevel(logging.INFO)

KEYBOARD_KEYS = [
    "\t",
    "\n",
    "\r",
    " ",
    "!",
    '"',
    "#",
    "$",
    "%",
    "&",
    "'",
    "(",
    ")",
    "*",
    "+",
    ",",
    "-",
    ".",
    "/",
    "0",
    "1",
    "2",
    "3",
    "4",
    "5",
    "6",
    "7",
    "8",
    "9",
    ":",
    ";",
    "<",
    "=",
    ">",
    "?",
    "@",
    "[",
    "\\",
    "]",
    "^",
    "_",
    "`",
    "a",
    "b",
    "c",
    "d",
    "e",
    "f",
    "g",
    "h",
    "i",
    "j",
    "k",
    "l",
    "m",
    "n",
    "o",
    "p",
    "q",
    "r",
    "s",
    "t",
    "u",
    "v",
    "w",
    "x",
    "y",
    "z",
    "{",
    "|",
    "}",
    "~",
    "accept",
    "add",
    "alt",
    "altleft",
    "altright",
    "apps",
    "backspace",
    "browserback",
    "browserfavorites",
    "browserforward",
    "browserhome",
    "browserrefresh",
    "browsersearch",
    "browserstop",
    "capslock",
    "clear",
    "convert",
    "ctrl",
    "ctrlleft",
    "ctrlright",
    "decimal",
    "del",
    "delete",
    "divide",
    "down",
    "end",
    "enter",
    "esc",
    "escape",
    "execute",
    "f1",
    "f10",
    "f11",
    "f12",
    "f13",
    "f14",
    "f15",
    "f16",
    "f17",
    "f18",
    "f19",
    "f2",
    "f20",
    "f21",
    "f22",
    "f23",
    "f24",
    "f3",
    "f4",
    "f5",
    "f6",
    "f7",
    "f8",
    "f9",
    "final",
    "fn",
    "hanguel",
    "hangul",
    "hanja",
    "help",
    "home",
    "insert",
    "junja",
    "kana",
    "kanji",
    "launchapp1",
    "launchapp2",
    "launchmail",
    "launchmediaselect",
    "left",
    "modechange",
    "multiply",
    "nexttrack",
    "nonconvert",
    "num0",
    "num1",
    "num2",
    "num3",
    "num4",
    "num5",
    "num6",
    "num7",
    "num8",
    "num9",
    "numlock",
    "pagedown",
    "pageup",
    "pause",
    "pgdn",
    "pgup",
    "playpause",
    "prevtrack",
    "print",
    "printscreen",
    "prntscrn",
    "prtsc",
    "prtscr",
    "return",
    "right",
    "scrolllock",
    "select",
    "separator",
    "shift",
    "shiftleft",
    "shiftright",
    "sleep",
    "space",
    "stop",
    "subtract",
    "tab",
    "up",
    "volumedown",
    "volumemute",
    "volumeup",
    "win",
    "winleft",
    "winright",
    "yen",
    "command",
    "option",
    "optionleft",
    "optionright",
]

Params = dict[str, int | str]


class ExecuteResponse(BaseModel):
    status: str
    output: str
    error: str
    returncode: int


class ExecuteError(BaseModel):
    status: str
    message: str


ExecuteResult = ExecuteResponse | ExecuteError


class DesktopEnvClient:
    """Client for interacting with the Android environment server"""

    def __init__(
        self,
        pkgs_prefix: str = "import pyautogui; import time; pyautogui.FAILSAFE = False; {command}",
        apt_packages: list[str] = ["xdotool", "xclip", "scrot", "firefox-esr"],
        path_to_vm: None | str = None,
        os_type: Literal["Ubuntu", "Windows"] = "Ubuntu",
    ):
        if path_to_vm is None:
            path_to_vm = get_osworld_vm_path(os_type)
        volumes = [f"{os.path.abspath(path_to_vm)}:/System.qcow2:ro"]
        logger.info("Setting up Android environment using Docker - Initial setup may take 5-10 minutes. Please wait...")
        self.provider = create_remote_env_provider(
            config=DockerProviderConfig(
                ports_to_forward={5000, 8006, 8080, 9222},
                image="happysixd/osworld-docker:latest",
                healthcheck_endpoint="/screenshot",
                healthcheck_port=5000,
                privileged=True,
                cap_add=["NET_ADMIN"],
                devices=["/dev/kvm"],
                volumes=volumes,
                user="root",
            )
        )
        self.provider.start_emulator()
        ip_addr = self.provider.get_ip_address()
        self.base_url = f"http://{ip_addr.ip_address}:{ip_addr.host_port[5000]}"
        self.retry_times = 10
        self.retry_interval = 5
        self.pkgs_prefix = pkgs_prefix
        self.apt_packages = apt_packages
        # self.install_apt_packages(self.apt_packages)

        self.chromium_port = ip_addr.host_port[9222]
        self.browser: Browser | None = None
        self.chromium_context: BrowserContext | None = None
        self._playwright: Playwright | None = None

    def install_apt_packages(self, apt_packages: list[str]):
        """
        Installs the apt packages on the server.
        """
        cmds = [
            "apt-get update",
            f"apt-get install -y {' '.join(apt_packages)}",
        ]
        for cmd in cmds:
            logger.info(f"Running command: {cmd}")
            self.execute_shell_command(cmd)

    # Chrome setup
    def _chrome_open_tabs_setup(self, urls_to_open: list[str]) -> None:
        host = self.provider.get_ip_address().ip_address
        port = self.chromium_port  # fixme: this port is hard-coded, need to be changed from config file

        remote_debugging_url = f"http://{host}:{port}"
        logger.info("Connect to Chrome @: %s", remote_debugging_url)
        logger.debug("PLAYWRIGHT ENV: %s", repr(os.environ))

        playwright = sync_playwright().start()
        for attempt in range(15):
            if attempt > 0:
                time.sleep(5)

            browser = None
            try:
                browser = playwright.chromium.connect_over_cdp(remote_debugging_url)
            except Exception as e:
                if attempt < 14:
                    logger.error(f"Attempt {attempt + 1}: Failed to connect, retrying. Error: {e}")
                    continue
                else:
                    logger.error(f"Failed to connect after multiple attempts: {e}")
                    playwright.stop()
                    raise e

            if not browser:
                playwright.stop()
                return None

            logger.info("Opening %s...", urls_to_open)
            for i, url in enumerate(urls_to_open):
                # Use the first context (which should be the only one if using default profile)
                if i == 0:
                    context = browser.contexts[0]
                    context.set_extra_http_headers({"Accept-Language": "en-US;q=0.7,en;q=0.6"})
                page = context.new_page()  # Create a new page (tab) within the existing context
                try:
                    page.goto(url, timeout=60000)
                except Exception as e:
                    logger.warning("Opening %s exceeds time limit", url)  # only for human test
                logger.info(f"Opened tab {i + 1}: {url}")

                if i == 0:
                    # clear the default tab
                    default_page = context.pages[0]
                    default_page.close()

            # Store playwright instance as instance variable so it can be cleaned up later
            self._playwright = playwright
            # Do not close the context or browser; they will remain open after script ends
            self.browser, self.chromium_context = browser, context

            break

    # ================================
    # Keyboard and mouse actions space
    # ================================
    def move_mouse(self, x: int, y: int) -> ExecuteResult:
        """
        Moves the mouse to the specified coordinates.
        """
        return self.execute_python_command(f"pyautogui.moveTo({x}, {y})")

    def left_click(self) -> ExecuteResult:
        """
        Clicks the left button of the mouse at the specified coordinates.
        """
        return self.execute_python_command("pyautogui.click(button='left')")

    def right_click(self) -> ExecuteResult:
        """
        Clicks the right button of the mouse at the specified coordinates.
        """
        return self.execute_python_command("pyautogui.click(button='right')")

    def double_click(self) -> ExecuteResult:
        """
        Double-clicks the left button of the mouse at the specified coordinates.
        """
        return self.execute_python_command("pyautogui.doubleClick(button='left')")

    def write(self, text: str | list[str], delay_in_ms: int = 75) -> ExecuteResult:
        """
        Writes the specified text at the current cursor position.
        """
        return self.execute_python_command(f"pyautogui.write('{text}', interval={delay_in_ms / 1000})")

    def press(self, key: str | list[str]) -> ExecuteResult:
        """
        Presses a keyboard key
        """
        if isinstance(key, str):
            return self.execute_python_command(f"pyautogui.press('{key}')")
        else:
            return self.execute_python_command(f"pyautogui.hotkey({', '.join([f'{key}' for key in key])})")

    def drag(self, start: tuple[int, int], end: tuple[int, int]) -> ExecuteResult:
        """
        Drags the mouse from the start position to the end position.
        """
        self.move_mouse(start[0], start[1])
        return self.execute_python_command(
            f"pyautogui.dragTo({end[0]}, {end[1]}, duration=1.0, button='left', mouseDownUp=True)"
        )

    def scroll(self, x: int, y: int, direction: Literal["up", "down"] = "down", amount: int = 2) -> ExecuteResult:
        """
        Scrolls the mouse wheel in the specified direction.
        """
        if direction == "down":
            amount = -amount
        return self.execute_python_command(f"pyautogui.scroll({amount}, x={x}, y={y})")

    def open_chrome(self, url: str) -> ExecuteResult:
        """
        Opens the specified URL in Chrome.
        """
        logger.info("Opening Chrome")
        if self.chromium_context is None:
            self.execute_shell_command(
                "google-chrome --remote-debugging-port=1337 --disable-features=Translate",
                background=True,
            )
            time.sleep(5)
            self.execute_shell_command("socat tcp-listen:9222,fork tcp:localhost:1337", background=True)
            self._chrome_open_tabs_setup([url])
            time.sleep(5)
        else:
            self.chromium_context.new_page().goto(url, timeout=60000)
        if self.chromium_context is None:
            return ExecuteError(status="error", message="Failed to open Chrome.")
        else:
            return ExecuteResponse(status="success", output="", error="", returncode=0)

    def open(self, url_or_file: str, sleep_time: int = 10) -> ExecuteResult:
        """
        Opens the specified URL or file in the default application.
        """

        if url_or_file.startswith(("http://", "https://")):
            response = self.open_chrome(url_or_file)
        else:
            response = self.execute_shell_command(f"xdg-open {url_or_file}", background=True)
            logger.info(f"Waiting for the application to open for {sleep_time} seconds")
            time.sleep(sleep_time)

        if isinstance(response, ExecuteError):
            return response

        # Try to maximize the window using the window title
        logger.info("Maximizing the window")
        response = self.execute_shell_command("wmctrl -r :ACTIVE: -b add,maximized_vert,maximized_horz")
        time.sleep(3)
        return response

    # ================================

    def execute_python_command(self, command: str) -> ExecuteResult:
        """
        Executes a python command on the server.
        It can be used to execute the pyautogui commands, or... any other python command. who knows?
        """
        # command_list = ["python", "-c", self.pkgs_prefix.format(command=command)]
        command_list = ["python", "-c", self.pkgs_prefix.format(command=command)]
        payload = json.dumps({"command": command_list, "shell": False})

        for _ in range(self.retry_times):
            try:
                response = requests.post(
                    self.base_url + "/execute",
                    headers={"Content-Type": "application/json"},
                    data=payload,
                    timeout=90,
                )
                if response.status_code == 200:
                    logger.info("Command executed successfully: %s", response.text)
                    result = response.json()
                    if result["status"] == "error":
                        return ExecuteError(status="error", message=result["message"])
                    else:
                        return ExecuteResponse(**result)
                else:
                    logger.error("Failed to execute command. Status code: %d", response.status_code)
                    logger.info("Retrying to execute command.")
            except requests.exceptions.ReadTimeout:
                break
            except Exception as e:
                logger.error("An error occurred while trying to execute the command: %s", e)
                logger.info("Retrying to execute command.")
            time.sleep(self.retry_interval)

        logger.error("Failed to execute command.")
        return ExecuteError(status="error", message=f"Failed to execute command {command}.")

    def execute_shell_command(self, command: str, background: bool = False, timeout: int = 120) -> ExecuteResult:
        """
        Executes a terminal command on the server.
        If the command ends with &, it will be executed in the background and return immediately.
        """
        command_list = [command]
        payload = json.dumps({"command": command_list, "shell": True})

        if command.strip().endswith("&"):
            command = command.strip()[:-1]
            background = True

        # If command ends with &, execute it in background
        if background:
            try:
                requests.post(
                    self.base_url + "/setup/launch",
                    headers={"Content-Type": "application/json"},
                    data=payload,
                    timeout=5,
                )
                return ExecuteResponse(status="success", output="", error="", returncode=0)
            except Exception as e:
                logger.error("An error occurred while trying to execute the background command: %s", e)
                return ExecuteError(status="error", message=f"Failed to execute background command {command}.")

        # For non-background commands, use the existing retry logic
        for _ in range(self.retry_times):
            try:
                response = requests.post(
                    self.base_url + "/execute",
                    headers={"Content-Type": "application/json"},
                    data=payload,
                    timeout=timeout,
                )
                if response.status_code == 200:
                    logger.info("Command executed successfully: %s", response.text)
                    result = response.json()
                    if result["status"] == "error":
                        return ExecuteError(status="error", message=result["message"])
                    else:
                        return ExecuteResponse(**result)
                else:
                    logger.error("Failed to execute command. Status code: %d", response.status_code)
                    logger.info("Retrying to execute command.")
            except requests.exceptions.ReadTimeout:
                break
            except Exception as e:
                logger.error("An error occurred while trying to execute the command: %s", e)
                logger.info("Retrying to execute command.")
            time.sleep(self.retry_interval)

        logger.error("Failed to execute command.")
        return ExecuteError(status="error", message=f"Failed to execute command {command}.")

    def get_terminal_output(self) -> str | None:
        """
        Gets the terminal output from the server. None -> no terminal output or unexpected error.
        """

        for _ in range(self.retry_times):
            try:
                response = requests.get(self.base_url + "/terminal")
                if response.status_code == 200:
                    logger.info("Got terminal output successfully")
                    return response.json()["output"]
                else:
                    logger.error("Failed to get terminal output. Status code: %d", response.status_code)
                    logger.info("Retrying to get terminal output.")
            except Exception as e:
                logger.error("An error occurred while trying to get the terminal output: %s", e)
                logger.info("Retrying to get terminal output.")
            time.sleep(self.retry_interval)

        logger.error("Failed to get terminal output.")
        return None

    def get_desktop_screenshot(self) -> bytes | None:
        """
        Gets a screenshot from the server. With the cursor. None -> no screenshot or unexpected error.
        """

        for _ in range(self.retry_times):
            try:
                response = requests.get(self.base_url + "/screenshot")
                if response.status_code == 200:
                    logger.info("Got screenshot successfully")
                    return response.content
                else:
                    logger.error("Failed to get screenshot. Status code: %d", response.status_code)
                    logger.info("Retrying to get screenshot.")
            except Exception as e:
                logger.error("An error occurred while trying to get the screenshot: %s", e)
                logger.info("Retrying to get screenshot.")
            time.sleep(self.retry_interval)

        logger.error("Failed to get screenshot.")
        return None

    def get_playwright_screenshot(self) -> bytes | None:
        """
        Gets a screenshot using Playwright from the active browser context.
        Returns None if browser context is not available.
        """
        if self.chromium_context is None:
            logger.warning("No browser context available for screenshot")
            return None

        try:
            # Get the active page
            page = self.chromium_context.pages[0]
            # Take screenshot
            screenshot_bytes = page.screenshot(type="png", full_page=True)
            return screenshot_bytes
        except Exception as e:
            logger.error(f"Failed to take screenshot using Playwright: {e}")
            return None

    def screenshot(self) -> bytes | None:
        """
        Gets a screenshot from the server. With the cursor. None -> no screenshot or unexpected error.
        """
        return self.get_desktop_screenshot()

    def get_obs(self, require_terminal: bool = False):
        # We provide screenshot, terminal (optional)
        # can be customized and scaled
        return {
            "screenshot": self.get_desktop_screenshot(),
            "terminal": self.get_terminal_output() if require_terminal else None,
        }

    def get_vm_platform(self) -> str | None:
        """
        Gets the size of the vm screen.
        """
        result = self.execute_python_command("import platform; print(platform.system())")
        if result is not None and isinstance(result, ExecuteResponse):
            return result.output.strip()
        else:
            return None

    def get_screen_size(self) -> dict[str, int] | None:
        """
        Gets the size of the vm screen.
        """

        for _ in range(self.retry_times):
            try:
                response = requests.post(self.base_url + "/screen_size")
                if response.status_code == 200:
                    logger.info("Got screen size successfully")
                    return response.json()
                else:
                    logger.error("Failed to get screen size. Status code: %d", response.status_code)
                    logger.info("Retrying to get screen size.")
            except Exception as e:
                logger.error("An error occurred while trying to get the screen size: %s", e)
                logger.info("Retrying to get screen size.")
            time.sleep(self.retry_interval)

        logger.error("Failed to get screen size.")
        return None

    def get_vm_desktop_path(self) -> str | None:
        """
        Gets the desktop path of the vm.
        """

        for _ in range(self.retry_times):
            try:
                response = requests.post(self.base_url + "/desktop_path")
                if response.status_code == 200:
                    logger.info("Got desktop path successfully")
                    return response.json()["desktop_path"]
                else:
                    logger.error("Failed to get desktop path. Status code: %d", response.status_code)
                    logger.info("Retrying to get desktop path.")
            except Exception as e:
                logger.error("An error occurred while trying to get the desktop path: %s", e)
                logger.info("Retrying to get desktop path.")
            time.sleep(self.retry_interval)

        logger.error("Failed to get desktop path.")
        return None

    def get_vm_directory_tree(self, path) -> dict[str, Any] | None:
        """
        Gets the directory tree of the vm.
        """
        payload = json.dumps({"path": path})

        for _ in range(self.retry_times):
            try:
                response = requests.post(
                    self.base_url + "/list_directory", headers={"Content-Type": "application/json"}, data=payload
                )
                if response.status_code == 200:
                    logger.info("Got directory tree successfully")
                    return response.json()["directory_tree"]
                else:
                    logger.error("Failed to get directory tree. Status code: %d", response.status_code)
                    logger.info("Retrying to get directory tree.")
            except Exception as e:
                logger.error("An error occurred while trying to get directory tree: %s", e)
                logger.info("Retrying to get directory tree.")
            time.sleep(self.retry_interval)

        logger.error("Failed to get directory tree.")
        return None

    # Record video
    def start_recording(self):
        """
        Starts recording the screen.
        """

        for _ in range(self.retry_times):
            try:
                response = requests.post(self.base_url + "/start_recording")
                if response.status_code == 200:
                    logger.info("Recording started successfully")
                    return
                else:
                    logger.error("Failed to start recording. Status code: %d", response.status_code)
                    logger.info("Retrying to start recording.")
            except Exception as e:
                logger.error("An error occurred while trying to start recording: %s", e)
                logger.info("Retrying to start recording.")
            time.sleep(self.retry_interval)

        logger.error("Failed to start recording.")

    def end_recording(self, dest: str):
        """
        Ends recording the screen.
        """

        for _ in range(self.retry_times):
            try:
                response = requests.post(self.base_url + "/end_recording")
                if response.status_code == 200:
                    logger.info("Recording stopped successfully")
                    with open(dest, "wb") as f:
                        for chunk in response.iter_content(chunk_size=8192):
                            if chunk:
                                f.write(chunk)
                    return
                else:
                    logger.error("Failed to stop recording. Status code: %d", response.status_code)
                    logger.info("Retrying to stop recording.")
            except Exception as e:
                logger.error("An error occurred while trying to stop recording: %s", e)
                logger.info("Retrying to stop recording.")
            time.sleep(self.retry_interval)

        logger.error("Failed to stop recording.")

    def close(self) -> None:
        """Close the environment"""
        if self._playwright is not None:
            self._playwright.stop()
        self.provider.stop_emulator()

    def kill(self) -> None:
        """Kill the environment"""
        self.close()

    def health(self) -> bool:
        """Check the health of the environment"""
        try:
            response = requests.get(f"{self.base_url}/screenshot")
            response.raise_for_status()
        except Exception as e:
            print(f"Environment is not healthy: {e}")
            return False
        return True


if __name__ == "__main__":
    client = DesktopEnvClient()
    # obs = client.get_obs(require_terminal=True)

    def save_desktop_screenshot():
        screenshot = client.get_desktop_screenshot()
        if screenshot:
            with open("screenshot_desktop.png", "wb") as f:
                f.write(screenshot)

    def save_playwright_screenshot():
        screenshot = client.get_playwright_screenshot()
        if screenshot:
            with open("screenshot_playwright.png", "wb") as f:
                f.write(screenshot)

    def shell():
        while True:
            inp = input(">$ ")
            result = client.execute_shell_command(inp)
            if isinstance(result, ExecuteError):
                print(result.message)
            else:
                if result.output:
                    print(result.output)
                else:
                    print(result.error)
            save_desktop_screenshot()
            save_playwright_screenshot()

    def open_url():
        logger.info("Opening URL")
        client.open("https://www.rentalcars.com/")
        logger.info("Saving screenshot")
        # save_screenshot()
        save_desktop_screenshot()

    def test_actions():
        actions = [
            ("Move mouse to (100, 100)", lambda: client.move_mouse(100, 100)),
            ("Left click", lambda: client.left_click()),
            ("Right click", lambda: client.right_click()),
            ("Move mouse to (500, 500)", lambda: client.move_mouse(100, 100)),
            ("Double click", lambda: client.double_click()),
            ("Write 'Hello, world!'", lambda: client.write("Hello, world!")),
            ("Press Enter", lambda: client.press("Enter")),
            ("Open rentalcars.com", lambda: client.open("https://www.rentalcars.com/")),
            ("Scroll at (100, 100)", lambda: client.scroll(100, 100)),
            ("Execute 'ls -l'", lambda: client.execute_shell_command("ls -l")),
            ("Execute Python print", lambda: client.execute_python_command("print('Hello, world!')")),
        ]

        for action_name, action_func in actions:
            logger.info(f"\nNext action: {action_name}")
            input("Press Enter to execute this action...")
            result = action_func()
            if isinstance(result, ExecuteError):
                print(f"Error: {result.message}")
            else:
                logger.info("Action completed successfully")
            save_desktop_screenshot()
            logger.info("Saved screenshot")
            input("Press Enter to continue to next action...")

    try:
        open_url()
        # test_actions()
        shell()
    except (KeyboardInterrupt, EOFError):
        logger.info("Exiting...")
    except Exception as e:
        print(e)

    client.close()
