"""Provider-agnostic agent tool example centered on vm_run."""

from __future__ import annotations

import json
from typing import Any

from pyro_mcp import Pyro
from pyro_mcp.vm_manager import (
    DEFAULT_ALLOW_HOST_COMPAT,
    DEFAULT_MEM_MIB,
    DEFAULT_TIMEOUT_SECONDS,
    DEFAULT_TTL_SECONDS,
    DEFAULT_VCPU_COUNT,
)

VM_RUN_TOOL: dict[str, Any] = {
    "name": "vm_run",
    "description": "Run one command in an ephemeral Firecracker VM and clean it up.",
    "input_schema": {
        "type": "object",
        "properties": {
            "environment": {"type": "string"},
            "command": {"type": "string"},
            "vcpu_count": {"type": "integer"},
            "mem_mib": {"type": "integer"},
            "timeout_seconds": {"type": "integer", "default": 30},
            "ttl_seconds": {"type": "integer", "default": 600},
            "network": {"type": "boolean", "default": False},
            "allow_host_compat": {"type": "boolean", "default": False},
        },
        "required": ["environment", "command"],
    },
}


def call_vm_run(arguments: dict[str, Any]) -> dict[str, Any]:
    pyro = Pyro()
    return pyro.run_in_vm(
        environment=str(arguments["environment"]),
        command=str(arguments["command"]),
        vcpu_count=int(arguments.get("vcpu_count", DEFAULT_VCPU_COUNT)),
        mem_mib=int(arguments.get("mem_mib", DEFAULT_MEM_MIB)),
        timeout_seconds=int(arguments.get("timeout_seconds", DEFAULT_TIMEOUT_SECONDS)),
        ttl_seconds=int(arguments.get("ttl_seconds", DEFAULT_TTL_SECONDS)),
        network=bool(arguments.get("network", False)),
        allow_host_compat=bool(arguments.get("allow_host_compat", DEFAULT_ALLOW_HOST_COMPAT)),
    )


def main() -> None:
    tool_arguments: dict[str, Any] = {
        "environment": "debian:12",
        "command": "git --version",
        "timeout_seconds": 30,
        "network": False,
    }
    tool_call: dict[str, Any] = {
        "name": "vm_run",
        "arguments": tool_arguments,
    }

    print(json.dumps({"tool": VM_RUN_TOOL, "tool_call": tool_call}, indent=2, sort_keys=True))
    result = call_vm_run(tool_arguments)
    print(json.dumps({"tool_result": result}, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()