"""Thin LangChain wrapper over the public vm_run contract.

Requirements:
- `pip install langchain-core` or `uv add langchain-core`

This example keeps the framework integration intentionally narrow: a single tool
that delegates straight to `Pyro.run_in_vm(...)`.
"""

from __future__ import annotations

import json
from typing import Any, Callable, TypeVar, cast

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,
)

F = TypeVar("F", bound=Callable[..., Any])


def run_vm_run_tool(
    *,
    environment: str,
    command: str,
    vcpu_count: int = DEFAULT_VCPU_COUNT,
    mem_mib: int = DEFAULT_MEM_MIB,
    timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
    ttl_seconds: int = DEFAULT_TTL_SECONDS,
    network: bool = False,
    allow_host_compat: bool = DEFAULT_ALLOW_HOST_COMPAT,
) -> str:
    pyro = Pyro()
    result = pyro.run_in_vm(
        environment=environment,
        command=command,
        vcpu_count=vcpu_count,
        mem_mib=mem_mib,
        timeout_seconds=timeout_seconds,
        ttl_seconds=ttl_seconds,
        network=network,
        allow_host_compat=allow_host_compat,
    )
    return json.dumps(result, sort_keys=True)


def build_langchain_vm_run_tool() -> Any:
    try:
        from langchain_core.tools import tool  # type: ignore[import-not-found]
    except ImportError as exc:  # pragma: no cover - dependency guard
        raise RuntimeError(
            "langchain-core is required for this example. Install it with "
            "`uv add langchain-core` or `pip install langchain-core`."
        ) from exc

    decorator = cast(Callable[[F], F], tool("vm_run"))

    @decorator
    def vm_run(
        environment: str,
        command: str,
        vcpu_count: int = DEFAULT_VCPU_COUNT,
        mem_mib: int = DEFAULT_MEM_MIB,
        timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
        ttl_seconds: int = DEFAULT_TTL_SECONDS,
        network: bool = False,
        allow_host_compat: bool = DEFAULT_ALLOW_HOST_COMPAT,
    ) -> str:
        """Run one command in an ephemeral Firecracker VM and clean it up."""
        return run_vm_run_tool(
            environment=environment,
            command=command,
            vcpu_count=vcpu_count,
            mem_mib=mem_mib,
            timeout_seconds=timeout_seconds,
            ttl_seconds=ttl_seconds,
            network=network,
            allow_host_compat=allow_host_compat,
        )

    return vm_run


def main() -> None:
    tool = build_langchain_vm_run_tool()
    print(tool)


if __name__ == "__main__":
    main()