from __future__ import annotations

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

from core.device_manager import DeviceManager
from core.errors import ElementNotFoundError
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.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),
        "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