from __future__ import annotations

from core.models import Bounds, Scene, SceneElement, Task
from runtime.ai_planner import AIPlanner
from runtime.executor import Executor, ExecutorConfig
from runtime.planner import PlannedStep, Planner
from runtime.planner_config import PlannerConfig
from runtime.task import TaskRunner, TaskRunnerConfig
from tests.fakes import PNG_10X20


class RaisingPlanner(Planner):
    def __init__(self, error: Exception) -> None:
        self.error = error
        self.calls = 0

    def plan(self, *, goal, scene, context):
        self.calls += 1
        raise self.error

    def goal_reached(self, *, goal, scene, context):
        return False


class NarrowSignaturePlanner(Planner):
    """Predates the `screenshot` parameter added to the base Planner.plan()."""

    def __init__(self) -> None:
        self.calls = 0

    def plan(self, *, goal, scene, context):
        self.calls += 1
        if context.step_results:
            return []
        return [PlannedStep(action="tap", description="tap")]

    def goal_reached(self, *, goal, scene, context):
        return bool(context.step_results)


class ScreenshotRecordingPlanner(Planner):
    def __init__(self) -> None:
        self.screenshots: list[bytes | None] = []

    def plan(self, *, goal, scene, context, screenshot=None):
        self.screenshots.append(screenshot)
        if context.step_results:
            return []
        return [PlannedStep(action="tap", description="tap")]

    def goal_reached(self, *, goal, scene, context):
        return bool(context.step_results)


def _scene() -> Scene:
    return Scene(
        width=10,
        height=20,
        elements=[SceneElement(id="send", type="button", text="Send", bounds=Bounds(1, 2, 3, 4))],
    )


def _runner(*, planner=None, planner_config=None, observer=None) -> TaskRunner:
    return TaskRunner(
        planner=planner,
        planner_config=planner_config,
        executor=Executor(
            tools={"tap": lambda **kwargs: {"ok": True}},
            config=ExecutorConfig(max_retries=1, backoff_seconds=0),
        ),
        config=TaskRunnerConfig(max_steps=5),
        observer=observer or (lambda device_id: _scene()),
        screenshot_provider=lambda device_id: PNG_10X20,
    )


def test_task_runner_marks_task_failed_when_planner_raises() -> None:
    planner = RaisingPlanner(RuntimeError("boom"))
    runner = _runner(planner=planner)

    result = runner.run(Task(goal="inspect", device_id="phone"))

    assert result.status == "failed"
    assert result.failure_reason == "RuntimeError: boom"
    assert planner.calls == 1


def test_task_runner_marks_task_failed_when_observer_raises() -> None:
    def failing_observer(device_id: str) -> Scene:
        raise RuntimeError("no device")

    runner = _runner(planner=Planner(), observer=failing_observer)

    result = runner.run(Task(goal="inspect", device_id="phone"))

    assert result.status == "failed"
    assert result.failure_reason == "RuntimeError: no device"


def test_task_runner_omits_screenshot_kwarg_for_narrow_signature_planner() -> None:
    planner = NarrowSignaturePlanner()
    runner = _runner(planner=planner)

    result = runner.run(Task(goal="inspect", device_id="phone"))

    assert result.status == "completed"
    assert planner.calls == 2


def test_task_runner_passes_screenshot_to_planner_that_declares_it() -> None:
    planner = ScreenshotRecordingPlanner()
    runner = _runner(planner=planner)

    result = runner.run(Task(goal="inspect", device_id="phone"))

    assert result.status == "completed"
    assert planner.screenshots == [PNG_10X20, PNG_10X20]


def test_task_runner_default_planner_is_stub_when_ai_planner_disabled() -> None:
    runner = _runner(planner=None, planner_config=PlannerConfig(enabled=False))

    assert type(runner.planner) is Planner


def test_task_runner_default_planner_is_ai_planner_when_enabled() -> None:
    runner = _runner(
        planner=None,
        planner_config=PlannerConfig(enabled=True, provider="anthropic", model="test-model"),
    )

    assert isinstance(runner.planner, AIPlanner)