#!/usr/bin/env python3
"""
Proxmox API Client Library for Jarvis PAI
Provides high-level interface for Proxmox VE operations
"""

import json
import os
from pathlib import Path
from typing import Dict, List, Optional, Any
from proxmoxer import ProxmoxAPI
import urllib3

# Disable SSL warnings for self-signed certificates
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


class ProxmoxClient:
    """High-level Proxmox VE API client"""

    def __init__(self, credentials_path: Optional[str] = None):
        """
        Initialize Proxmox client with credentials

        Args:
            credentials_path: Path to credentials JSON file
                            Defaults to ~/.claude/secrets/proxmox.json
        """
        if credentials_path is None:
            credentials_path = os.path.expanduser("~/.claude/secrets/proxmox.json")

        self.creds = self._load_credentials(credentials_path)
        self.proxmox = self._connect()
        self.node = self.creds.get("node", "pve")

    def _load_credentials(self, path: str) -> Dict[str, Any]:
        """Load credentials from JSON file"""
        with open(path, 'r') as f:
            return json.load(f)

    def _connect(self) -> ProxmoxAPI:
        """Establish connection to Proxmox API"""
        # Split token_id into user and token parts
        # Format: "root@pam!jarvis" -> user="root@pam", token="jarvis"
        token_id_parts = self.creds["token_id"].split("!")
        user = token_id_parts[0]
        token_name = token_id_parts[1]

        return ProxmoxAPI(
            self.creds["host"],
            port=self.creds.get("port", 8006),
            user=user,
            token_name=token_name,
            token_value=self.creds["token_secret"],
            verify_ssl=self.creds.get("verify_ssl", False)
        )

    # === VM Lifecycle Operations ===

    def list_vms(self, node: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        List all VMs on a node

        Args:
            node: Proxmox node name (defaults to configured node)

        Returns:
            List of VM dictionaries with vmid, name, status, etc.
        """
        node = node or self.node
        return self.proxmox.nodes(node).qemu.get()

    def get_vm(self, vmid: int, node: Optional[str] = None) -> Dict[str, Any]:
        """
        Get detailed information about a specific VM

        Args:
            vmid: VM ID
            node: Proxmox node name

        Returns:
            Dictionary with VM configuration and status
        """
        node = node or self.node
        config = self.proxmox.nodes(node).qemu(vmid).config.get()
        status = self.proxmox.nodes(node).qemu(vmid).status.current.get()
        return {**config, **status}

    def start_vm(self, vmid: int, node: Optional[str] = None) -> str:
        """
        Start a VM

        Args:
            vmid: VM ID
            node: Proxmox node name

        Returns:
            Task ID (UPID)
        """
        node = node or self.node
        return self.proxmox.nodes(node).qemu(vmid).status.start.post()

    def stop_vm(self, vmid: int, node: Optional[str] = None, force: bool = False) -> str:
        """
        Stop a VM gracefully (or force shutdown)

        Args:
            vmid: VM ID
            node: Proxmox node name
            force: If True, force immediate shutdown

        Returns:
            Task ID (UPID)
        """
        node = node or self.node
        if force:
            return self.proxmox.nodes(node).qemu(vmid).status.stop.post(forceStop=1)
        return self.proxmox.nodes(node).qemu(vmid).status.stop.post()

    def restart_vm(self, vmid: int, node: Optional[str] = None) -> str:
        """
        Restart a VM

        Args:
            vmid: VM ID
            node: Proxmox node name

        Returns:
            Task ID (UPID)
        """
        node = node or self.node
        return self.proxmox.nodes(node).qemu(vmid).status.reboot.post()

    def shutdown_vm(self, vmid: int, node: Optional[str] = None, timeout: int = 60) -> str:
        """
        Graceful VM shutdown with timeout

        Args:
            vmid: VM ID
            node: Proxmox node name
            timeout: Seconds to wait before forcing shutdown

        Returns:
            Task ID (UPID)
        """
        node = node or self.node
        return self.proxmox.nodes(node).qemu(vmid).status.shutdown.post(timeout=timeout)

    # === VM Creation and Cloning ===

    def create_vm(self, vmid: int, name: str, **kwargs) -> str:
        """
        Create a new VM

        Args:
            vmid: VM ID
            name: VM name
            **kwargs: Additional VM parameters (memory, cores, etc.)

        Returns:
            Task ID (UPID)
        """
        node = kwargs.pop("node", self.node)
        params = {"vmid": vmid, "name": name, **kwargs}
        return self.proxmox.nodes(node).qemu.post(**params)

    def clone_vm(self, vmid: int, newid: int, name: Optional[str] = None,
                 node: Optional[str] = None, full: bool = True) -> str:
        """
        Clone an existing VM

        Args:
            vmid: Source VM ID
            newid: New VM ID
            name: Name for cloned VM
            node: Proxmox node name
            full: If True, create full clone (not linked)

        Returns:
            Task ID (UPID)
        """
        node = node or self.node
        params = {"newid": newid, "full": 1 if full else 0}
        if name:
            params["name"] = name
        return self.proxmox.nodes(node).qemu(vmid).clone.post(**params)

    def delete_vm(self, vmid: int, node: Optional[str] = None, purge: bool = True) -> str:
        """
        Delete a VM

        Args:
            vmid: VM ID
            node: Proxmox node name
            purge: If True, also delete from backup jobs and HA

        Returns:
            Task ID (UPID)
        """
        node = node or self.node
        return self.proxmox.nodes(node).qemu(vmid).delete(purge=1 if purge else 0)

    # === Snapshot Management ===

    def list_snapshots(self, vmid: int, node: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        List all snapshots for a VM

        Args:
            vmid: VM ID
            node: Proxmox node name

        Returns:
            List of snapshot dictionaries
        """
        node = node or self.node
        return self.proxmox.nodes(node).qemu(vmid).snapshot.get()

    def create_snapshot(self, vmid: int, snapname: str, description: str = "",
                       vmstate: bool = False, node: Optional[str] = None) -> str:
        """
        Create a VM snapshot

        Args:
            vmid: VM ID
            snapname: Snapshot name
            description: Snapshot description
            vmstate: If True, include RAM state (for running VMs)
            node: Proxmox node name

        Returns:
            Task ID (UPID)
        """
        node = node or self.node
        params = {
            "snapname": snapname,
            "description": description,
            "vmstate": 1 if vmstate else 0
        }
        return self.proxmox.nodes(node).qemu(vmid).snapshot.post(**params)

    def delete_snapshot(self, vmid: int, snapname: str, node: Optional[str] = None) -> str:
        """
        Delete a VM snapshot

        Args:
            vmid: VM ID
            snapname: Snapshot name
            node: Proxmox node name

        Returns:
            Task ID (UPID)
        """
        node = node or self.node
        return self.proxmox.nodes(node).qemu(vmid).snapshot(snapname).delete()

    def rollback_snapshot(self, vmid: int, snapname: str, node: Optional[str] = None) -> str:
        """
        Rollback VM to a snapshot

        Args:
            vmid: VM ID
            snapname: Snapshot name
            node: Proxmox node name

        Returns:
            Task ID (UPID)
        """
        node = node or self.node
        return self.proxmox.nodes(node).qemu(vmid).snapshot(snapname).rollback.post()

    # === Resource Monitoring ===

    def get_vm_status(self, vmid: int, node: Optional[str] = None) -> Dict[str, Any]:
        """
        Get current VM status and resource usage

        Args:
            vmid: VM ID
            node: Proxmox node name

        Returns:
            Dictionary with status, CPU, memory, disk, network stats
        """
        node = node or self.node
        return self.proxmox.nodes(node).qemu(vmid).status.current.get()

    def get_node_status(self, node: Optional[str] = None) -> Dict[str, Any]:
        """
        Get node status and resource usage

        Args:
            node: Proxmox node name

        Returns:
            Dictionary with node stats (CPU, memory, storage, uptime)
        """
        node = node or self.node
        return self.proxmox.nodes(node).status.get()

    def list_nodes(self) -> List[Dict[str, Any]]:
        """
        List all nodes in the Proxmox cluster

        Returns:
            List of node dictionaries
        """
        return self.proxmox.nodes.get()

    # === Container (LXC) Operations ===

    def list_containers(self, node: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        List all LXC containers on a node

        Args:
            node: Proxmox node name

        Returns:
            List of container dictionaries
        """
        node = node or self.node
        return self.proxmox.nodes(node).lxc.get()

    def start_container(self, vmid: int, node: Optional[str] = None) -> str:
        """Start an LXC container"""
        node = node or self.node
        return self.proxmox.nodes(node).lxc(vmid).status.start.post()

    def stop_container(self, vmid: int, node: Optional[str] = None) -> str:
        """Stop an LXC container"""
        node = node or self.node
        return self.proxmox.nodes(node).lxc(vmid).status.stop.post()

    def get_container(self, vmid: int, node: Optional[str] = None) -> Dict[str, Any]:
        """
        Get detailed information about a specific LXC container

        Args:
            vmid: Container ID
            node: Proxmox node name

        Returns:
            Dictionary with container configuration and status
        """
        node = node or self.node
        config = self.proxmox.nodes(node).lxc(vmid).config.get()
        status = self.proxmox.nodes(node).lxc(vmid).status.current.get()
        return {**config, **status}

    def get_container_status(self, vmid: int, node: Optional[str] = None) -> Dict[str, Any]:
        """
        Get current container status and resource usage

        Args:
            vmid: Container ID
            node: Proxmox node name

        Returns:
            Dictionary with status, CPU, memory, disk, network stats
        """
        node = node or self.node
        return self.proxmox.nodes(node).lxc(vmid).status.current.get()

    def create_container(self, vmid: int, ostemplate: str, hostname: str,
                        storage: str = "local-lvm", password: Optional[str] = None,
                        ssh_public_keys: Optional[str] = None,
                        node: Optional[str] = None, **kwargs) -> str:
        """
        Create a new LXC container

        Args:
            vmid: Container ID
            ostemplate: OS template (e.g., "local:vztmpl/ubuntu-22.04-standard_22.04-1_amd64.tar.zst")
            hostname: Container hostname
            storage: Storage for container root filesystem
            password: Root password (optional, SSH keys preferred)
            ssh_public_keys: SSH public keys for root access
            node: Proxmox node name
            **kwargs: Additional container parameters (memory, cores, rootfs, net0, etc.)

        Returns:
            Task ID (UPID)
        """
        node = node or self.node
        params = {
            "vmid": vmid,
            "ostemplate": ostemplate,
            "hostname": hostname,
            "storage": storage,
            **kwargs
        }

        if password:
            params["password"] = password
        if ssh_public_keys:
            params["ssh-public-keys"] = ssh_public_keys

        return self.proxmox.nodes(node).lxc.post(**params)

    def configure_container_for_docker(self, vmid: int, node: Optional[str] = None) -> bool:
        """
        Configure an LXC container to support Docker (nested containers)

        Args:
            vmid: Container ID
            node: Proxmox node name

        Returns:
            True if configuration successful
        """
        node = node or self.node

        # Required features for Docker support
        docker_config = {
            "features": "nesting=1,keyctl=1",
            "unprivileged": 0,  # Privileged container for Docker
        }

        try:
            self.proxmox.nodes(node).lxc(vmid).config.put(**docker_config)
            return True
        except Exception as e:
            print(f"Error configuring container for Docker: {e}")
            return False

    def restart_container(self, vmid: int, node: Optional[str] = None) -> str:
        """
        Restart an LXC container

        Args:
            vmid: Container ID
            node: Proxmox node name

        Returns:
            Task ID (UPID)
        """
        node = node or self.node
        return self.proxmox.nodes(node).lxc(vmid).status.reboot.post()

    def shutdown_container(self, vmid: int, node: Optional[str] = None, timeout: int = 60) -> str:
        """
        Graceful container shutdown with timeout

        Args:
            vmid: Container ID
            node: Proxmox node name
            timeout: Seconds to wait before forcing shutdown

        Returns:
            Task ID (UPID)
        """
        node = node or self.node
        return self.proxmox.nodes(node).lxc(vmid).status.shutdown.post(timeout=timeout)

    def delete_container(self, vmid: int, node: Optional[str] = None, purge: bool = True) -> str:
        """
        Delete an LXC container

        Args:
            vmid: Container ID
            node: Proxmox node name
            purge: If True, also delete from backup jobs

        Returns:
            Task ID (UPID)
        """
        node = node or self.node
        return self.proxmox.nodes(node).lxc(vmid).delete(purge=1 if purge else 0)

    def clone_container(self, vmid: int, newid: int, hostname: Optional[str] = None,
                       node: Optional[str] = None, full: bool = True) -> str:
        """
        Clone an existing LXC container

        Args:
            vmid: Source container ID
            newid: New container ID
            hostname: Hostname for cloned container (optional)
            node: Proxmox node name
            full: If True, create full clone (not linked)

        Returns:
            Task ID (UPID)
        """
        node = node or self.node
        params = {"newid": newid, "full": 1 if full else 0}
        if hostname:
            params["hostname"] = hostname
        return self.proxmox.nodes(node).lxc(vmid).clone.post(**params)

    def list_container_snapshots(self, vmid: int, node: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        List all snapshots for an LXC container

        Args:
            vmid: Container ID
            node: Proxmox node name

        Returns:
            List of snapshot dictionaries
        """
        node = node or self.node
        return self.proxmox.nodes(node).lxc(vmid).snapshot.get()

    def create_container_snapshot(self, vmid: int, snapname: str, description: str = "",
                                 node: Optional[str] = None) -> str:
        """
        Create an LXC container snapshot

        Args:
            vmid: Container ID
            snapname: Snapshot name
            description: Snapshot description
            node: Proxmox node name

        Returns:
            Task ID (UPID)
        """
        node = node or self.node
        params = {
            "snapname": snapname,
            "description": description
        }
        return self.proxmox.nodes(node).lxc(vmid).snapshot.post(**params)

    def delete_container_snapshot(self, vmid: int, snapname: str, node: Optional[str] = None) -> str:
        """
        Delete an LXC container snapshot

        Args:
            vmid: Container ID
            snapname: Snapshot name
            node: Proxmox node name

        Returns:
            Task ID (UPID)
        """
        node = node or self.node
        return self.proxmox.nodes(node).lxc(vmid).snapshot(snapname).delete()

    def rollback_container_snapshot(self, vmid: int, snapname: str, node: Optional[str] = None) -> str:
        """
        Rollback container to a snapshot

        Args:
            vmid: Container ID
            snapname: Snapshot name
            node: Proxmox node name

        Returns:
            Task ID (UPID)
        """
        node = node or self.node
        return self.proxmox.nodes(node).lxc(vmid).snapshot(snapname).rollback.post()

    def get_all_containers_status(self, node: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        Get status summary for all LXC containers

        Args:
            node: Proxmox node name

        Returns:
            List of container status dictionaries
        """
        containers = self.list_containers(node)
        result = []
        for ct in containers:
            status = self.get_container_status(ct["vmid"], node)
            result.append({
                "vmid": ct["vmid"],
                "name": ct.get("name", ""),
                "status": status.get("status", "unknown"),
                "cpu": status.get("cpu", 0),
                "mem": status.get("mem", 0),
                "maxmem": status.get("maxmem", 0),
                "disk": status.get("disk", 0),
                "maxdisk": status.get("maxdisk", 0),
                "uptime": status.get("uptime", 0)
            })
        return result

    def list_container_templates(self, node: Optional[str] = None, storage: str = "local") -> List[Dict[str, Any]]:
        """
        List available LXC templates

        Args:
            node: Proxmox node name
            storage: Storage containing templates

        Returns:
            List of available templates
        """
        node = node or self.node
        return self.proxmox.nodes(node).storage(storage).content.get(content="vztmpl")

    # === Storage Operations ===

    def list_storage(self, node: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        List all storage on a node

        Args:
            node: Proxmox node name

        Returns:
            List of storage dictionaries
        """
        node = node or self.node
        return self.proxmox.nodes(node).storage.get()

    def get_storage_status(self, storage: str, node: Optional[str] = None) -> Dict[str, Any]:
        """
        Get storage status and usage

        Args:
            storage: Storage name
            node: Proxmox node name

        Returns:
            Dictionary with storage stats
        """
        node = node or self.node
        return self.proxmox.nodes(node).storage(storage).status.get()

    # === Task Management ===

    def get_task_status(self, upid: str, node: Optional[str] = None) -> Dict[str, Any]:
        """
        Get status of a Proxmox task

        Args:
            upid: Task ID (UPID)
            node: Proxmox node name

        Returns:
            Dictionary with task status and progress
        """
        node = node or self.node
        return self.proxmox.nodes(node).tasks(upid).status.get()

    def wait_for_task(self, upid: str, node: Optional[str] = None,
                     timeout: int = 300, poll_interval: int = 2) -> bool:
        """
        Wait for a task to complete

        Args:
            upid: Task ID (UPID)
            node: Proxmox node name
            timeout: Maximum seconds to wait
            poll_interval: Seconds between status checks

        Returns:
            True if task succeeded, False if failed or timeout
        """
        import time
        node = node or self.node
        elapsed = 0

        while elapsed < timeout:
            status = self.get_task_status(upid, node)
            if status.get("status") == "stopped":
                return status.get("exitstatus") == "OK"
            time.sleep(poll_interval)
            elapsed += poll_interval

        return False

    # === Backup Operations ===

    def create_backup(self, vmid: int, storage: str = "local",
                     mode: str = "snapshot", compress: str = "zstd",
                     node: Optional[str] = None) -> str:
        """
        Create a VM backup

        Args:
            vmid: VM ID
            storage: Backup storage location
            mode: Backup mode (snapshot, suspend, stop)
            compress: Compression algorithm (zstd, gzip, lzo)
            node: Proxmox node name

        Returns:
            Task ID (UPID)
        """
        node = node or self.node
        return self.proxmox.nodes(node).vzdump.post(
            vmid=vmid,
            storage=storage,
            mode=mode,
            compress=compress
        )

    def list_backups(self, storage: str = "local",
                    node: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        List backups in storage

        Args:
            storage: Storage name
            node: Proxmox node name

        Returns:
            List of backup dictionaries
        """
        node = node or self.node
        return self.proxmox.nodes(node).storage(storage).content.get(content="backup")

    # === Network Configuration ===

    def list_networks(self, node: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        List network interfaces on a node

        Args:
            node: Proxmox node name

        Returns:
            List of network interface dictionaries
        """
        node = node or self.node
        return self.proxmox.nodes(node).network.get()

    # === Convenience Methods ===

    def get_vm_by_name(self, name: str, node: Optional[str] = None) -> Optional[Dict[str, Any]]:
        """
        Find a VM by name

        Args:
            name: VM name
            node: Proxmox node name

        Returns:
            VM dictionary or None if not found
        """
        vms = self.list_vms(node)
        for vm in vms:
            if vm.get("name") == name:
                return self.get_vm(vm["vmid"], node)
        return None

    def get_all_vms_status(self, node: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        Get status summary for all VMs

        Args:
            node: Proxmox node name

        Returns:
            List of VM status dictionaries with vmid, name, status, resources
        """
        vms = self.list_vms(node)
        result = []
        for vm in vms:
            status = self.get_vm_status(vm["vmid"], node)
            result.append({
                "vmid": vm["vmid"],
                "name": vm.get("name", ""),
                "status": status.get("status", "unknown"),
                "cpu": status.get("cpu", 0),
                "mem": status.get("mem", 0),
                "maxmem": status.get("maxmem", 0),
                "disk": status.get("disk", 0),
                "maxdisk": status.get("maxdisk", 0),
                "uptime": status.get("uptime", 0)
            })
        return result


def main():
    """CLI interface for testing"""
    import sys

    if len(sys.argv) < 2:
        print("Usage: proxmox_client.py <command> [args...]")
        print("\nCommands:")
        print("  list - List all VMs")
        print("  status <vmid> - Get VM status")
        print("  start <vmid> - Start a VM")
        print("  stop <vmid> - Stop a VM")
        print("  nodes - List all nodes")
        sys.exit(1)

    client = ProxmoxClient()
    command = sys.argv[1]

    if command == "list":
        vms = client.get_all_vms_status()
        for vm in vms:
            print(f"VM {vm['vmid']}: {vm['name']} - {vm['status']}")

    elif command == "status" and len(sys.argv) > 2:
        vmid = int(sys.argv[2])
        status = client.get_vm_status(vmid)
        print(json.dumps(status, indent=2))

    elif command == "start" and len(sys.argv) > 2:
        vmid = int(sys.argv[2])
        upid = client.start_vm(vmid)
        print(f"Starting VM {vmid}... Task ID: {upid}")

    elif command == "stop" and len(sys.argv) > 2:
        vmid = int(sys.argv[2])
        upid = client.stop_vm(vmid)
        print(f"Stopping VM {vmid}... Task ID: {upid}")

    elif command == "nodes":
        nodes = client.list_nodes()
        for node in nodes:
            print(f"Node: {node['node']} - Status: {node['status']}")

    else:
        print(f"Unknown command: {command}")
        sys.exit(1)


if __name__ == "__main__":
    main()