import asyncio
import base64
import os
import shlex
import shutil
from enum import StrEnum
from pathlib import Path
from typing import Literal, TypedDict, cast, get_args
from uuid import uuid4

from anthropic.types.beta import BetaToolComputerUse20241022Param, BetaToolUnionParam

from .base import BaseAnthropicTool, ToolError, ToolResult
from .run import run

OUTPUT_DIR = "/tmp/outputs"

TYPING_DELAY_MS = 12
TYPING_GROUP_SIZE = 50

Action_20241022 = Literal[
    "key",
    "type",
    "mouse_move",
    "left_click",
    "left_click_drag",
    "right_click",
    "middle_click",
    "double_click",
    "screenshot",
    "cursor_position",
]

Action_20250124 = (
    Action_20241022
    | Literal[
        "left_mouse_down",
        "left_mouse_up",
        "scroll",
        "hold_key",
        "wait",
        "triple_click",
    ]
)

Action_20251124 = Action_20250124 | Literal["zoom"]

ScrollDirection = Literal["up", "down", "left", "right"]


class Resolution(TypedDict):
    width: int
    height: int


# sizes above XGA/WXGA are not recommended (see README.md)
# scale down to one of these targets if ComputerTool._scaling_enabled is set
MAX_SCALING_TARGETS: dict[str, Resolution] = {
    "XGA": Resolution(width=1024, height=768),  # 4:3
    "WXGA": Resolution(width=1280, height=800),  # 16:10
    "FWXGA": Resolution(width=1366, height=768),  # ~16:9
}

CLICK_BUTTONS = {
    "left_click": 1,
    "right_click": 3,
    "middle_click": 2,
    "double_click": "--repeat 2 --delay 10 1",
    "triple_click": "--repeat 3 --delay 10 1",
}


class ScalingSource(StrEnum):
    COMPUTER = "computer"
    API = "api"


class ComputerToolOptions(TypedDict):
    display_height_px: int
    display_width_px: int
    display_number: int | None


def chunks(s: str, chunk_size: int) -> list[str]:
    return [s[i : i + chunk_size] for i in range(0, len(s), chunk_size)]


