from __future__ import annotations

from core.models import Bounds, Scene, SceneElement, Task
from runtime.executor import Executor, ExecutorConfig
from runtime.planner import PlannedStep, Planner
from runtime.task import TaskRunner, TaskRunnerConfig
from storage.artifact_store import ArtifactStore
from storage.task_metadata import TaskMetadataStore
from storage.timeline import Timeline
from tests.fakes import PNG_10X20


class ScriptedPlanner(Planner):
    def __init__(self, steps: list[PlannedStep]) -> None:
        self.steps = steps

    def plan(self, *, goal, scene, context):
        if len(context.step_results) >= len(self.steps):
            return []
        return [self.steps[len(context.step_results)]]

    def goal_reached(self, *, goal, scene, context):
        return len(context.step_results) >= len(self.steps) and all(
            result.success for result in context.step_results
        )


def test_task_runner_executes_loop_and_writes_timeline(tmp_path) -> None:
    scene = Scene(
        width=10,
        height=20,
        elements=[
            SceneElement(
                id="search",
                type="input",
                text="Search",
                bounds=Bounds(1, 2, 4, 4),
            )
        ],
    )
    planner = ScriptedPlanner(
        [
            PlannedStep(action="tap", description="tap search", args={"x": 3, "y": 4}),
            PlannedStep(
                action="input_text",
                description="type query",
                args={"text": "Mac mini"},
            ),
        ]
    )
    executor = Executor(
        tools={
            "tap": lambda **kwargs: {"ok": True, **kwargs},
            "input_text": lambda **kwargs: {"ok": True, **kwargs},
        },
        config=ExecutorConfig(max_retries=1, backoff_seconds=0),
    )
    metadata = TaskMetadataStore(tmp_path / "tasks.sqlite3")
    timeline = Timeline(ArtifactStore(tmp_path / "history"))
    task = Task(goal="open app and search", device_id="iphone-1")
    metadata.create_task(task)

    runner = TaskRunner(
        planner=planner,
        executor=executor,
        metadata_store=metadata,
        timeline=timeline,
        config=TaskRunnerConfig(max_steps=5),
        observer=lambda device_id: scene,
        screenshot_provider=lambda device_id: PNG_10X20,
    )

    result = runner.run(task)

    assert result.status == "completed"
    assert len(timeline.read(task.id)) == 2
    assert metadata.get_task(task.id)["status"] == "completed"


def test_task_runner_stops_before_the_next_planned_action() -> None:
    scene = Scene(width=10, height=20, elements=[])
    stop_requested = False
    actions: list[str] = []

    def record_action(**kwargs):
        nonlocal stop_requested
        actions.append("tap")
        stop_requested = True
        return {"ok": True}

    runner = TaskRunner(
        planner=ScriptedPlanner(
            [
                PlannedStep(action="tap", description="first", args={}),
                PlannedStep(action="tap", description="second", args={}),
            ]
        ),
        executor=Executor(
            tools={"tap": record_action},
            config=ExecutorConfig(max_retries=1, backoff_seconds=0),
        ),
        config=TaskRunnerConfig(max_steps=5),
        observer=lambda device_id: scene,
        screenshot_provider=lambda device_id: PNG_10X20,
    )

    result = runner.run(
        Task(goal="perform two actions", device_id="phone"),
        should_stop=lambda: stop_requested,
    )

    assert result.status == "failed"
    assert result.failure_reason == "execution interrupted"
    assert actions == ["tap"]