import pytest

from agent.core.presets import AgentPreset, register_preset
from agent.core.runner import AgentRunner, RunConfig
from agent.orchestration.models import AgentRole, CompletionPolicy
from agent.trace.models import Trace
from agent.trace.store import FileSystemTraceStore


class StubCoordinator:
    def __init__(self, snapshots):
        self.snapshots = list(snapshots)
        self.ensure_calls = []
        self.completion_calls = 0

    async def ensure_ledger(self, trace_id, mission, *, allow_create=True):
        self.ensure_calls.append((trace_id, mission))

    async def mission_completion(self, root_trace_id):
        self.completion_calls += 1
        if len(self.snapshots) > 1:
            return self.snapshots.pop(0)
        return self.snapshots[0]


class RootFirstCoordinator(StubCoordinator):
    def __init__(self):
        super().__init__([{"status": "legacy"}])
        self.root_calls = 0

    async def root_completion(self, root_trace_id):
        self.root_calls += 1
        return {"status": "completed", "root_task_id": "root-task"}


def _knowledge_off():
    from agent.tools.builtin.knowledge import KnowledgeConfig

    return KnowledgeConfig(
        enable_extraction=False,
        enable_completion_extraction=False,
        enable_injection=False,
    )


@pytest.mark.asyncio
async def test_runner_prefers_root_completion_and_keeps_legacy_fallback(tmp_path):
    root_first = RootFirstCoordinator()
    runner = AgentRunner(
        trace_store=FileSystemTraceStore(str(tmp_path / "root-first")),
        llm_call=None,
        task_coordinator=root_first,
    )
    assert (await runner._get_mission_completion("root"))["status"] == "completed"
    assert root_first.root_calls == 1
    assert root_first.completion_calls == 0

    legacy = StubCoordinator([{"status": "blocked"}])
    runner.task_coordinator = legacy
    assert (await runner._get_mission_completion("root"))["status"] == "blocked"
    assert legacy.completion_calls == 1


def _register_role_preset(name, role, max_iterations):
    register_preset(
        name,
        AgentPreset(
            role=role,
            allowed_tools=[],
            max_iterations=max_iterations,
            skills=[],
        ),
    )


def _explicit_config(name, max_iterations):
    root_task_spec = {
        "objective": "complete the mission",
        "acceptance_criteria": [{"description": "mission result is validated"}],
    }
    return RunConfig(
        agent_type=name,
        completion_policy=CompletionPolicy.EXPLICIT_VALIDATION,
        max_iterations=max_iterations,
        tools=[],
        tool_groups=[],
        root_task_spec=root_task_spec,
        knowledge=_knowledge_off(),
    )


@pytest.mark.asyncio
async def test_planner_no_tool_exit_is_persisted_and_loop_continues(tmp_path):
    preset = "test_completion_gate_planner"
    _register_role_preset(preset, AgentRole.PLANNER, 3)
    coordinator = StubCoordinator([
        {
            "status": "needs_replan",
            "root_task_id": "root-task",
            "pending_tasks": [{"task_id": "child", "status": "pending"}],
        },
        {
            "status": "completed",
            "root_task_id": "root-task",
            "pending_tasks": [],
            "result_summary": "accepted root result",
        },
    ])
    calls = []

    async def llm_call(**kwargs):
        calls.append(kwargs["messages"])
        return {"content": "", "tool_calls": None, "finish_reason": "stop"}

    store = FileSystemTraceStore(str(tmp_path))
    runner = AgentRunner(
        trace_store=store,
        llm_call=llm_call,
        task_coordinator=coordinator,
    )
    events = [
        event
        async for event in runner.run(
            [{"role": "user", "content": "run mission"}],
            _explicit_config(preset, 3),
        )
    ]

    final = events[-1]
    assert isinstance(final, Trace)
    assert final.status == "completed"
    assert final.result_summary == "accepted root result"
    assert coordinator.ensure_calls[0][1]["objective"] == "complete the mission"
    assert len(calls) == 2
    assert any(
        message.get("role") == "user"
        and "[FRAMEWORK_MISSION_INCOMPLETE]" in message.get("content", "")
        for message in calls[1]
    )
    assert calls[1][0]["role"] == "system"
    assert "Framework role contract: Planner" in calls[1][0]["content"]
    persisted = await store.get_trace_messages(final.trace_id)
    controls = [
        message for message in persisted
        if message.role == "user"
        and "[FRAMEWORK_MISSION_INCOMPLETE]" in str(message.content)
    ]
    assert len(controls) == 1


@pytest.mark.asyncio
async def test_incomplete_planner_resume_restores_explicit_identity(tmp_path):
    preset = "test_completion_gate_resume"
    _register_role_preset(preset, AgentRole.PLANNER, 1)
    coordinator = StubCoordinator([{
        "status": "needs_replan",
        "root_task_id": "root-task",
        "pending_tasks": [{"task_id": "root-task", "status": "needs_replan"}],
    }])

    async def llm_call(**_kwargs):
        return {"content": "", "tool_calls": None, "finish_reason": "stop"}

    store = FileSystemTraceStore(str(tmp_path))
    runner = AgentRunner(
        trace_store=store,
        llm_call=llm_call,
        task_coordinator=coordinator,
    )
    first = await runner.run_result(
        [{"role": "user", "content": "run mission"}],
        _explicit_config(preset, 1),
    )
    assert first["status"] == "incomplete"

    resume_config = RunConfig(
        trace_id=first["trace_id"],
        knowledge=_knowledge_off(),
    )
    resumed = await runner.run_result([], resume_config)

    assert resumed["status"] == "incomplete"
    assert resumed["error"] == "iteration_limit_reached"
    assert resume_config.agent_type == preset
    assert resume_config.completion_policy == CompletionPolicy.EXPLICIT_VALIDATION


