"""Contract tests for InMemoryTaskDispatcher."""

from __future__ import annotations

import asyncio

import pytest

from cyberbrain26.execution_plane.secure_gateway import (
    AuthorizedExecutionService,
    ExecutionRefused,
)
from cyberbrain26.extension_layer.agent_registry import Agent, AgentRegistry, TaskDispatcher
from cyberbrain26.extension_layer.registry_impl import InMemoryAgentRegistry, InMemoryTaskDispatcher
from cyberbrain26.security_plane.approval import ApprovalManager
from cyberbrain26.security_plane.identity import AgentIdentity, PermissionScope, TrustLevel
from cyberbrain26.security_plane.permission import PermissionGateway, SecurityError
from cyberbrain26.security_plane.policy import PolicyEngine, PolicyMode


class MockAgent:
    """Mock agent implementing Agent protocol for testing."""

    def __init__(
        self,
        role: str,
        agent_id: str | None = None,
        identity: AgentIdentity | None = None,
        permission_gateway: PermissionGateway | None = None,
    ) -> None:
        self._role = role
        self._agent_id = agent_id or f"agent-{role}"
        self.identity = identity
        self.permission_gateway = permission_gateway
        self.executed_payloads: list[dict] = []

    @property
    def role(self) -> str:
        return self._role

    @property
    def agent_id(self) -> str:
        return self._agent_id

    async def execute(self, payload: dict) -> dict:
        self.executed_payloads.append(payload)
        return {"status": "ok", "payload": payload}

    def check_permission(self, action: str, resource: str | None = None) -> bool:
        if self.identity is None or self.permission_gateway is None:
            return False
        result = self.permission_gateway.evaluate(self.identity, action, resource)
        return result.decision.value == "allow"


def test_agent_registry_protocol_compliance() -> None:
    """InMemoryAgentRegistry must implement AgentRegistry protocol."""
    registry: AgentRegistry = InMemoryAgentRegistry()
    agent = MockAgent("coder")
    registry.register("agent-1", agent)
    assert registry.get("agent-1") is agent
    assert registry.exists("agent-1") is True
    assert registry.unregister("agent-1") is True
    assert registry.exists("agent-1") is False


def test_task_dispatcher_protocol_compliance() -> None:
    """InMemoryTaskDispatcher must implement TaskDispatcher protocol."""
    dispatcher: TaskDispatcher = InMemoryTaskDispatcher()
    agent = MockAgent("coder")
    dispatcher.register_agent("agent-1", agent)
    # Protocol compliance verified by type checker


@pytest.mark.asyncio
async def test_dispatch_uses_task_id() -> None:
    """dispatch must accept and track task_id."""
    dispatcher = InMemoryTaskDispatcher()
    identity = AgentIdentity("agent-1", "coder", trust_level=TrustLevel.TRUSTED)
    agent = MockAgent("coder", identity=identity)
    dispatcher.register_agent("agent-1", agent)

    result = await dispatcher.dispatch("task-123", "agent-1", {"action": "plan"})
    assert result["status"] == "ok"
    assert len(agent.executed_payloads) == 1
    assert agent.executed_payloads[0]["action"] == "plan"


@pytest.mark.asyncio
async def test_dispatch_to_role_deterministic_round_robin() -> None:
    """dispatch_to_role must select agents deterministically (round-robin)."""
    dispatcher = InMemoryTaskDispatcher()
    agent1 = MockAgent(
        "coder", "agent-1", identity=AgentIdentity("agent-1", "coder", trust_level=TrustLevel.TRUSTED)
    )
    agent2 = MockAgent(
        "coder", "agent-2", identity=AgentIdentity("agent-2", "coder", trust_level=TrustLevel.TRUSTED)
    )
    dispatcher.register_agent("agent-1", agent1)
    dispatcher.register_agent("agent-2", agent2)

    # First call -> agent1
    await dispatcher.dispatch_to_role("task-1", "coder", {"action": "plan"})
    assert len(agent1.executed_payloads) == 1
    assert len(agent2.executed_payloads) == 0

    # Second call -> agent2
    await dispatcher.dispatch_to_role("task-2", "coder", {"action": "plan"})
    assert len(agent1.executed_payloads) == 1
    assert len(agent2.executed_payloads) == 1

    # Third call -> agent1 (round-robin)
    await dispatcher.dispatch_to_role("task-3", "coder", {"action": "plan"})
    assert len(agent1.executed_payloads) == 2
    assert len(agent2.executed_payloads) == 1


