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 tests.fakes import PNG_10X20
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),
        screenshot_provider=lambda device_id: PNG_10X20,
    )

    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
    record = timeline.read(task.id)[0]
    assert record["before_screenshot_path"]
    assert record["after_screenshot_path"]
    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()