#!/usr/bin/env python3

import os
from typing import Optional

from dotenv import load_dotenv
from proxmoxer import ProxmoxAPI
from mcp.server.fastmcp import FastMCP

load_dotenv("/opt/hermes/MCP_proxmox/.env")

mcp = FastMCP("proxmox-mcp")

PROXMOX_HOST = os.getenv("PROXMOX_HOST")
PROXMOX_USER = os.getenv("PROXMOX_USER")
PROXMOX_TOKEN_NAME = os.getenv("PROXMOX_TOKEN_NAME")
PROXMOX_TOKEN_VALUE = os.getenv("PROXMOX_TOKEN_VALUE")

VERIFY_SSL = (
    os.getenv("PROXMOX_VERIFY_SSL", "true").lower() == "true"
)

ALLOW_POWER_ACTIONS = (
    os.getenv("ALLOW_POWER_ACTIONS", "false").lower() == "true"
)

ALLOW_DELETE = (
    os.getenv("ALLOW_DELETE", "false").lower() == "true"
)

ALLOW_ROLLBACK = (
    os.getenv("ALLOW_ROLLBACK", "false").lower() == "true"
)

def get_proxmox():
    return ProxmoxAPI(
        PROXMOX_HOST,
        user=PROXMOX_USER,
        token_name=PROXMOX_TOKEN_NAME,
        token_value=PROXMOX_TOKEN_VALUE,
        verify_ssl=VERIFY_SSL,
        backend="https",
    )
def power_allowed():
    if not ALLOW_POWER_ACTIONS:
        raise PermissionError(
            "Les actions d'alimentation sont désactivées."
        )


@mcp.tool()
def list_nodes():
    """
    Liste tous les nodes du cluster.
    """

    proxmox = get_proxmox()

    return proxmox.nodes.get()


@mcp.tool()
def cluster_status():
    """
    Retourne l'état général du cluster.
    """

    proxmox = get_proxmox()

    return proxmox.cluster.status.get()


@mcp.tool()
def list_vms(node: Optional[str] = None):
    """
    Liste les VMs QEMU.
    """

    proxmox = get_proxmox()

    result = []

    if node:
        nodes = [node]
    else:
        nodes = [
            n["node"]
            for n in proxmox.nodes.get()
        ]

    for current_node in nodes:

        try:
            vms = proxmox.nodes(current_node).qemu.get()

            for vm in vms:

                result.append({
                    "node": current_node,
                    "vmid": vm.get("vmid"),
                    "name": vm.get("name"),
                    "status": vm.get("status"),
                    "cpus": vm.get("cpus"),
                    "mem": vm.get("mem"),
                    "maxmem": vm.get("maxmem"),
                    "uptime": vm.get("uptime")
                })

        except Exception as e:

            result.append({
                "node": current_node,
                "error": str(e)
            })

    return result


@mcp.tool()
def list_lxc(node: Optional[str] = None):
    """
    Liste les conteneurs LXC.
    """

    proxmox = get_proxmox()

    result = []

    if node:
        nodes = [node]
    else:
        nodes = [
            n["node"]
            for n in proxmox.nodes.get()
        ]

    for current_node in nodes:

        try:
            lxcs = proxmox.nodes(current_node).lxc.get()

            for ct in lxcs:

                result.append({
                    "node": current_node,
                    "vmid": ct.get("vmid"),
                    "name": ct.get("name"),
                    "status": ct.get("status"),
                    "cpus": ct.get("cpus"),
                    "mem": ct.get("mem"),
                    "maxmem": ct.get("maxmem"),
                })

        except Exception as e:

            result.append({
                "node": current_node,
                "error": str(e)
            })

    return result


@mcp.tool()
def vm_status(node: str, vmid: int):
    """
    Statut détaillé d'une VM.
    """

    proxmox = get_proxmox()

    return proxmox.nodes(node).qemu(vmid).status.current.get()


@mcp.tool()
def lxc_status(node: str, vmid: int):
    """
    Statut détaillé d'un LXC.
    """

    proxmox = get_proxmox()

    return proxmox.nodes(node).lxc(vmid).status.current.get()


@mcp.tool()
def node_resources(node: str):
    """
    CPU / RAM / stockage du node.
    """

    proxmox = get_proxmox()

    return proxmox.nodes(node).status.get()


@mcp.tool()
def start_vm(node: str, vmid: int):

    power_allowed()

    proxmox = get_proxmox()

    task = (
        proxmox.nodes(node)
        .qemu(vmid)
        .status.start.post()
    )

    return {
        "action": "start",
        "vmid": vmid,
        "task": task
    }


@mcp.tool()
def stop_vm(node: str, vmid: int):

    power_allowed()

    proxmox = get_proxmox()

    task = (
        proxmox.nodes(node)
        .qemu(vmid)
        .status.stop.post()
    )

    return {
        "action": "stop",
        "vmid": vmid,
        "task": task
    }


@mcp.tool()
def shutdown_vm(node: str, vmid: int):

    power_allowed()

    proxmox = get_proxmox()

    task = (
        proxmox.nodes(node)
        .qemu(vmid)
        .status.shutdown.post()
    )

    return {
        "action": "shutdown",
        "vmid": vmid,
        "task": task
    }


@mcp.tool()
def reboot_vm(node: str, vmid: int):

    power_allowed()

    proxmox = get_proxmox()

    task = (
        proxmox.nodes(node)
        .qemu(vmid)
        .status.reboot.post()
    )

    return {
        "action": "reboot",
        "vmid": vmid,
        "task": task
    }


@mcp.tool()
def create_snapshot(
    node: str,
    vmid: int,
    snapshot_name: str,
    description: str = ""
):
    """
    Création snapshot VM.
    """

    proxmox = get_proxmox()

    task = (
        proxmox.nodes(node)
        .qemu(vmid)
        .snapshot.post(
            snapname=snapshot_name,
            description=description
        )
    )

    return {
        "action": "snapshot",
        "vmid": vmid,
        "snapshot": snapshot_name,
        "task": task
    }


@mcp.tool()
def list_snapshots(node: str, vmid: int):

    proxmox = get_proxmox()

    return (
        proxmox.nodes(node)
        .qemu(vmid)
        .snapshot.get()
    )


@mcp.tool()
def rollback_snapshot(
    node: str,
    vmid: int,
    snapshot_name: str
):
    """
    Rollback snapshot.
    Désactivé par défaut.
    """

    if not ALLOW_ROLLBACK:
        raise PermissionError(
            "Rollback désactivé."
        )

    proxmox = get_proxmox()

    task = (
        proxmox.nodes(node)
        .qemu(vmid)
        .snapshot(snapshot_name)
        .rollback.post()
    )

    return {
        "action": "rollback",
        "vmid": vmid,
        "snapshot": snapshot_name,
        "task": task
    }


if __name__ == "__main__":
    mcp.run()