@pytest.mark.asyncio
async def test_dispatch_security_boundary_identity_policy_permission() -> None:
    """dispatch must enforce Identity → Policy → Permission before execute."""
    policy_engine = PolicyEngine(mode=PolicyMode.SAFE)
    gateway = PermissionGateway(policy_engine)

    identity = AgentIdentity(
        agent_id="secure-agent",
        role="coder",
        trust_level=TrustLevel.TRUSTED,
        scopes={PermissionScope.READ_FILE},
    )

    dispatcher = InMemoryTaskDispatcher(policy_engine=policy_engine)
    agent = MockAgent("coder", "secure-agent", identity=identity, permission_gateway=gateway)
    dispatcher.register_agent("secure-agent", agent)

    # Allowed action (has READ_FILE scope, non-sensitive)
    result = await dispatcher.dispatch("task-1", "secure-agent", {"action": "read_file"})
    assert result["status"] == "ok"

    # Denied action (lacks WRITE_FILE scope)
    with pytest.raises(SecurityError):
        await dispatcher.dispatch("task-2", "secure-agent", {"action": "write_file"})


@pytest.mark.asyncio
async def test_dispatch_records_audit_event() -> None:
    """dispatch must record audit events when audit_log is provided."""
    from cyberbrain26.security_plane.audit import AuditEventType, AuditLog, AuditQuery

    audit = AuditLog()
    policy_engine = PolicyEngine(mode=PolicyMode.SAFE)
    dispatcher = InMemoryTaskDispatcher(policy_engine=policy_engine, audit_log=audit)

    identity = AgentIdentity(
        agent_id="audit-agent",
        role="coder",
        trust_level=TrustLevel.TRUSTED,
        scopes={PermissionScope.READ_FILE},
    )
    agent = MockAgent("coder", "audit-agent", identity=identity)
    dispatcher.register_agent("audit-agent", agent)

    await dispatcher.dispatch("task-1", "audit-agent", {"action": "read_file"})

    perm_events = audit.query(AuditQuery(event_type=AuditEventType.PERMISSION_CHECK))
    assert len(perm_events) >= 1
    assert perm_events[0].action == "read_file"


@pytest.mark.asyncio
async def test_dispatch_to_role_security_boundary() -> None:
    """dispatch_to_role must enforce security boundary."""
    policy_engine = PolicyEngine(mode=PolicyMode.SAFE)
    gateway = PermissionGateway(policy_engine)

    identity = AgentIdentity(
        agent_id="secure-agent-1",
        role="coder",
        trust_level=TrustLevel.TRUSTED,
        scopes={PermissionScope.SHELL},
    )

    dispatcher = InMemoryTaskDispatcher(policy_engine=policy_engine)
    agent = MockAgent("coder", "secure-agent-1", identity=identity, permission_gateway=gateway)
    dispatcher.register_agent("secure-agent-1", agent)

    # Sensitive action requires approval in SAFE mode
    with pytest.raises(SecurityError):
        await dispatcher.dispatch_to_role("task-1", "coder", {"action": "execute_shell"})


@pytest.mark.asyncio
async def test_dispatch_unregistered_agent_raises() -> None:
    """dispatch must raise KeyError for unregistered agent."""
    dispatcher = InMemoryTaskDispatcher()
    with pytest.raises(KeyError, match="agent not found"):
        await dispatcher.dispatch("task-1", "unknown", {"action": "test"})


@pytest.mark.asyncio
async def test_dispatch_to_role_no_agent_raises() -> None:
    """dispatch_to_role must raise KeyError for unknown role."""
    dispatcher = InMemoryTaskDispatcher()
    with pytest.raises(KeyError, match="no agent with role"):
        await dispatcher.dispatch_to_role("task-1", "unknown-role", {"action": "test"})


@pytest.mark.asyncio
async def test_dispatch_without_action_is_refused() -> None:
    """A payload with no declared action must be refused, not defaulted."""
    dispatcher = InMemoryTaskDispatcher()
    identity = AgentIdentity("agent-1", "coder", trust_level=TrustLevel.TRUSTED)
    agent = MockAgent("coder", identity=identity)
    dispatcher.register_agent("agent-1", agent)

    with pytest.raises(SecurityError, match="no default action is inferred"):
        await dispatcher.dispatch("task-1", "agent-1", {"payload": "anything"})
    assert agent.executed_payloads == []


@pytest.mark.asyncio
async def test_dispatch_to_role_without_action_is_refused() -> None:
    """dispatch_to_role must not invent an action either."""
    dispatcher = InMemoryTaskDispatcher()
    identity = AgentIdentity("agent-1", "coder", trust_level=TrustLevel.TRUSTED)
    agent = MockAgent("coder", identity=identity)
    dispatcher.register_agent("agent-1", agent)

    with pytest.raises(SecurityError, match="no default action is inferred"):
        await dispatcher.dispatch_to_role("task-1", "coder", {})
    assert agent.executed_payloads == []