class BaseComputerTool:
    """
    A tool that allows the agent to interact with the screen, keyboard, and mouse of the current computer.
    The tool parameters are defined by Anthropic and are not editable.
    """

    name: Literal["computer"] = "computer"
    width: int
    height: int
    display_num: int | None

    _screenshot_delay = 2.0
    _scaling_enabled = True

    @property
    def options(self) -> ComputerToolOptions:
        width, height = self.scale_coordinates(
            ScalingSource.COMPUTER, self.width, self.height
        )
        return {
            "display_width_px": width,
            "display_height_px": height,
            "display_number": self.display_num,
        }

    def __init__(self):
        super().__init__()

        self.width = int(os.getenv("WIDTH") or 0)
        self.height = int(os.getenv("HEIGHT") or 0)
        assert self.width and self.height, "WIDTH, HEIGHT must be set"
        if (display_num := os.getenv("DISPLAY_NUM")) is not None:
            self.display_num = int(display_num)
            self._display_prefix = f"DISPLAY=:{self.display_num} "
        else:
            self.display_num = None
            self._display_prefix = ""

        self.xdotool = f"{self._display_prefix}xdotool"

    async def __call__(
        self,
        *,
        action: Action_20241022,
        text: str | None = None,
        coordinate: tuple[int, int] | None = None,
        start_coordinate: tuple[int, int] | None = None,
        **kwargs,
    ):
        if action in ("mouse_move", "left_click_drag"):
            if coordinate is None:
                raise ToolError(f"coordinate is required for {action}")
            if text is not None:
                raise ToolError(f"text is not accepted for {action}")

            if action == "left_click_drag":
                if start_coordinate is None:
                    raise ToolError(f"start_coordinate is required for {action}")
                start_x, start_y = self.validate_and_get_coordinates(start_coordinate)
                end_x, end_y = self.validate_and_get_coordinates(coordinate)
                command_parts = [
                    self.xdotool,
                    f"mousemove --sync {start_x} {start_y} mousedown 1 mousemove --sync {end_x} {end_y} mouseup 1",
                ]
                return await self.shell(" ".join(command_parts))
            elif action == "mouse_move":
                x, y = self.validate_and_get_coordinates(coordinate)
                command_parts = [self.xdotool, f"mousemove --sync {x} {y}"]
                return await self.shell(" ".join(command_parts))

        if action in ("key", "type"):
            if text is None:
                raise ToolError(f"text is required for {action}")
            if coordinate is not None:
                raise ToolError(f"coordinate is not accepted for {action}")
            if not isinstance(text, str):
                raise ToolError(output=f"{text} must be a string")

            if action == "key":
                command_parts = [self.xdotool, f"key -- {text}"]
                return await self.shell(" ".join(command_parts))
            elif action == "type":
                results: list[ToolResult] = []
                for chunk in chunks(text, TYPING_GROUP_SIZE):
                    command_parts = [
                        self.xdotool,
                        f"type --delay {TYPING_DELAY_MS} -- {shlex.quote(chunk)}",
                    ]
                    results.append(
                        await self.shell(" ".join(command_parts), take_screenshot=False)
                    )
                screenshot_base64 = (await self.screenshot()).base64_image
                return ToolResult(
                    output="".join(result.output or "" for result in results),
                    error="".join(result.error or "" for result in results),
                    base64_image=screenshot_base64,
                )

        if action in (
            "left_click",
            "right_click",
            "double_click",
            "middle_click",
            "screenshot",
            "cursor_position",
        ):
            if text is not None:
                raise ToolError(f"text is not accepted for {action}")
            if coordinate is not None:
                raise ToolError(f"coordinate is not accepted for {action}")

            if action == "screenshot":
                return await self.screenshot()
            elif action == "cursor_position":
                command_parts = [self.xdotool, "getmouselocation --shell"]
                result = await self.shell(
                    " ".join(command_parts),
                    take_screenshot=False,
                )
                output = result.output or ""
                x, y = self.scale_coordinates(
                    ScalingSource.COMPUTER,
                    int(output.split("X=")[1].split("\n")[0]),
                    int(output.split("Y=")[1].split("\n")[0]),
                )
                return result.replace(output=f"X={x},Y={y}")
            else:
                command_parts = [self.xdotool, f"click {CLICK_BUTTONS[action]}"]
                return await self.shell(" ".join(command_parts))

        raise ToolError(f"Invalid action: {action}")

    def validate_and_get_coordinates(self, coordinate: tuple[int, int] | None = None):
        if not isinstance(coordinate, list) or len(coordinate) != 2:
            raise ToolError(f"{coordinate} must be a tuple of length 2")
        if not all(isinstance(i, int) and i >= 0 for i in coordinate):
            raise ToolError(f"{coordinate} must be a tuple of non-negative ints")

        return self.scale_coordinates(ScalingSource.API, coordinate[0], coordinate[1])

    async def screenshot(self):
        """Take a screenshot of the current screen and return the base64 encoded image."""
        path = await self._new_output_path("screenshot")
        result = await self._capture(path)
        if self._scaling_enabled:
            x, y = self.screenshot_size()
            await self.shell(
                f"convert {path} -resize {x}x{y}! {path}", take_screenshot=False
            )

        if path.exists():
            return result.replace(
                base64_image=base64.b64encode(path.read_bytes()).decode()
            )
        raise ToolError(f"Failed to take screenshot: {result.error}")

    def screenshot_size(self) -> tuple[int, int]:
        """The dimensions of the screenshots sent to the model.

        This is the coordinate frame the model emits coordinates in; it differs
        from the physical display whenever scaling applies.
        """
        return self.scale_coordinates(ScalingSource.COMPUTER, self.width, self.height)

    async def _new_output_path(self, prefix: str) -> Path:
        output_dir = Path(OUTPUT_DIR)
        await asyncio.to_thread(output_dir.mkdir, parents=True, exist_ok=True)
        return output_dir / f"{prefix}_{uuid4().hex}.png"

    async def _capture(self, path: Path) -> ToolResult:
        """Capture the display at physical resolution into `path`."""
        # Try gnome-screenshot first
        if shutil.which("gnome-screenshot"):
            screenshot_cmd = f"{self._display_prefix}gnome-screenshot -f {path} -p"
        else:
            # Fall back to scrot if gnome-screenshot isn't available
            screenshot_cmd = f"{self._display_prefix}scrot -p {path}"
        return await self.shell(screenshot_cmd, take_screenshot=False)

    async def shell(self, command: str, take_screenshot=True) -> ToolResult:
        """Run a shell command and return the output, error, and optionally a screenshot."""
        _, stdout, stderr = await run(command)
        base64_image = None

        if take_screenshot:
            # delay to let things settle before taking a screenshot
            await asyncio.sleep(self._screenshot_delay)
            base64_image = (await self.screenshot()).base64_image

        return ToolResult(output=stdout, error=stderr, base64_image=base64_image)

    def scale_coordinates(self, source: ScalingSource, x: int, y: int):
        """Scale coordinates to a target maximum resolution."""
        if not self._scaling_enabled:
            return x, y
        ratio = self.width / self.height
        target_dimension = None
        for dimension in MAX_SCALING_TARGETS.values():
            # allow some error in the aspect ratio - not ratios are exactly 16:9
            if abs(dimension["width"] / dimension["height"] - ratio) < 0.02:
                if dimension["width"] < self.width:
                    target_dimension = dimension
                break
        if target_dimension is None:
            return x, y
        # should be less than 1
        x_scaling_factor = target_dimension["width"] / self.width
        y_scaling_factor = target_dimension["height"] / self.height
        if source == ScalingSource.API:
            if x > self.width or y > self.height:
                raise ToolError(f"Coordinates {x}, {y} are out of bounds")
            # scale up
            return round(x / x_scaling_factor), round(y / y_scaling_factor)
        # scale down
        return round(x * x_scaling_factor), round(y * y_scaling_factor)


