#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Script de gestion des ressources DRBD Linstor pour les VMs K3s
Exécuté avant le déploiement Terraform pour s'assurer que les ressources
de stockage DRBD sont créées et dimensionnées correctement.
"""

import subprocess
import sys
import json
import argparse
import os
from typing import Dict, Optional, Tuple


# Clé SSH globale (peut être définie via variable d'environnement)
SSH_KEY_PATH = os.environ.get("SSH_KEY_PATH", None)

# Noms de ressources DRBD fixes pour chaque nœud
RESOURCE_NAMES = {
    1000: "pm-a7f3c8e1",  # acemagician - k3s-server-1
    1001: "pm-b4d2f9a3",  # elitedesk - k3s-server-2
}

# Configuration des nœuds Proxmox
NODE_CONFIG = {
    1000: {"node": "acemagician", "vm_name": "k3s-server-1"},
    1001: {"node": "elitedesk", "vm_name": "k3s-server-2"},
}


def run_ssh_command(command: str, host: str = "192.168.100.30", ssh_key: Optional[str] = None) -> Tuple[int, str, str]:
    """
    Exécute une commande SSH sur le contrôleur Linstor (thinkpad - 192.168.100.30).

    Args:
        command: Commande à exécuter
        host: Hôte sur lequel exécuter la commande (défaut: 192.168.100.30)
        ssh_key: Chemin vers la clé SSH privée (optionnel)

    Returns:
        Tuple (code_retour, stdout, stderr)
    """
    ssh_cmd = ["ssh", "-o", "StrictHostKeyChecking=no"]

    # Ajouter la clé SSH si spécifiée
    if ssh_key:
        ssh_cmd.extend(["-i", ssh_key])

    ssh_cmd.extend([f"root@{host}", command])

    try:
        result = subprocess.run(
            ssh_cmd,
            capture_output=True,
            text=True,
            timeout=30
        )
        return result.returncode, result.stdout, result.stderr
    except subprocess.TimeoutExpired:
        return 1, "", "Timeout lors de l'exécution de la commande SSH"
    except Exception as e:
        return 1, "", f"Erreur lors de l'exécution SSH: {str(e)}"


def check_resource_exists(resource_name: str, ssh_key: Optional[str] = None) -> bool:
    """
    Vérifie si une ressource DRBD existe.

    Args:
        resource_name: Nom de la ressource à vérifier
        ssh_key: Chemin vers la clé SSH privée (optionnel)

    Returns:
        True si la ressource existe, False sinon
    """
    returncode, stdout, stderr = run_ssh_command(
        f"linstor resource list --resource {resource_name} --machine-readable",
        ssh_key=ssh_key
    )

    if returncode == 0:
        # Vérifie si la sortie contient des informations sur la ressource
        return len(stdout.strip()) > 0 and "resources" in stdout

    return False


def get_resource_size(resource_name: str, ssh_key: Optional[str] = None) -> Optional[int]:
    """
    Récupère la taille actuelle d'une ressource DRBD en GiB.

    Args:
        resource_name: Nom de la ressource
        ssh_key: Chemin vers la clé SSH privée (optionnel)

    Returns:
        Taille en GiB ou None si erreur
    """
    returncode, stdout, stderr = run_ssh_command(
        f"linstor volume-definition list --resource {resource_name} --machine-readable",
        ssh_key=ssh_key
    )

    if returncode != 0:
        return None

    try:
        # Parse la sortie JSON de Linstor
        data = json.loads(stdout)
        if data and len(data) > 0:
            volume_defs = data[0].get("volume_definitions", [])
            if volume_defs and len(volume_defs) > 0:
                # Taille en KiB, conversion en GiB
                size_kib = volume_defs[0].get("size_kib", 0)
                return size_kib // (1024 * 1024)
    except (json.JSONDecodeError, KeyError, IndexError) as e:
        print(f"Erreur lors du parsing de la taille: {e}", file=sys.stderr)
        return None

    return None


def create_resource(resource_name: str, size_gib: int, nodes: list, ssh_key: Optional[str] = None) -> bool:
    """
    Crée une nouvelle ressource DRBD avec réplication.

    Args:
        resource_name: Nom de la ressource à créer
        size_gib: Taille en GiB
        nodes: Liste des nœuds pour la réplication
        ssh_key: Chemin vers la clé SSH privée (optionnel)

    Returns:
        True si succès, False sinon
    """
    print(f"Création de la ressource {resource_name} avec {size_gib}GiB...")

    # Étape 1: Créer la définition de ressource
    returncode, stdout, stderr = run_ssh_command(
        f"linstor resource-definition create {resource_name}",
        ssh_key=ssh_key
    )

    if returncode != 0:
        print(f"Erreur lors de la création de la définition: {stderr}", file=sys.stderr)
        return False

    # Étape 2: Créer la définition de volume
    returncode, stdout, stderr = run_ssh_command(
        f"linstor volume-definition create {resource_name} {size_gib}GiB",
        ssh_key=ssh_key
    )

    if returncode != 0:
        print(f"Erreur lors de la création du volume: {stderr}", file=sys.stderr)
        return False

    # Étape 3: Déployer la ressource sur les nœuds avec réplication
    for node in nodes:
        returncode, stdout, stderr = run_ssh_command(
            f"linstor resource create {node} {resource_name} --storage-pool linstor_thinpool",
            ssh_key=ssh_key
        )

        if returncode != 0:
            print(f"Erreur lors du déploiement sur {node}: {stderr}", file=sys.stderr)
            return False

    print(f"✓ Ressource {resource_name} créée avec succès")
    return True


def resize_resource(resource_name: str, new_size_gib: int, ssh_key: Optional[str] = None) -> bool:
    """
    Augmente la taille d'une ressource DRBD existante.

    Args:
        resource_name: Nom de la ressource à redimensionner
        new_size_gib: Nouvelle taille en GiB
        ssh_key: Chemin vers la clé SSH privée (optionnel)

    Returns:
        True si succès, False sinon
    """
    print(f"Redimensionnement de la ressource {resource_name} à {new_size_gib}GiB...")

    returncode, stdout, stderr = run_ssh_command(
        f"linstor volume-definition set-size {resource_name} 0 {new_size_gib}GiB",
        ssh_key=ssh_key
    )

    if returncode != 0:
        print(f"Erreur lors du redimensionnement: {stderr}", file=sys.stderr)
        return False

    print(f"✓ Ressource {resource_name} redimensionnée avec succès")
    return True


def manage_vm_resource(vmid: int, size_gib: int, dry_run: bool = False, ssh_key: Optional[str] = None) -> bool:
    """
    Gère la ressource DRBD pour une VM spécifique.

    Args:
        vmid: ID de la VM
        size_gib: Taille souhaitée en GiB
        dry_run: Si True, affiche les actions sans les exécuter
        ssh_key: Chemin vers la clé SSH privée (optionnel)

    Returns:
        True si succès, False sinon
    """
    if vmid not in RESOURCE_NAMES:
        print(f"VMID {vmid} non configuré", file=sys.stderr)
        return False

    resource_name = RESOURCE_NAMES[vmid]
    node_info = NODE_CONFIG[vmid]

    print(f"\n{'='*60}")
    print(f"Gestion de la ressource pour VM {vmid} ({node_info['vm_name']})")
    print(f"Ressource DRBD: {resource_name}")
    print(f"Nœud Proxmox: {node_info['node']}")
    print(f"Taille souhaitée: {size_gib}GiB")
    print(f"{'='*60}\n")

    # Vérifie si la ressource existe
    resource_exists = check_resource_exists(resource_name, ssh_key=ssh_key)

    if not resource_exists:
        print(f"La ressource {resource_name} n'existe pas.")

        if dry_run:
            print(f"[DRY-RUN] Créerait la ressource {resource_name} avec {size_gib}GiB")
            return True

        # Créer la ressource sur les 3 nœuds pour réplication
        nodes = ["acemagician", "elitedesk", "thinkpad"]
        return create_resource(resource_name, size_gib, nodes, ssh_key=ssh_key)

    else:
        print(f"La ressource {resource_name} existe déjà.")

        # Vérifie la taille actuelle
        current_size = get_resource_size(resource_name, ssh_key=ssh_key)

        if current_size is None:
            print("Impossible de récupérer la taille actuelle", file=sys.stderr)
            return False

        print(f"Taille actuelle: {current_size}GiB")

        if current_size == size_gib:
            print(f"✓ La taille correspond déjà ({size_gib}GiB), aucune action nécessaire")
            return True

        elif current_size < size_gib:
            print(f"La taille doit être augmentée de {current_size}GiB à {size_gib}GiB")

            if dry_run:
                print(f"[DRY-RUN] Redimensionnerait {resource_name} à {size_gib}GiB")
                return True

            return resize_resource(resource_name, size_gib, ssh_key=ssh_key)

        else:
            print(f"⚠ La taille actuelle ({current_size}GiB) est supérieure à la taille souhaitée ({size_gib}GiB)")
            print("La réduction de taille n'est pas supportée, conservation de la taille actuelle")
            return True


def main():
    """Point d'entrée principal du script."""
    parser = argparse.ArgumentParser(
        description="Gestion des ressources DRBD Linstor pour les VMs K3s"
    )
    parser.add_argument(
        "--vmid",
        type=int,
        required=True,
        choices=[1000, 1001],
        help="ID de la VM (1000=acemagician, 1001=elitedesk)"
    )
    parser.add_argument(
        "--size",
        type=int,
        required=True,
        help="Taille du disque en GiB"
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Affiche les actions sans les exécuter"
    )
    parser.add_argument(
        "--ssh-key",
        type=str,
        default=SSH_KEY_PATH,
        help="Chemin vers la clé SSH privée (défaut: variable d'environnement SSH_KEY_PATH)"
    )

    args = parser.parse_args()

    # Exécute la gestion de la ressource
    success = manage_vm_resource(args.vmid, args.size, args.dry_run, ssh_key=args.ssh_key)

    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()