from __future__ import annotations

from typing import Any
from unittest.mock import MagicMock, patch

from agents.collab_runner import CollaborativeTaskRunner, CollaborativeTaskRunnerConfig
from agents.config import CollaborationConfig
from agents.models import Observation, ReflectionAction, ReflectionOutcome, VerificationVerdict
from core.models import Bounds, Scene, SceneElement, Task
from runtime.executor import StepResult
from runtime.planner import PlannedStep
from runtime.task import TaskRunner, TaskRunnerConfig
from storage.artifact_store import ArtifactStore
from storage.task_metadata import TaskMetadataStore
from storage.timeline import Timeline
from world.config import WorldConfig


def _scene() -> Scene:
    return Scene(
        width=1080,
        height=1920,
        elements=[SceneElement(id="btn1", type="button", bounds=Bounds(10, 20, 100, 50), text="OK")],
    )


def _task() -> Task:
    return Task(goal="Test goal", device_id="dev1")


def _planned_step() -> PlannedStep:
    return PlannedStep(action="describe_screen", description="Observe", args={})


def _success_step_result() -> StepResult:
    return StepResult(step=_planned_step(), success=True, attempts=1, result="ok")


def _observation() -> Observation:
    return Observation(scene_summary="Screen 1080x1920")


def _achieved_verdict() -> VerificationVerdict:
    return VerificationVerdict(result="achieved", reasoning="Effect observed.")


def _not_achieved_verdict() -> VerificationVerdict:
    return VerificationVerdict(result="not_achieved", reasoning="No change.")


def _replan_outcome() -> ReflectionOutcome:
    return ReflectionOutcome(replan=True, reasoning="No recovery.")


def _recovery_outcome() -> ReflectionOutcome:
    return ReflectionOutcome(
        replan=False,
        action=ReflectionAction(action="swipe", description="Scroll", args={"direction": "up"}),
        reasoning="Try scrolling.",
    )


def test_completes_without_verification_failure() -> None:
    """Task completes normally when verifier reports achieved."""
    planner = MagicMock()
    planner.plan.return_value = [_planned_step()]
    planner.goal_reached.return_value = True

    executor = MagicMock()
    executor.execute.return_value = _success_step_result()

    observer = MagicMock()
    observer.observe.return_value = _observation()

    verifier = MagicMock()
    verifier.verify.return_value = _achieved_verdict()

    reflector = MagicMock()

    runner = CollaborativeTaskRunner(
        planner=planner,
        executor=executor,
        observer=observer,
        verifier=verifier,
        reflector=reflector,
        config=CollaborativeTaskRunnerConfig(max_steps=5),
        collaboration_config=CollaborationConfig(enabled=True, max_recovery_attempts=3),
    )

    with patch("tools.describe_screen.describe_screen", return_value=_scene()):
        task = runner.run(_task())

    assert task.status == "completed"
    reflector.reflect.assert_not_called()


def test_recovers_via_reflector_action() -> None:
    """Task recovers when reflector proposes a distinct action."""
    call_count = 0

    def verifier_side_effect(**kwargs: Any) -> VerificationVerdict:
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            return _not_achieved_verdict()
        return _achieved_verdict()

    planner = MagicMock()
    planner.plan.return_value = [_planned_step()]
    planner.goal_reached.side_effect = [False, True]

    executor = MagicMock()
    executor.execute.return_value = _success_step_result()

    observer = MagicMock()
    observer.observe.return_value = _observation()

    verifier = MagicMock()
    verifier.verify.side_effect = verifier_side_effect

    reflector = MagicMock()
    reflector.reflect.return_value = _recovery_outcome()

    runner = CollaborativeTaskRunner(
        planner=planner,
        executor=executor,
        observer=observer,
        verifier=verifier,
        reflector=reflector,
        config=CollaborativeTaskRunnerConfig(max_steps=5),
        collaboration_config=CollaborationConfig(enabled=True, max_recovery_attempts=3),
    )

    with patch("tools.describe_screen.describe_screen", return_value=_scene()):
        task = runner.run(_task())

    assert task.status == "completed"
    reflector.reflect.assert_called_once()
    executor.execute.assert_called()  # recovery action also executed


def test_exhausts_recovery_ceiling() -> None:
    """Task fails when reflection-recovery ceiling is reached."""
    planner = MagicMock()
    planner.plan.return_value = [_planned_step()]
    planner.goal_reached.return_value = False

    executor = MagicMock()
    executor.execute.return_value = _success_step_result()

    observer = MagicMock()
    observer.observe.return_value = _observation()

    verifier = MagicMock()
    verifier.verify.return_value = _not_achieved_verdict()

    reflector = MagicMock()
    reflector.reflect.return_value = _replan_outcome()

    runner = CollaborativeTaskRunner(
        planner=planner,
        executor=executor,
        observer=observer,
        verifier=verifier,
        reflector=reflector,
        config=CollaborativeTaskRunnerConfig(max_steps=10, max_recovery_attempts=2),
        collaboration_config=CollaborationConfig(enabled=True, max_recovery_attempts=2),
    )

    with patch("tools.describe_screen.describe_screen", return_value=_scene()):
        task = runner.run(_task())

    assert task.status == "failed"
    assert "ceiling" in (task.failure_reason or "").lower()


