from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass, field
from time import sleep
from typing import Any

from core.errors import ElementNotFoundError
from device.manager import DeviceManager
from runtime.context import TaskContext
from runtime.planner import PlannedStep

ToolCallable = Callable[..., Any]


@dataclass(frozen=True)
class StepResult:
    step: PlannedStep
    success: bool
    attempts: int
    result: Any = None
    error: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "step": {
                "action": self.step.action,
                "description": self.step.description,
                "args": dict(self.step.args),
                "expected_text": self.step.expected_text,
            },
            "success": self.success,
            "attempts": self.attempts,
            "result": self.result,
            "error": self.error,
        }


@dataclass
class ExecutorConfig:
    max_retries: int = 3
    backoff_seconds: float = 0.25


class Executor:
    def __init__(
        self,
        *,
        tools: dict[str, ToolCallable] | None = None,
        config: ExecutorConfig | None = None,
    ) -> None:
        self.tools = tools or default_tool_registry()
        self.config = config or ExecutorConfig()

    def execute(
        self,
        step: PlannedStep,
        *,
        context: TaskContext | None = None,
    ) -> StepResult:
        if self.config.max_retries < 1:
            raise ValueError("max_retries must be at least 1")

        last_error: Exception | None = None
        for attempt in range(1, self.config.max_retries + 1):
            try:
                result = self._execute_once(step, context=context)
                return StepResult(
                    step=step,
                    success=True,
                    attempts=attempt,
                    result=result,
                )
            except Exception as exc:
                last_error = exc
                if attempt < self.config.max_retries:
                    sleep(self.config.backoff_seconds)

        return StepResult(
            step=step,
            success=False,
            attempts=self.config.max_retries,
            error=str(last_error) if last_error else "step failed",
        )

    def _execute_once(
        self,
        step: PlannedStep,
        *,
        context: TaskContext | None,
    ) -> Any:
        if step.action == "wait_for_text":
            query = step.args["query"]
            scene = context.latest_scene if context else None
            if scene is None:
                raise ElementNotFoundError("element not found")
            result = self.tools["find_text"](scene=scene, query=query)
            if not result.get("found"):
                raise ElementNotFoundError("element not found")
            return result

        tool = self.tools.get(step.action)
        if tool is None:
            raise KeyError(f"unknown tool action {step.action}")
        return tool(**step.args)


def default_tool_registry(
    *,
    manager: DeviceManager | None = None,
) -> dict[str, ToolCallable]:
    from tools.describe_screen import describe_screen
    from tools.describe_screen_semantic import describe_screen_semantic
    from tools.find_icon import find_icon, find_icon_on_screen
    from tools.find_text import find_text, find_text_on_screen
    from tools.input_text import input_text
    from tools.launch_app import launch_app, terminate_app
    from tools.screenshot import take_screenshot
    from tools.swipe import swipe
    from tools.tap import tap
    from tools.ui_tree import get_ui_tree

    return {
        "take_screenshot": _bind_manager(take_screenshot, manager),
        "screenshot": _bind_manager(take_screenshot, manager),
        "tap": _bind_manager(tap, manager),
        "swipe": _bind_manager(swipe, manager),
        "input_text": _bind_manager(input_text, manager),
        "launch_app": _bind_manager(launch_app, manager),
        "terminate_app": _bind_manager(terminate_app, manager),
        "get_ui_tree": _bind_manager(get_ui_tree, manager),
        "ui_tree": _bind_manager(get_ui_tree, manager),
        "describe_screen": _bind_manager(describe_screen, manager),
        "describe_screen_semantic": _bind_manager(describe_screen_semantic, manager),
        "find_text": find_text,
        "find_text_on_screen": _bind_manager(find_text_on_screen, manager),
        "find_icon": find_icon,
        "find_icon_on_screen": _bind_manager(find_icon_on_screen, manager),
    }


def _bind_manager(func: ToolCallable, manager: DeviceManager | None) -> ToolCallable:
    if manager is None:
        return func

    def wrapper(*args: Any, **kwargs: Any) -> Any:
        kwargs.setdefault("manager", manager)
        return func(*args, **kwargs)

    return wrapper