from __future__ import annotations

from collections.abc import Iterator
from typing import Any

from core.models import Bounds, Scene, SceneElement, Task
from runtime.context import TaskContext
from runtime.executor import Executor, ExecutorConfig
from runtime.planner import PlannedStep, Planner
from runtime.task import TaskRunner, TaskRunnerConfig
from semantic.models import SemanticScene, SemanticWidget
from world.config import WorldConfig
from world.model import WorldModel


class RecordingPlanner(Planner):
    def __init__(self, steps: list[PlannedStep]) -> None:
        self.steps = steps
        self.worlds: list[Any] = []
        self.world_snapshots: list[dict[str, Any] | None] = []

    def plan(self, *, goal, scene, context, world=None):
        self.worlds.append(world)
        self.world_snapshots.append(world.to_dict() if world is not None else None)
        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
        )


class KwargRecordingPlanner(Planner):
    def __init__(self) -> None:
        self.received_world_kwarg: list[bool] = []

    def plan(self, *, goal, scene, context, **kwargs):
        self.received_world_kwarg.append("world" in kwargs)
        if context.step_results:
            return []
        return [PlannedStep(action="tap", description="tap")]

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


class NoWorldPlanner(Planner):
    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)


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


def _semantic_scene(page: str) -> SemanticScene:
    return SemanticScene(
        page=page,
        intents=["send a message"],
        widgets=[SemanticWidget(element_id="send", purpose="send message")],
    )


def _runner(
    *,
    planner: Planner,
    executor: Executor,
    scenes: list[Scene] | None = None,
    world_model: WorldModel | None = None,
    world_config: WorldConfig | None = None,
) -> TaskRunner:
    scene_iter: Iterator[Scene] = iter(scenes or [_scene(), _scene(), _scene()])

    def observer(device_id: str) -> Scene:
        return next(scene_iter, _scene())

    return TaskRunner(
        planner=planner,
        executor=executor,
        config=TaskRunnerConfig(max_steps=5),
        observer=observer,
        world_model=world_model,
        world_config=world_config,
    )


def test_planner_plan_accepts_calls_with_and_without_world() -> None:
    scene = _scene()
    planner = Planner()
    context = TaskContext(
        task_id="task",
        goal="inspect",
    )
    world_model = WorldModel(config=WorldConfig(history_size=2))

    without_world = planner.plan(goal="inspect", scene=scene, context=context)
    with_world = planner.plan(
        goal="inspect",
        scene=scene,
        context=context,
        world=world_model.state,
    )

    assert with_world == without_world


def test_task_runner_populates_world_after_step_when_enabled_by_default() -> None:
    planner = RecordingPlanner([PlannedStep(action="tap", description="tap")])
    runner = _runner(
        planner=planner,
        executor=Executor(
            tools={"tap": lambda **kwargs: {"ok": True}},
            config=ExecutorConfig(max_retries=1, backoff_seconds=0),
        ),
    )

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

    assert result.status == "completed"
    assert planner.worlds[0] is not None
    assert planner.world_snapshots[1] is not None
    assert len(planner.world_snapshots[1]["history"]) == 1


def test_task_runner_keeps_existing_planner_subclasses_without_world_working() -> None:
    planner = NoWorldPlanner()
    runner = _runner(
        planner=planner,
        executor=Executor(
            tools={"tap": lambda **kwargs: {"ok": True}},
            config=ExecutorConfig(max_retries=1, backoff_seconds=0),
        ),
    )

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

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


def test_task_runner_disabled_world_config_does_not_pass_world_kwarg() -> None:
    planner = KwargRecordingPlanner()
    runner = _runner(
        planner=planner,
        executor=Executor(
            tools={"tap": lambda **kwargs: {"ok": True}},
            config=ExecutorConfig(max_retries=1, backoff_seconds=0),
        ),
        world_config=WorldConfig(enabled=False),
    )

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

    assert result.status == "completed"
    assert planner.received_world_kwarg == [False, False]


def test_task_runner_passes_current_world_to_planner_by_second_step() -> None:
    planner = RecordingPlanner(
        [
            PlannedStep(action="tap", description="first"),
            PlannedStep(action="tap", description="second"),
        ]
    )
    runner = _runner(
        planner=planner,
        executor=Executor(
            tools={"tap": lambda **kwargs: {"ok": True}},
            config=ExecutorConfig(max_retries=1, backoff_seconds=0),
        ),
    )

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

    assert planner.worlds[1] is not None
    assert planner.world_snapshots[1] is not None
    assert planner.world_snapshots[2] is not None
    assert len(planner.world_snapshots[1]["history"]) == 1
    assert len(planner.world_snapshots[2]["history"]) == 2


def test_task_runner_world_tracks_app_page_and_history_across_steps() -> None:
    world_model = WorldModel(config=WorldConfig(history_size=3))
    planner = RecordingPlanner(
        [
            PlannedStep(
                action="launch_app",
                description="launch chat",
                args={"app_id": "com.example.chat"},
            ),
            PlannedStep(action="tap", description="open chat"),
        ]
    )

    def launch_app(**kwargs):
        return {"ok": True, "semantic_scene": _semantic_scene("Home")}

    def tap(**kwargs):
        return {"ok": True, "semantic_scene": _semantic_scene("Chat")}

    runner = _runner(
        planner=planner,
        executor=Executor(
            tools={"launch_app": launch_app, "tap": tap},
            config=ExecutorConfig(max_retries=1, backoff_seconds=0),
        ),
        world_model=world_model,
    )

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

    assert result.status == "completed"
    assert world_model.state.current_app == "com.example.chat"
    assert world_model.state.current_page == "Chat"
    assert [event.action for event in world_model.state.history] == [
        "launch_app",
        "tap",
    ]
    assert planner.world_snapshots[1]["current_page"] == "Home"
    assert planner.world_snapshots[2]["current_page"] == "Chat"


def test_task_runner_world_tracking_failures_do_not_interrupt_loop() -> None:
    world_model = WorldModel(config=WorldConfig(history_size=3))
    planner = RecordingPlanner(
        [
            PlannedStep(action="launch_app", description="missing app args"),
            PlannedStep(
                action="tap",
                description="bad remember",
                args={"remember": "not-a-dict"},
            ),
        ]
    )
    runner = _runner(
        planner=planner,
        executor=Executor(
            tools={
                "launch_app": lambda **kwargs: {"ok": True},
                "tap": lambda **kwargs: {"ok": True},
            },
            config=ExecutorConfig(max_retries=1, backoff_seconds=0),
        ),
        world_model=world_model,
    )

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

    assert result.status == "completed"
    assert world_model.state.current_app is None
    assert world_model.state.current_page is None
    assert world_model.state.variables == {}
    assert len(world_model.state.history) == 2