class ComputerTool20241022(BaseComputerTool, BaseAnthropicTool):
    api_type: Literal["computer_20241022"] = "computer_20241022"

    def to_params(self) -> BetaToolComputerUse20241022Param:
        return {"name": self.name, "type": self.api_type, **self.options}


class ComputerTool20250124(BaseComputerTool, BaseAnthropicTool):
    api_type: Literal["computer_20250124"] = "computer_20250124"

    def to_params(self):
        return cast(
            BetaToolUnionParam,
            {"name": self.name, "type": self.api_type, **self.options},
        )

    async def __call__(
        self,
        *,
        action: Action_20250124,
        text: str | None = None,
        coordinate: tuple[int, int] | None = None,
        start_coordinate: tuple[int, int] | None = None,
        scroll_direction: ScrollDirection | None = None,
        scroll_amount: int | None = None,
        duration: int | float | None = None,
        key: str | None = None,
        **kwargs,
    ):
        if action in ("left_mouse_down", "left_mouse_up"):
            if coordinate is not None:
                raise ToolError(f"coordinate is not accepted for {action=}.")
            command_parts = [
                self.xdotool,
                f"{'mousedown' if action == 'left_mouse_down' else 'mouseup'} 1",
            ]
            return await self.shell(" ".join(command_parts))
        if action == "scroll":
            if scroll_direction is None or scroll_direction not in get_args(
                ScrollDirection
            ):
                raise ToolError(
                    f"{scroll_direction=} must be 'up', 'down', 'left', or 'right'"
                )
            if not isinstance(scroll_amount, int) or scroll_amount < 0:
                raise ToolError(f"{scroll_amount=} must be a non-negative int")
            mouse_move_part = ""
            if coordinate is not None:
                x, y = self.validate_and_get_coordinates(coordinate)
                mouse_move_part = f"mousemove --sync {x} {y}"
            scroll_button = {
                "up": 4,
                "down": 5,
                "left": 6,
                "right": 7,
            }[scroll_direction]

            command_parts = [self.xdotool, mouse_move_part]
            if text:
                command_parts.append(f"keydown {text}")
            command_parts.append(f"click --repeat {scroll_amount} {scroll_button}")
            if text:
                command_parts.append(f"keyup {text}")

            return await self.shell(" ".join(command_parts))

        if action in ("hold_key", "wait"):
            if duration is None or not isinstance(duration, (int, float)):
                raise ToolError(f"{duration=} must be a number")
            if duration < 0:
                raise ToolError(f"{duration=} must be non-negative")
            if duration > 100:
                raise ToolError(f"{duration=} is too long.")

            if action == "hold_key":
                if text is None:
                    raise ToolError(f"text is required for {action}")
                escaped_keys = shlex.quote(text)
                command_parts = [
                    self.xdotool,
                    f"keydown {escaped_keys}",
                    f"sleep {duration}",
                    f"keyup {escaped_keys}",
                ]
                return await self.shell(" ".join(command_parts))

            if action == "wait":
                await asyncio.sleep(duration)
                return await self.screenshot()

        if action in (
            "left_click",
            "right_click",
            "double_click",
            "triple_click",
            "middle_click",
        ):
            if text is not None:
                raise ToolError(f"text is not accepted for {action}")
            mouse_move_part = ""
            if coordinate is not None:
                x, y = self.validate_and_get_coordinates(coordinate)
                mouse_move_part = f"mousemove --sync {x} {y}"

            command_parts = [self.xdotool, mouse_move_part]
            if key:
                command_parts.append(f"keydown {key}")
            command_parts.append(f"click {CLICK_BUTTONS[action]}")
            if key:
                command_parts.append(f"keyup {key}")

            return await self.shell(" ".join(command_parts))

        return await super().__call__(
            action=action,
            text=text,
            coordinate=coordinate,
            start_coordinate=start_coordinate,
            key=key,
            **kwargs,
        )


