from pathlib import Path
from genie.testbed import load
from fastmcp import FastMCP


testbed = load(str(Path(__file__).parent / "testbed.yml"))

mcp = FastMCP(
    name="Network MCP Server",
    instructions=(
        "You are a network automation assistant. "
        "Use your tools to answer questions about devices R1, R2, and R3."
    ),
)


def run_command(device_name: str, command: str) -> str:
    device = testbed.devices[device_name]
    device.connect(via="cli", log_stdout=False)
    output = device.execute(command)
    device.disconnect()
    return output


@mcp.tool()
def get_version(device: str) -> str:
    """
    Get the software version and uptime of a device.
    Use this when the user asks about a device's version, uptime, or general info.
    """
    return run_command(device, "show version")


@mcp.tool()
def get_interfaces(device: str) -> dict:
    """
    Get the status of all interfaces on a device.
    Use this when the user asks about interfaces, IPs, or link status.
    """
    device_obj = testbed.devices[device]
    device_obj.connect(via="cli", log_stdout=False)
    parsed = device_obj.parse("show ip interface brief")
    device_obj.disconnect()

    interfaces = []
    for name, data in parsed.get("interface", {}).items():
        interfaces.append({
            "name": name,
            "ip": data.get("ip_address", "unassigned"),
            "status": data.get("status", "unknown"),
            "protocol": data.get("protocol", "unknown"),
        })

    return {"device": device, "interfaces": interfaces}


@mcp.tool()
def run_show_command(device: str, command: str) -> str:
    """
    Run any 'show' command on a device and return the output.
    Use this when the user asks a question not covered by other tools.
    Only use for show commands (read-only). Never use for configuration commands.
    """
    if not command.strip().startswith("show"):
        return "Error: Only 'show' commands are allowed for safety."
    return run_command(device, command)


if __name__ == "__main__":
    mcp.run(transport="streamable-http", host="0.0.0.0", port=8000)