"""Computer control tool for VNC desktop interaction.

VNC-based implementation of Anthropic's computer_20251124 native tool.

CRITICAL vncdotool syntax:
- Use :: (double colon) for port numbers: '172.17.0.1::5900'
- Single colon means display number (port = display + 5900)
- vncdotool API is synchronous, wrapped in asyncio.to_thread()
"""

import asyncio
import base64
import tempfile
from pathlib import Path
from typing import Literal, Any

from loguru import logger

try:
    from vncdotool import api as vnc_api
except ImportError:
    vnc_api = None

from nanobot.agent.tools.anthropic.base import BaseAnthropicTool, ToolResult


class ComputerTool20251124(BaseAnthropicTool):
    """Computer control via VNC for desktop interaction.

    Supports keyboard input, mouse control, and screenshots.
    """

    api_type: Literal["computer_20251124"] = "computer_20251124"
    name: Literal["computer"] = "computer"
    beta_flag: str = "computer-use-2025-11-24"

    def __init__(
        self,
        vnc_host: str = "172.17.0.1",
        vnc_port: int = 5900,
        vnc_username: str = "deckedmoth",
        vnc_password: str = "123",
        display_width_px: int = 1024,
        display_height_px: int = 768,
    ):
        """Initialize computer tool.

        Args:
            vnc_host: VNC server hostname/IP
            vnc_port: VNC server port
            vnc_username: VNC username (if required)
            vnc_password: VNC password (if required)
            display_width_px: Display width for screenshots
            display_height_px: Display height for screenshots
        """
        if vnc_api is None:
            raise ImportError(
                "vncdotool is required for computer tool. "
                "Install with: pip install vncdotool"
            )

        self.vnc_host = vnc_host
        self.vnc_port = vnc_port
        self.vnc_username = vnc_username
        self.vnc_password = vnc_password
        self.display_width_px = display_width_px
        self.display_height_px = display_height_px

    def to_params(self):
        """Return tool definition for API.

        NOTE: display_width_px, display_height_px, and enable_zoom are NOT
        valid parameters for computer_20251124 and cause API hangs if sent.
        """
        return {
            "type": self.api_type,
            "name": self.name,
        }

    async def __call__(
        self,
        action: Literal[
            # Basic actions
            "key", "type", "mouse_move", "screenshot", "cursor_position",
            # Click actions
            "left_click", "right_click", "middle_click", "double_click", "triple_click",
            # Advanced mouse
            "left_mouse_down", "left_mouse_up", "left_click_drag",
            # Scroll
            "scroll",
            # Advanced keyboard
            "hold_key", "paste",  # paste bypasses keyboard layout issues
            # Utility
            "wait",
            # Zoom (computer_20251124)
            "zoom"
        ] | None = None,
        coordinate: list[int] | None = None,
        text: str | None = None,
        # Additional parameters for specific actions
        start_coordinate: list[int] | None = None,  # For left_click_drag
        scroll_direction: Literal["up", "down", "left", "right"] | None = None,  # For scroll
        scroll_amount: int | None = None,  # For scroll
        duration: float | None = None,  # For hold_key, wait
        region: list[int] | None = None,  # For zoom [x1, y1, x2, y2]
        key: str | None = None,  # Modifier key for clicks/scroll
        **kwargs,
    ) -> ToolResult:
        """Execute computer control action.

        Args:
            action: Action to perform
            coordinate: [x, y] coordinates for mouse actions
            text: Text to type or key name to press

        Returns:
            ToolResult with action result or screenshot
        """
        if not action:
            return ToolResult(error="No action provided")

        try:
            # Connect with correct syntax: double colon (::) for port number
            result = await asyncio.to_thread(
                self._execute_vnc_action,
                action,
                coordinate,
                text,
                start_coordinate,
                scroll_direction,
                scroll_amount,
                duration,
                region,
                key
            )
            return result
        except Exception as e:
            logger.error(f"Computer tool error: {e}")
            return ToolResult(error=str(e))

    def _execute_vnc_action(
        self,
        action: str,
        coordinate: list[int] | None,
        text: str | None,
        start_coordinate: list[int] | None,
        scroll_direction: str | None,
        scroll_amount: int | None,
        duration: float | None,
        region: list[int] | None,
        modifier_key: str | None
    ) -> ToolResult:
        """Execute VNC action in thread (vncdotool is synchronous).

        CRITICAL: vncdotool syntax requires :: (double colon) for port numbers!
        Single colon means display number: 172.17.0.1:5900 = display 5900 (port 11800)
        Double colon means port number: 172.17.0.1::5900 = port 5900
        """
        # Connect with DOUBLE colon for port
        server = f"{self.vnc_host}::{self.vnc_port}"
        client = vnc_api.connect(server, username=self.vnc_username, password=self.vnc_password)

        try:
            # Basic actions
            if action == "screenshot":
                return self._screenshot(client)
            elif action == "key":
                return self._key(client, text or "")
            elif action == "type":
                return self._type(client, text or "")
            elif action == "mouse_move":
                return self._mouse_move(client, coordinate or [0, 0])
            elif action == "cursor_position":
                return ToolResult(output="Cursor position tracking not implemented")

            # Click actions
            elif action == "left_click":
                return self._left_click(client, coordinate, modifier_key)
            elif action == "right_click":
                return self._right_click(client, coordinate, modifier_key)
            elif action == "middle_click":
                return self._middle_click(client, coordinate, modifier_key)
            elif action == "double_click":
                return self._double_click(client, coordinate, modifier_key)
            elif action == "triple_click":
                return self._triple_click(client, coordinate, modifier_key)

            # Advanced mouse
            elif action == "left_mouse_down":
                return self._left_mouse_down(client)
            elif action == "left_mouse_up":
                return self._left_mouse_up(client)
            elif action == "left_click_drag":
                return self._left_click_drag(client, start_coordinate, coordinate)

            # Scroll
            elif action == "scroll":
                return self._scroll(client, coordinate, scroll_direction, scroll_amount, modifier_key)

            # Advanced keyboard
            elif action == "hold_key":
                return self._hold_key(client, text, duration)
            elif action == "paste":
                return self._paste(client, text)

            # Utility
            elif action == "wait":
                return self._wait(duration)

            # Zoom
            elif action == "zoom":
                return self._zoom(client, region)

            else:
                return ToolResult(error=f"Unknown action: {action}")
        finally:
            client.disconnect()

    def _screenshot(self, client) -> ToolResult:
        """Capture screenshot.

        captureScreen() requires a file path, can't use BytesIO without format.
        Use temp file then read as bytes.

        IMPORTANT: VNC display may be in sleep mode. Wake it up before screenshot.
        """
        import time

        # Wake up display (move mouse + press space to wake screensaver)
        client.mouseMove(self.display_width_px // 2, self.display_height_px // 2)
        time.sleep(0.1)
        client.keyPress('space')
        time.sleep(0.5)  # Wait for display to wake

        # Request framebuffer update
        client.refreshScreen()
        time.sleep(0.5)  # Wait for framebuffer refresh

        # Capture screenshot
        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
            tmp_path = tmp.name

        client.captureScreen(tmp_path)
        png_data = Path(tmp_path).read_bytes()
        Path(tmp_path).unlink()  # Clean up

        base64_data = base64.b64encode(png_data).decode()
        return ToolResult(base64_image=base64_data)

    def _key(self, client, text: str) -> ToolResult:
        """Press a key.

        Use lowercase names from KEYMAP: 'esc', 'return', 'tab', etc.
        Single characters work directly: 'a', 'b', '1', etc.
        """
        client.keyPress(text.lower())
        return ToolResult(output=f"Pressed key: {text}")

    def _type(self, client, text: str) -> ToolResult:
        """Type text character by character."""
        for char in text:
            client.keyPress(char)
        return ToolResult(output=f"Typed: {text}")

    def _mouse_move(self, client, coordinate: list[int]) -> ToolResult:
        """Move mouse to coordinate."""
        x, y = coordinate[0], coordinate[1]
        client.mouseMove(x, y)
        return ToolResult(output=f"Moved mouse to ({x}, {y})")

    def _left_click(self, client, coordinate: list[int] | None = None, modifier_key: str | None = None) -> ToolResult:
        """Left click at coordinate (or current position)."""
        if coordinate:
            client.mouseMove(coordinate[0], coordinate[1])
        if modifier_key:
            client.keyDown(modifier_key.lower())
        client.mousePress(1)  # 1 = left button
        if modifier_key:
            client.keyUp(modifier_key.lower())
        return ToolResult(output="Left clicked")

    def _right_click(self, client, coordinate: list[int] | None = None, modifier_key: str | None = None) -> ToolResult:
        """Right click at coordinate (or current position)."""
        if coordinate:
            client.mouseMove(coordinate[0], coordinate[1])
        if modifier_key:
            client.keyDown(modifier_key.lower())
        client.mousePress(3)  # 3 = right button
        if modifier_key:
            client.keyUp(modifier_key.lower())
        return ToolResult(output="Right clicked")

    def _middle_click(self, client, coordinate: list[int] | None = None, modifier_key: str | None = None) -> ToolResult:
        """Middle click at coordinate (or current position)."""
        if coordinate:
            client.mouseMove(coordinate[0], coordinate[1])
        if modifier_key:
            client.keyDown(modifier_key.lower())
        client.mousePress(2)  # 2 = middle button
        if modifier_key:
            client.keyUp(modifier_key.lower())
        return ToolResult(output="Middle clicked")

    def _double_click(self, client, coordinate: list[int] | None = None, modifier_key: str | None = None) -> ToolResult:
        """Double click at coordinate (or current position)."""
        if coordinate:
            client.mouseMove(coordinate[0], coordinate[1])
        if modifier_key:
            client.keyDown(modifier_key.lower())
        client.mousePress(1)
        import time
        time.sleep(0.01)  # 10ms delay between clicks
        client.mousePress(1)
        if modifier_key:
            client.keyUp(modifier_key.lower())
        return ToolResult(output="Double clicked")

    def _triple_click(self, client, coordinate: list[int] | None = None, modifier_key: str | None = None) -> ToolResult:
        """Triple click at coordinate (or current position)."""
        if coordinate:
            client.mouseMove(coordinate[0], coordinate[1])
        if modifier_key:
            client.keyDown(modifier_key.lower())
        import time
        for _ in range(3):
            client.mousePress(1)
            time.sleep(0.01)  # 10ms delay between clicks
        if modifier_key:
            client.keyUp(modifier_key.lower())
        return ToolResult(output="Triple clicked")

    def _left_mouse_down(self, client) -> ToolResult:
        """Press and hold left mouse button."""
        client.mouseDown(1)
        return ToolResult(output="Left mouse button down")

    def _left_mouse_up(self, client) -> ToolResult:
        """Release left mouse button."""
        client.mouseUp(1)
        return ToolResult(output="Left mouse button up")

    def _left_click_drag(self, client, start_coordinate: list[int] | None, end_coordinate: list[int] | None) -> ToolResult:
        """Drag from start to end coordinate."""
        if not start_coordinate or not end_coordinate:
            return ToolResult(error="Both start_coordinate and coordinate required for left_click_drag")

        start_x, start_y = start_coordinate[0], start_coordinate[1]
        end_x, end_y = end_coordinate[0], end_coordinate[1]

        client.mouseMove(start_x, start_y)
        client.mouseDown(1)
        client.mouseDrag(end_x, end_y)  # vncdotool's mouseDrag method
        client.mouseUp(1)
        return ToolResult(output=f"Dragged from ({start_x}, {start_y}) to ({end_x}, {end_y})")

    def _scroll(
        self,
        client,
        coordinate: list[int] | None,
        scroll_direction: str | None,
        scroll_amount: int | None,
        modifier_key: str | None
    ) -> ToolResult:
        """Scroll in specified direction."""
        if not scroll_direction or scroll_direction not in ("up", "down", "left", "right"):
            return ToolResult(error=f"scroll_direction must be 'up', 'down', 'left', or 'right'")

        amount = scroll_amount or 5  # Default scroll amount

        # Move to coordinate if specified
        if coordinate:
            client.mouseMove(coordinate[0], coordinate[1])

        # VNC scroll buttons: 4=up, 5=down, 6=left, 7=right
        scroll_button = {"up": 4, "down": 5, "left": 6, "right": 7}[scroll_direction]

        # Hold modifier key if specified
        if modifier_key:
            client.keyDown(modifier_key.lower())

        # Scroll by pressing scroll button multiple times
        import time
        for _ in range(amount):
            client.mousePress(scroll_button)
            time.sleep(0.05)  # Small delay between scroll events

        if modifier_key:
            client.keyUp(modifier_key.lower())

        return ToolResult(output=f"Scrolled {scroll_direction} {amount} times")

    def _hold_key(self, client, text: str | None, duration: float | None) -> ToolResult:
        """Hold a key for specified duration."""
        if not text:
            return ToolResult(error="text (key name) required for hold_key")

        hold_duration = duration or 1.0  # Default 1 second
        if hold_duration < 0 or hold_duration > 100:
            return ToolResult(error="duration must be between 0 and 100 seconds")

        import time
        client.keyDown(text.lower())
        time.sleep(hold_duration)
        client.keyUp(text.lower())

        return ToolResult(output=f"Held key '{text}' for {hold_duration}s")

    def _paste(self, client, text: str | None) -> ToolResult:
        """Paste text via clipboard (bypasses keyboard layout issues).

        This uses VNC clipboard to send text, avoiding keyboard layout mismatches
        where characters like ':' become ';' due to different keyboard mappings.
        """
        if not text:
            return ToolResult(error="text required for paste")

        # Send text via clipboard and trigger paste
        client.paste(text)
        return ToolResult(output=f"Pasted via clipboard: {text[:50]}{'...' if len(text) > 50 else ''}")

    def _wait(self, duration: float | None) -> ToolResult:
        """Wait for specified duration."""
        wait_duration = duration or 1.0
        if wait_duration < 0 or wait_duration > 100:
            return ToolResult(error="duration must be between 0 and 100 seconds")

        import time
        time.sleep(wait_duration)
        return ToolResult(output=f"Waited {wait_duration}s")

    def _zoom(self, client, region: list[int] | None) -> ToolResult:
        """Zoom into specified region and capture screenshot.

        Region format: [x1, y1, x2, y2] - top-left and bottom-right corners.
        """
        if not region or len(region) != 4:
            return ToolResult(error="region must be [x1, y1, x2, y2]")

        # Take full screenshot first
        import time
        from PIL import Image

        # Wake up display
        client.mouseMove(self.display_width_px // 2, self.display_height_px // 2)
        time.sleep(0.1)
        client.keyPress('space')
        time.sleep(0.5)
        client.refreshScreen()
        time.sleep(0.5)

        # Capture screenshot
        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
            tmp_path = tmp.name

        client.captureScreen(tmp_path)

        # Crop to region
        img = Image.open(tmp_path)
        x1, y1, x2, y2 = region
        cropped = img.crop((x1, y1, x2, y2))

        # Save cropped image
        cropped_path = tmp_path.replace('.png', '_cropped.png')
        cropped.save(cropped_path)

        # Read and encode
        png_data = Path(cropped_path).read_bytes()
        Path(tmp_path).unlink()  # Clean up original
        Path(cropped_path).unlink()  # Clean up cropped

        base64_data = base64.b64encode(png_data).decode()
        return ToolResult(base64_image=base64_data)