class ComputerTool20251124(ComputerTool20250124):
    api_type: Literal["computer_20251124"] = "computer_20251124"  # pyright: ignore[reportIncompatibleVariableOverride]

    @property
    def options(self) -> ComputerToolOptions:  # pyright: ignore[reportIncompatibleMethodOverride]
        return {**super().options, "enable_zoom": True}  # pyright: ignore[reportReturnType]

    async def __call__(
        self,
        *,
        action: Action_20251124,
        text: str | None = None,
        coordinate: tuple[int, int] | None = None,
        scroll_direction: ScrollDirection | None = None,
        scroll_amount: int | None = None,
        duration: int | float | None = None,
        key: str | None = None,
        region: tuple[int, int, int, int] | None = None,
        **kwargs,
    ):
        if action == "zoom":
            if (
                region is None
                or not isinstance(region, (list, tuple))
                or len(region) != 4
            ):
                raise ToolError(
                    f"{region=} must be a tuple of 4 coordinates (x0, y0, x1, y1)"
                )
            if not all(isinstance(c, int) and c >= 0 for c in region):
                raise ToolError(f"{region=} must contain non-negative integers")

            if region[2] <= region[0] or region[3] <= region[1]:
                raise ToolError(f"{region=} must satisfy x0 < x1 and y0 < y1")
            return await self.zoom(region)

        return await super().__call__(
            action=action,
            text=text,
            coordinate=coordinate,
            scroll_direction=scroll_direction,
            scroll_amount=scroll_amount,
            duration=duration,
            key=key,
            **kwargs,
        )

    async def zoom(self, region: tuple[int, int, int, int]) -> ToolResult:
        """Return a magnified capture of `region`, given in screenshot pixels.

        The region is mapped onto a physical-resolution capture (so the crop
        carries more detail than the same area of a scaled-down screenshot) and
        the crop is then resized, aspect preserved, to fit within the
        screenshot frame: one axis fills it, the other is at most as large.
        The model keeps emitting coordinates in screenshot pixels afterwards;
        the zoom image itself is never a coordinate frame.
        """
        # scale_coordinates(API, ...) also bounds-checks against the display.
        x0, y0 = self.scale_coordinates(ScalingSource.API, region[0], region[1])
        x1, y1 = self.scale_coordinates(ScalingSource.API, region[2], region[3])
        frame_w, frame_h = self.screenshot_size()

        raw_path = await self._new_output_path("zoom_raw")
        zoom_path = await self._new_output_path("zoom")
        capture = await self._capture(raw_path)
        if not raw_path.exists():
            raise ToolError(f"Failed to take screenshot for zoom: {capture.error}")
        try:
            # +repage drops the crop offset so the resize sees a plain image;
            # a bare WxH geometry resizes to fit within the frame, aspect kept.
            await run(
                f"convert {raw_path} -crop {x1 - x0}x{y1 - y0}+{x0}+{y0} +repage "
                f"-resize {frame_w}x{frame_h} {zoom_path}"
            )
            if not zoom_path.exists():
                raise ToolError("Failed to crop screenshot for zoom")
            return ToolResult(
                base64_image=base64.b64encode(zoom_path.read_bytes()).decode()
            )
        finally:
            raw_path.unlink(missing_ok=True)
            zoom_path.unlink(missing_ok=True)


# The member tools of computer_toolset_20260801, one per trained action: the
# trained action names are the member names, and the model calls each member
# as its own tool (tool_use.name is the member name). The member set is fixed
# per dated version; a later version's new members arrive only with a new
# dated type.
COMPUTER_TOOLSET_MEMBERS: frozenset[str] = frozenset(
    [
        "key",
        "hold_key",
        "type",
        "cursor_position",
        "mouse_move",
        "left_mouse_down",
        "left_mouse_up",
        "left_click",
        "left_click_drag",
        "right_click",
        "middle_click",
        "double_click",
        "triple_click",
        "scroll",
        "wait",
        "screenshot",
        "zoom",
    ]
)