@pytest.mark.asyncio
async def test_dispatch_unmapped_action_is_refused() -> None:
    """An action outside the canonical vocabulary must be denied."""
    dispatcher = InMemoryTaskDispatcher()
    identity = AgentIdentity(
        "agent-1",
        "coder",
        trust_level=TrustLevel.TRUSTED,
        scopes={PermissionScope.READ_FILE, PermissionScope.SHELL, PermissionScope.PYTHON},
    )
    agent = MockAgent("coder", identity=identity)
    dispatcher.register_agent("agent-1", agent)

    # Every scope the agent holds is still not a statement about "execute":
    # the action is unregistered, so it is denied before any gate is consulted.
    with pytest.raises(SecurityError):
        await dispatcher.dispatch("task-1", "agent-1", {"action": "execute"})
    assert agent.executed_payloads == []


@pytest.mark.asyncio
async def test_agent_without_identity_is_denied() -> None:
    """An agent without identity must fail closed before execution."""
    dispatcher = InMemoryTaskDispatcher()
    agent = MockAgent("coder")  # No identity
    dispatcher.register_agent("agent-1", agent)

    with pytest.raises(SecurityError, match="identity is required"):
        await dispatcher.dispatch("task-1", "agent-1", {"action": "execute_shell"})
    assert agent.executed_payloads == []


@pytest.mark.asyncio
async def test_untrusted_agent_denied() -> None:
    """UNTRUSTED agents must be denied by trust policy."""
    policy_engine = PolicyEngine(mode=PolicyMode.SAFE)
    gateway = PermissionGateway(policy_engine)

    identity = AgentIdentity(
        agent_id="untrusted-agent",
        role="coder",
        trust_level=TrustLevel.UNTRUSTED,
        scopes={PermissionScope.READ_FILE, PermissionScope.SHELL},
    )

    dispatcher = InMemoryTaskDispatcher(policy_engine=policy_engine)
    agent = MockAgent("coder", "untrusted-agent", identity=identity, permission_gateway=gateway)
    dispatcher.register_agent("untrusted-agent", agent)

    with pytest.raises(SecurityError):
        await dispatcher.dispatch("task-1", "untrusted-agent", {"action": "read_file"})


# ----------------------------------------------------------------------
# One approval per step, and a payload that changes after it is refused
# ----------------------------------------------------------------------
class _RecordingProvider:
    """Execution provider that records whatever it is asked to run."""

    provider_name = "recording"

    def __init__(self) -> None:
        self.calls: list[tuple[object, ...]] = []

    async def execute(self, *args: object, **kwargs: object) -> dict[str, object]:
        self.calls.append(args)
        return {"return_code": 0}


class _SwappingCoder(MockAgent):
    """Coder whose body swaps the code after the dispatch approved the original.

    This models the one thing the single-approval flow must never tolerate: the
    payload that reached the operator is not the payload that reaches the
    sandbox. It is deliberately a *legitimate-looking* agent - it holds the
    right identity, the right scope and calls the real execution service.
    """

    def __init__(
        self,
        role: str,
        identity: AgentIdentity,
        service: AuthorizedExecutionService,
    ) -> None:
        super().__init__(role, agent_id=identity.agent_id, identity=identity)
        self.service = service

    async def execute(self, payload: dict) -> dict:
        payload["code"] = "print('swapped after the operator approved the step')"
        await self.service.execute(
            self.identity,
            payload["code"],
            payload["language"],
            task_id=str(payload.get("task_id", "dispatched-task")),
            workspace=payload.get("workspace"),
        )
        return {"status": "ok"}


@pytest.mark.asyncio
async def test_approval_for_the_original_payload_does_not_authorise_a_changed_payload() -> None:
    """A granted approval must not travel from the approved code to other code.

    The dispatcher obtains one human approval for the step as declared, then
    hands control to the agent. If the agent changes the payload in between, the
    inner gateway has to notice: reusing the approval would run code nobody
    approved, and asking again would hide the swap. So the execution is refused,
    the provider is never reached, and the refusal is audited.
    """
    approvals = ApprovalManager()
    identity = AgentIdentity(
        agent_id="coder-1",
        role="coder",
        trust_level=TrustLevel.TRUSTED,
        scopes={PermissionScope.PYTHON},
    )
    provider = _RecordingProvider()
    service = AuthorizedExecutionService(
        provider,  # type: ignore[arg-type]
        PermissionGateway(PolicyEngine()),
        approval_manager=approvals,
        approval_timeout_seconds=0.05,
    )
    dispatcher = InMemoryTaskDispatcher(approval_manager=approvals)
    agent = _SwappingCoder("coder", identity, service)
    dispatcher.register_agent("coder-1", agent)

    async def approve_once() -> None:
        while not approvals.list_pending():
            await asyncio.sleep(0.005)
        approvals.approve(approvals.list_pending()[0].request_id, "reviewer", "approved")

    approver = asyncio.create_task(approve_once())
    with pytest.raises(ExecutionRefused, match="does not cover this execution"):
        await dispatcher.dispatch(
            "task-1",
            "coder-1",
            {
                "action": "execute_python",
                "language": "python",
                "code": "print('approved by the reviewer')",
            },
        )
    await approver

    assert provider.calls == [], "the sandbox must not be reached with unapproved code"
    assert len(approvals.list_records()) == 1, "exactly one approval, and it was not reused"


