"""
Android Control Tool

uiautomator2-based Android device control via ADB host bridge.
On Mac/Windows, ADB server runs on host, Docker connects via TCP.
"""

import base64
import logging
import os
import time
from typing import Dict, Any, Optional

from .base import BaseTool, ToolMetadata, ToolParameter, ToolPermission, ToolExecutionContext, ToolResult

logger = logging.getLogger(__name__)

# Active device sessions
_device_sessions: Dict[str, Any] = {}


class AndroidTool(BaseTool):
    """Control an Android device via uiautomator2."""

    def get_metadata(self) -> ToolMetadata:
        return ToolMetadata(
            name="ANDROID",
            description="Control an Android device. Actions: connect, tap, type, swipe, screenshot, shell, app_start, app_stop, press, info.",
            category="android",
            parameters=[
                ToolParameter("action", "string", "Action to perform", required=True,
                              enum=["connect", "tap", "type", "swipe", "screenshot",
                                    "shell", "app_start", "app_stop", "press", "info",
                                    "disconnect"]),
                ToolParameter("device_id", "string", "Device serial or IP", default="auto"),
                ToolParameter("x", "integer", "X coordinate for tap"),
                ToolParameter("y", "integer", "Y coordinate for tap"),
                ToolParameter("text", "string", "Text for type action"),
                ToolParameter("command", "string", "Shell command to run"),
                ToolParameter("package", "string", "App package name"),
                ToolParameter("key", "string", "Key to press (home, back, recent, etc.)"),
                ToolParameter("start_x", "integer", "Swipe start X"),
                ToolParameter("start_y", "integer", "Swipe start Y"),
                ToolParameter("end_x", "integer", "Swipe end X"),
                ToolParameter("end_y", "integer", "Swipe end Y"),
                ToolParameter("duration", "number", "Swipe duration in seconds", default=0.5),
            ],
            permissions=[ToolPermission.EXECUTE, ToolPermission.SYSTEM],
        )

    def execute(self, params: Dict[str, Any], context: ToolExecutionContext) -> ToolResult:
        action = params.get("action", "")
        device_id = params.get("device_id", "auto")

        try:
            import uiautomator2 as u2
        except ImportError:
            return ToolResult.fail("uiautomator2 not installed")

        try:
            if action == "connect":
                return self._connect(u2, device_id, context)

            # Get existing session
            d = _device_sessions.get(device_id)
            if not d:
                # Try auto-connect
                result = self._connect(u2, device_id, context)
                if not result.success:
                    return result
                d = _device_sessions.get(device_id)
                if not d:
                    return ToolResult.fail("No device connected. Use action='connect' first.")

            if action == "tap":
                x, y = params.get("x", 0), params.get("y", 0)
                d.click(x, y)
                time.sleep(0.3)
                return ToolResult.ok(f"Tapped ({x}, {y})")

            elif action == "type":
                text = params.get("text", "")
                d.send_keys(text)
                return ToolResult.ok(f"Typed: {text[:50]}")

            elif action == "swipe":
                d.swipe(
                    params.get("start_x", 500), params.get("start_y", 1500),
                    params.get("end_x", 500), params.get("end_y", 500),
                    duration=params.get("duration", 0.5)
                )
                return ToolResult.ok("Swiped")

            elif action == "screenshot":
                screenshot_dir = os.path.join(context.workspace_path, ".screenshots")
                os.makedirs(screenshot_dir, exist_ok=True)
                img = d.screenshot()
                path = os.path.join(screenshot_dir, f"android_{int(time.time())}.png")
                img.save(path)
                import io
                buf = io.BytesIO()
                img.save(buf, format="PNG")
                b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
                return ToolResult(success=True, output=f"Screenshot saved: {path}",
                                  base64_image=b64)

            elif action == "shell":
                command = params.get("command", "")
                result = d.shell(command)
                output = result.output if hasattr(result, "output") else str(result)
                return ToolResult.ok(output)

            elif action == "app_start":
                package = params.get("package", "")
                d.app_start(package)
                return ToolResult.ok(f"Started: {package}")

            elif action == "app_stop":
                package = params.get("package", "")
                d.app_stop(package)
                return ToolResult.ok(f"Stopped: {package}")

            elif action == "press":
                key = params.get("key", "home")
                d.press(key)
                return ToolResult.ok(f"Pressed: {key}")

            elif action == "info":
                info = d.info
                device_info = d.device_info
                return ToolResult.ok(
                    f"Display: {info.get('displayWidth', '?')}x{info.get('displayHeight', '?')}\n"
                    f"SDK: {info.get('sdkInt', '?')}\n"
                    f"Product: {device_info.get('productName', '?')}\n"
                    f"Battery: {d.battery_info.get('level', '?')}%"
                )

            elif action == "disconnect":
                _device_sessions.pop(device_id, None)
                return ToolResult.ok(f"Disconnected: {device_id}")

            else:
                return ToolResult.fail(f"Unknown action: {action}")

        except Exception as e:
            logger.exception(f"Android error: {e}")
            return ToolResult.fail(f"Android error: {e}")

    def _connect(self, u2, device_id: str, context: ToolExecutionContext) -> ToolResult:
        """Connect to a device. Uses host ADB bridge on Mac/Windows."""
        try:
            adb_host = os.environ.get("ADB_HOST", "host.docker.internal")

            if device_id == "auto":
                # Connect via host ADB bridge
                d = u2.connect()
            else:
                d = u2.connect(device_id)

            _device_sessions[device_id] = d
            info = d.info
            return ToolResult.ok(
                f"Connected to device\n"
                f"Display: {info.get('displayWidth', '?')}x{info.get('displayHeight', '?')}\n"
                f"SDK: {info.get('sdkInt', '?')}"
            )
        except Exception as e:
            return ToolResult.fail(f"Failed to connect to device: {e}")


def get_all_tools():
    return [AndroidTool]