def test_collaborative_run_shares_bookkeeping_with_task_runner(tmp_path) -> None:
    """A completed collaborative task records Timeline/TaskMetadataStore/
    on_task_succeeded exactly like a plain TaskRunner.run() would for the
    same scripted single-step scenario, and populates context.world along
    the way (verified indirectly via what the Observer is called with)."""
    planner = MagicMock()
    planner.plan.return_value = [_planned_step()]
    planner.goal_reached.side_effect = [False, True]

    executor = MagicMock()
    executor.execute.return_value = _success_step_result()

    observer = MagicMock()
    observer.observe.return_value = _observation()

    verifier = MagicMock()
    verifier.verify.return_value = _achieved_verdict()

    reflector = MagicMock()

    metadata = TaskMetadataStore(tmp_path / "tasks.sqlite3")
    timeline = Timeline(ArtifactStore(tmp_path / "history"))
    on_task_succeeded = MagicMock()

    shared_task_runner = TaskRunner(
        metadata_store=metadata,
        timeline=timeline,
        on_task_succeeded=on_task_succeeded,
        world_config=WorldConfig(enabled=True),
        config=TaskRunnerConfig(max_steps=5),
    )

    runner = CollaborativeTaskRunner(
        planner=planner,
        executor=executor,
        observer=observer,
        verifier=verifier,
        reflector=reflector,
        task_runner=shared_task_runner,
        config=CollaborativeTaskRunnerConfig(max_steps=5),
        collaboration_config=CollaborationConfig(enabled=True, max_recovery_attempts=3),
    )

    task = _task()
    metadata.create_task(task)

    with patch("tools.describe_screen.describe_screen", return_value=_scene()):
        result = runner.run(task)

    assert result.status == "completed"
    assert len(timeline.read(task.id)) == 1
    assert metadata.get_task(task.id)["status"] == "completed"
    on_task_succeeded.assert_called_once_with(task.id, task.goal, timeline)

    # context.world was populated (not None) throughout the run: every
    # Observer.observe call received a non-None `world` kwarg.
    assert observer.observe.call_args_list
    for call in observer.observe.call_args_list:
        assert call.kwargs["world"] is not None


def test_second_step_verified_against_post_step_one_observation() -> None:
    """When a single Planner.plan() call returns a 2-element plan, the
    second step's Verifier.verify() call uses the first step's post
    observation as its pre_observation, not the stale pre-plan observation."""
    step_one = _planned_step()
    step_two = PlannedStep(action="tap", description="tap next", args={"x": 1, "y": 2})

    planner = MagicMock()
    planner.plan.return_value = [step_one, step_two]
    planner.goal_reached.side_effect = [False, True]

    executor = MagicMock()
    executor.execute.return_value = _success_step_result()

    obs_pre_plan = Observation(scene_summary="obs0-pre-plan")
    obs_after_step_one = Observation(scene_summary="obs1-after-step-one")
    obs_after_step_two = Observation(scene_summary="obs2-after-step-two")
    obs_next_outer_iter = Observation(scene_summary="obs3-next-outer-iter")

    observer = MagicMock()
    observer.observe.side_effect = [
        obs_pre_plan,
        obs_after_step_one,
        obs_after_step_two,
        obs_next_outer_iter,
    ]

    verifier = MagicMock()
    verifier.verify.return_value = _achieved_verdict()

    reflector = MagicMock()

    runner = CollaborativeTaskRunner(
        planner=planner,
        executor=executor,
        observer=observer,
        verifier=verifier,
        reflector=reflector,
        config=CollaborativeTaskRunnerConfig(max_steps=5),
        collaboration_config=CollaborationConfig(enabled=True, max_recovery_attempts=3),
    )

    with patch("tools.describe_screen.describe_screen", return_value=_scene()):
        task = runner.run(_task())

    assert task.status == "completed"
    reflector.reflect.assert_not_called()
    assert verifier.verify.call_count == 2

    first_call_kwargs = verifier.verify.call_args_list[0].kwargs
    second_call_kwargs = verifier.verify.call_args_list[1].kwargs

    assert first_call_kwargs["pre_observation"] is obs_pre_plan
    assert first_call_kwargs["post_observation"] is obs_after_step_one

    # The bug being guarded against: the second step must not be verified
    # against the stale pre-whole-plan observation.
    assert second_call_kwargs["pre_observation"] is obs_after_step_one
    assert second_call_kwargs["pre_observation"] is not obs_pre_plan
    assert second_call_kwargs["post_observation"] is obs_after_step_two


def test_disabled_collaboration_runs_plain() -> None:
    """When collaboration is disabled, delegates to plain TaskRunner."""
    planner = MagicMock()
    executor = MagicMock()

    mock_task_runner = MagicMock()
    expected_task = _task()
    expected_task.status = "completed"
    mock_task_runner.run.return_value = expected_task

    runner = CollaborativeTaskRunner(
        planner=planner,
        executor=executor,
        task_runner=mock_task_runner,
        collaboration_config=CollaborationConfig(enabled=False),
    )

    result = runner.run(_task())
    assert result.status == "completed"
    mock_task_runner.run.assert_called_once()