from __future__ import annotations

from typing import Any

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 runtime.tool_calling_client import ToolCallDecision
from storage.artifact_store import ArtifactStore
from storage.timeline import Timeline
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)


# ---------------------------------------------------------------------------
# D9: per-step prompt recording
# ---------------------------------------------------------------------------


class ScriptedToolCallingClient:
    """Returns a sequence of decisions, capturing the actual user_prompt each call."""

    def __init__(self, decisions: list[ToolCallDecision]) -> None:
        self._decisions = list(decisions)
        self.calls: list[dict[str, Any]] = []

    def decide(
        self,
        *,
        system_prompt: str,
        user_prompt: str,
        screenshot: bytes | None,
        tools: list[Any],
        timeout: float,
    ) -> ToolCallDecision:
        index = len(self.calls)
        self.calls.append(
            {
                "system_prompt": system_prompt,
                "user_prompt": user_prompt,
            }
        )
        return self._decisions[index]


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


def test_multi_step_timeline_records_actual_per_step_prompts(tmp_path) -> None:
    """When AIPlanner is used, each timeline step's prompt is the real
    per-step user prompt (containing scene JSON), not the bare task goal."""
    decisions = [
        ToolCallDecision(tool_name="tap", arguments={"x": 1, "y": 2}),
        ToolCallDecision(tool_name="finish_task", arguments={"success": True}),
    ]
    client = ScriptedToolCallingClient(decisions)
    planner = AIPlanner(client=client)
    executor = Executor(
        tools={"tap": lambda **kwargs: {"ok": True, **kwargs}},
        config=ExecutorConfig(max_retries=1, backoff_seconds=0),
    )
    timeline = Timeline(ArtifactStore(tmp_path / "history"))
    task = Task(goal="tap the button", device_id="phone")

    runner = TaskRunner(
        planner=planner,
        executor=executor,
        timeline=timeline,
        config=TaskRunnerConfig(max_steps=5),
        observer=lambda device_id: _multi_step_scene("Send"),
        screenshot_provider=lambda device_id: PNG_10X20,
    )

    runner.run(task)

    records = timeline.read(task.id)
    # Step 1 was recorded (step 2 was finish_task, which returns empty plan
    # and completes the task without a timeline append).
    assert len(records) == 1
    prompt = records[0]["prompt"]
    # The per-step prompt is NOT the bare task goal.
    assert prompt != "tap the button"
    # It contains scene-specific content that only the real planner_user_prompt
    # would include.
    assert "Current Scene (JSON)" in prompt
    assert "Call exactly one tool" in prompt
    assert "tap the button" in prompt


def test_non_ai_planner_falls_back_to_task_goal_for_prompt(tmp_path) -> None:
    """A non-LLM planner (no step.prompt) keeps recording task.goal as the
    timeline prompt — backward compat with pre-D9 behavior."""
    scene = _multi_step_scene("Search")
    planner = ScriptedPlannerForTimeline(
        [PlannedStep(action="tap", description="tap", args={"x": 1, "y": 2})]
    )
    executor = Executor(
        tools={"tap": lambda **kwargs: {"ok": True, **kwargs}},
        config=ExecutorConfig(max_retries=1, backoff_seconds=0),
    )
    timeline = Timeline(ArtifactStore(tmp_path / "history"))
    task = Task(goal="search for something", device_id="phone")

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

    runner.run(task)

    records = timeline.read(task.id)
    assert len(records) >= 1
    # Non-AI planner: prompt falls back to task goal.
    assert records[0]["prompt"] == "search for something"


class ScriptedPlannerForTimeline(Planner):
    """Simple planner that returns a fixed list of steps then signals done."""

    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(
            r.success for r in context.step_results
        )