from __future__ import annotations

import asyncio

from beaver.skills.learning.replay import ReplayToolExecutor, ReplayToolPolicy, classify_tool_mode
from beaver.tools.base import BaseTool, ToolContext, ToolResult, ToolSpec
from beaver.tools.registry.tool_registry import ToolRegistry
from beaver.tools.runtime.executor import ToolExecutor


class FakeTool(BaseTool):
    def __init__(self, name: str, *, toolset: str = "filesystem", metadata: dict | None = None) -> None:
        self._spec = ToolSpec(
            name=name,
            description=f"{name} tool",
            input_schema={"type": "object", "properties": {"path": {"type": "string"}}},
            toolset=toolset,
            metadata=metadata or {},
        )

    @property
    def spec(self) -> ToolSpec:
        return self._spec

    async def invoke(self, arguments: dict, context: ToolContext) -> ToolResult:
        return ToolResult(success=True, content=f"executed:{arguments}", tool_name=self.spec.name)


def _executor(*tools: FakeTool) -> ReplayToolExecutor:
    registry = ToolRegistry()
    for tool in tools:
        registry.register(tool)
    return ReplayToolExecutor(ToolExecutor(registry), registry=registry, policy=ReplayToolPolicy())


def test_classify_tool_modes_from_spec() -> None:
    assert classify_tool_mode(FakeTool("read_file").spec) == "executed"
    assert classify_tool_mode(FakeTool("write_file").spec) == "executed"
    assert classify_tool_mode(FakeTool("mcp_outlook_send_email", toolset="mcp", metadata={"transport": "mcp"}).spec) == "surrogate"
    assert classify_tool_mode(FakeTool("delete_account", toolset="mcp", metadata={"transport": "mcp"}).spec) == "blocked"


def test_replay_executor_executes_safe_tool_and_records_trace() -> None:
    executor = _executor(FakeTool("write_file"))

    result = asyncio.run(executor.execute("write_file", {"path": "a.txt"}, context=ToolContext(workspace="/tmp/replay")))

    assert result.success is True
    assert result.content.startswith("executed:")
    assert executor.traces[0]["mode"] == "executed"
    assert executor.traces[0]["tool_name"] == "write_file"


def test_replay_executor_surrogates_external_write_and_blocks_destructive() -> None:
    executor = _executor(
        FakeTool("mcp_outlook_send_email", toolset="mcp", metadata={"transport": "mcp"}),
        FakeTool("delete_account", toolset="mcp", metadata={"transport": "mcp"}),
    )

    send = asyncio.run(executor.execute("mcp_outlook_send_email", {"to": "ada@example.com"}, context=ToolContext()))
    delete = asyncio.run(executor.execute("delete_account", {"id": "1"}, context=ToolContext()))

    assert send.success is True
    assert send.error == "replay_surrogate"
    assert delete.success is False
    assert delete.error == "replay_blocked"
    assert [trace["mode"] for trace in executor.traces] == ["surrogate", "blocked"]