@pytest.mark.asyncio
async def test_explicit_planner_rewind_and_wrong_role_resume_are_rejected_before_mutation(
    tmp_path,
):
    planner_preset = "test_completion_gate_rewind_planner"
    worker_preset = "test_completion_gate_wrong_role_worker"
    _register_role_preset(planner_preset, AgentRole.PLANNER, 2)
    _register_role_preset(worker_preset, AgentRole.WORKER, 2)
    store = FileSystemTraceStore(str(tmp_path))
    planner_trace = Trace(
        trace_id="planner-trace",
        mode="agent",
        task="mission",
        agent_type=planner_preset,
        agent_role=AgentRole.PLANNER.value,
        status="completed",
        head_sequence=0,
        last_sequence=10,
    )
    worker_trace = Trace(
        trace_id="worker-trace",
        mode="agent",
        task="worker",
        agent_type=worker_preset,
        agent_role=AgentRole.WORKER.value,
        status="completed",
    )
    await store.create_trace(planner_trace)
    await store.create_trace(worker_trace)
    coordinator = StubCoordinator([{"status": "completed"}])
    runner = AgentRunner(
        trace_store=store,
        llm_call=lambda **_kwargs: None,
        task_coordinator=coordinator,
    )

    with pytest.raises(ValueError, match="do not support rewind"):
        await runner.run_result(
            [],
            RunConfig(
                trace_id=planner_trace.trace_id,
                agent_type=planner_preset,
                completion_policy=CompletionPolicy.EXPLICIT_VALIDATION,
                after_sequence=5,
            ),
        )
    with pytest.raises(ValueError, match="Trace role mismatch"):
        await runner.run_result(
            [],
            RunConfig(
                trace_id=worker_trace.trace_id,
                agent_type=planner_preset,
                completion_policy=CompletionPolicy.EXPLICIT_VALIDATION,
            ),
        )

    assert (await store.get_trace(planner_trace.trace_id)).status == "completed"
    assert (await store.get_trace(worker_trace.trace_id)).status == "completed"
    assert coordinator.ensure_calls == []


@pytest.mark.asyncio
async def test_planner_iteration_exhaustion_is_incomplete(tmp_path):
    preset = "test_completion_gate_exhaustion"
    _register_role_preset(preset, AgentRole.PLANNER, 2)
    coordinator = StubCoordinator([{
        "status": "waiting_children",
        "root_task_id": "root-task",
        "pending_tasks": ["child"],
    }])

    async def llm_call(**_kwargs):
        return {"content": "", "tool_calls": None, "finish_reason": "stop"}

    runner = AgentRunner(
        trace_store=FileSystemTraceStore(str(tmp_path)),
        llm_call=llm_call,
        task_coordinator=coordinator,
    )
    result = await runner.run_result(
        [{"role": "user", "content": "run mission"}],
        _explicit_config(preset, 2),
    )

    assert result["status"] == "incomplete"
    assert result["error"] == "iteration_limit_reached"
    assert result["summary"] is None


@pytest.mark.asyncio
async def test_blocked_root_stops_planner_as_incomplete(tmp_path):
    preset = "test_completion_gate_blocked"
    _register_role_preset(preset, AgentRole.PLANNER, 3)
    coordinator = StubCoordinator([{
        "status": "blocked",
        "root_task_id": "root-task",
        "blocked_reason": "external dependency",
        "pending_tasks": [],
    }])
    calls = 0

    async def llm_call(**_kwargs):
        nonlocal calls
        calls += 1
        return {"content": "blocked", "tool_calls": None, "finish_reason": "stop"}

    runner = AgentRunner(
        trace_store=FileSystemTraceStore(str(tmp_path)),
        llm_call=llm_call,
        task_coordinator=coordinator,
    )
    result = await runner.run_result(
        [{"role": "user", "content": "run mission"}],
        _explicit_config(preset, 3),
    )

    assert calls == 1
    assert result["status"] == "incomplete"
    assert result["error"] == "mission_blocked"


@pytest.mark.asyncio
@pytest.mark.parametrize("role", [AgentRole.WORKER, AgentRole.VALIDATOR])
async def test_explicit_subtrace_without_terminal_submit_is_protocol_failure(
    tmp_path, role
):
    preset = "test_completion_gate_" + role.value
    _register_role_preset(preset, role, 2)

    async def llm_call(**_kwargs):
        return {"content": "done", "tool_calls": None, "finish_reason": "stop"}

    runner = AgentRunner(
        trace_store=FileSystemTraceStore(str(tmp_path / role.value)),
        llm_call=llm_call,
        task_coordinator=StubCoordinator([{"status": "completed"}]),
    )
    result = await runner.run_result(
        [{"role": "user", "content": "sub task"}],
        _explicit_config(preset, 2),
    )

    assert result["status"] == "failed"
    assert result["error"] == "protocol_violation"