KEY_REPEAT_MIN = 1
KEY_REPEAT_MAX = 100


class ComputerToolset20260801(ComputerTool20251124):
    """The computer toolset (`computer_toolset_20260801`).

    Where earlier dated versions declare a single "computer" tool and receive
    every action through its `action` parameter, the toolset is declared as
    one nameless `tools[]` entry and the model calls each action as its own
    member tool: `tool_use.name` is the member name (`left_click`,
    `screenshot`, ...), the block carries `toolset_name: "computer"`, and the
    input is that action's parameter set with no `action` discriminator.
    Every tool_result answering a member call must carry the same
    `toolset_name` (see `loop._make_api_tool_result`).

    Member semantics match the corresponding actions of
    `computer_20251124`, with three differences:

    * `zoom` is an always-declared member (enabled by default) instead of an
      `enable_zoom` opt-in flag.
    * `key` accepts an optional `repeat` count (1-100).
    * Pointer members (clicks, drags) spell a held modifier chord `text`,
      where `computer_20250124` spelled it `key`.

    A turn may carry several member tool_use blocks; the caller executes them
    sequentially in block order and stops on the first failure (see
    `loop.sampling_loop`). Coordinates stay in screenshot pixels, as before.
    """

    api_type: Literal["computer_toolset_20260801"] = "computer_toolset_20260801"  # pyright: ignore[reportIncompatibleVariableOverride]
    toolset_name: Literal["computer"] = "computer"
    member_names: frozenset[str] = COMPUTER_TOOLSET_MEMBERS

    def to_params(self):
        # The toolset entry is nameless and carries no display_* options (the
        # API rejects them): coordinates are exchanged in screenshot pixels,
        # so the display geometry travels in the screenshots themselves. Per-
        # member overrides could be declared here via "configs"; this demo
        # keeps every member at its default (all enabled).
        return cast(BetaToolUnionParam, {"type": self.api_type})

    async def __call__(
        self,
        *,
        action: Action_20251124,
        text: str | None = None,
        repeat: int | None = None,
        **kwargs,
    ):
        if action not in self.member_names:
            raise ToolError(f"Invalid member tool: {action}")

        if action == "key":
            # The member's text is a single key or +-joined chord and repeat
            # defaults to 1, so an omitted repeat takes the same quoted path
            # rather than the inherited unquoted `xdotool key -- {text}`.
            if text is None:
                raise ToolError(f"text is required for {action}")
            if kwargs.get("coordinate") is not None:
                raise ToolError(f"coordinate is not accepted for {action}")
            if repeat is None:
                repeat = KEY_REPEAT_MIN
            if (
                not isinstance(repeat, int)
                or not KEY_REPEAT_MIN <= repeat <= KEY_REPEAT_MAX
            ):
                raise ToolError(
                    f"{repeat=} must be an integer between {KEY_REPEAT_MIN} and {KEY_REPEAT_MAX}"
                )
            keys = " ".join([shlex.quote(text)] * repeat)
            return await self.shell(f"{self.xdotool} key -- {keys}")

        if action in CLICK_BUTTONS and text is not None:
            # The toolset spells the held modifier chord "text"; the
            # implementation inherited from ComputerTool20250124 spells it
            # "key" on click actions.
            return await super().__call__(action=action, key=text, **kwargs)

        if action == "left_click_drag" and text is not None:
            start_coordinate = kwargs.pop("start_coordinate", None)
            coordinate = kwargs.pop("coordinate", None)
            if start_coordinate is None or coordinate is None:
                raise ToolError(
                    f"start_coordinate and coordinate are required for {action}"
                )
            start_x, start_y = self.validate_and_get_coordinates(start_coordinate)
            end_x, end_y = self.validate_and_get_coordinates(coordinate)
            escaped_keys = shlex.quote(text)
            return await self.shell(
                f"{self.xdotool} keydown {escaped_keys} "
                f"mousemove --sync {start_x} {start_y} mousedown 1 "
                f"mousemove --sync {end_x} {end_y} mouseup 1 "
                f"keyup {escaped_keys}"
            )

        return await super().__call__(action=action, text=text, **kwargs)