@pytest.mark.asyncio
async def test_unchanged_payload_lets_the_second_gate_proceed_on_one_approval() -> None:
    """The same dispatch with an untouched payload needs no second approval.

    The gateway still re-runs its own scope and policy check; what it does not
    do is raise a second request the operator would have to answer again.
    """

    class _Coder(MockAgent):
        async def execute(self, payload: dict) -> dict:
            await service.execute(
                self.identity,
                payload["code"],
                payload["language"],
                task_id=str(payload.get("task_id", "dispatched-task")),
            )
            return {"status": "ok"}

    approvals = ApprovalManager()
    identity = AgentIdentity(
        agent_id="coder-1",
        role="coder",
        trust_level=TrustLevel.TRUSTED,
        scopes={PermissionScope.PYTHON},
    )
    provider = _RecordingProvider()
    service = AuthorizedExecutionService(
        provider,  # type: ignore[arg-type]
        PermissionGateway(PolicyEngine()),
        approval_manager=approvals,
        approval_timeout_seconds=0.05,
    )
    dispatcher = InMemoryTaskDispatcher(approval_manager=approvals)
    agent = _Coder("coder", agent_id="coder-1", identity=identity)
    dispatcher.register_agent("coder-1", agent)

    async def approve_once() -> None:
        while not approvals.list_pending():
            await asyncio.sleep(0.005)
        approvals.approve(approvals.list_pending()[0].request_id, "reviewer", "approved")

    approver = asyncio.create_task(approve_once())
    result = await dispatcher.dispatch(
        "task-1",
        "coder-1",
        {"action": "execute_python", "language": "python", "code": "print('approved')"},
    )
    await approver

    assert result == {"status": "ok"}
    assert len(provider.calls) == 1
    assert len(approvals.list_records()) == 1


@pytest.mark.asyncio
async def test_outer_allow_mints_no_grant_for_the_inner_gate() -> None:
    """A policy "allow" must not be laundered into an approval.

    ``plan`` needs no human decision at the outer gate. If that allow were
    turned into a grant, an agent could spend it to skip the approval the real
    execution action requires - so no grant is minted and the inner gate asks.
    """
    approvals = ApprovalManager()
    identity = AgentIdentity(
        agent_id="coder-1",
        role="coder",
        trust_level=TrustLevel.TRUSTED,
        scopes={PermissionScope.PYTHON},
    )
    provider = _RecordingProvider()
    service = AuthorizedExecutionService(
        provider,  # type: ignore[arg-type]
        PermissionGateway(PolicyEngine()),
        approval_manager=approvals,
        approval_timeout_seconds=0.05,
    )
    dispatcher = InMemoryTaskDispatcher(approval_manager=approvals)
    agent = _SwappingCoder("coder", identity, service)
    dispatcher.register_agent("coder-1", agent)

    # "code" is a known non-mutating label, so the outer chain allows it outright.
    # The agent body then tries to run shell, which is a real action that was
    # never approved: the grant, if one existed, would name a different action.
    with pytest.raises(ExecutionRefused, match="execute_shell"):
        await dispatcher.dispatch(
            "task-1",
            "coder-1",
            {"action": "code", "language": "shell", "code": "echo hi"},
        )

    assert provider.calls == []
    assert approvals.list_records() == []


class TestAgentProtocol:
    """Verify BaseAgent implements Agent protocol."""

    def test_base_agent_has_required_attributes(self) -> None:
        from cyberbrain26.agents.orchestrator import BaseAgent
        from cyberbrain26.control_plane import AsyncEventBus, MemoryTaskQueue

        class ConcreteAgent(BaseAgent):
            @property
            def name(self) -> str:
                return "Test"

            @property
            def role(self) -> str:
                return "test"

            async def execute(self, task_data: dict) -> None:
                pass

        bus = AsyncEventBus()
        queue = MemoryTaskQueue()
        agent = ConcreteAgent(bus, queue)

        # Verify protocol attributes
        assert hasattr(agent, "role")
        assert hasattr(agent, "identity")
        assert hasattr(agent, "permission_gateway")
        assert hasattr(agent, "execute")
        assert hasattr(agent, "check_permission")
        assert isinstance(agent, Agent)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])