#!/usr/bin/env python3
"""
Script de Gestion des Ressources LINSTOR pour Proxmox

Auteur: BENE Maël
Version: 1.0
Date: 2025-11-27

Description:
    Ce script gère automatiquement les ressources LINSTOR pour les VMs Proxmox.
    Il assure que les ressources existent avec la taille correcte avant le déploiement.

Fonctionnalités:
    - Vérifie l'existence d'une ressource
    - Crée la ressource si elle n'existe pas
    - Redimensionne la ressource si la taille ne correspond pas (uniquement augmentation)
    - Opérations idempotentes (peut être exécuté plusieurs fois en toute sécurité)
"""

import subprocess
import json
import sys
import argparse
import re
import os
from pathlib import Path


# Configuration des ressources par défaut
# Ces valeurs peuvent être modifiées selon vos besoins
RESOURCE_CONFIG = {
    'vm-1000-disk-0': {
        'node': 'acemagician',
        'size': '100G',
        'storage_pool': 'pve-storage'
    },
    'vm-1001-disk-0': {
        'node': 'elitedesk',
        'size': '100G',
        'storage_pool': 'pve-storage'
    },
    'vm-1002-disk-0': {
        'node': 'thinkpad',
        'size': '20G',
        'storage_pool': 'local-lvm'
    }
}


class LinstorManager:
    """Gestionnaire des ressources LINSTOR"""

    def __init__(self, verbose=False, remote_host=None):
        self.verbose = verbose
        self.remote_host = remote_host  # Hôte sur lequel exécuter les commandes LINSTOR via SSH

    def log(self, message):
        """Affiche un message de log si le mode verbose est activé"""
        if self.verbose:
            print(f"[INFO] {message}")

    def run_command(self, command, remote_host=None):
        """Exécute une commande shell et retourne la sortie

        Args:
            command: Liste de commandes à exécuter
            remote_host: Si spécifié, exécute la commande via SSH sur cet hôte
        """
        if remote_host:
            # Construit la commande SSH avec options pour désactiver la vérification de clé d'hôte
            ssh_command = [
                'ssh',
                '-o', 'StrictHostKeyChecking=no',
                '-o', 'UserKnownHostsFile=/dev/null',
                f'root@{remote_host}'
            ] + command
            self.log(f"Exécution SSH sur {remote_host}: {' '.join(command)}")
            command = ssh_command
        else:
            self.log(f"Exécution locale: {' '.join(command)}")

        try:
            result = subprocess.run(
                command,
                capture_output=True,
                text=True,
                check=True
            )
            return result.stdout.strip()
        except subprocess.CalledProcessError as e:
            print(f"[ERREUR] Échec de la commande: {' '.join(command)}", file=sys.stderr)
            print(f"[ERREUR] Code de sortie: {e.returncode}", file=sys.stderr)
            print(f"[ERREUR] Stdout: {e.stdout}", file=sys.stderr)
            print(f"[ERREUR] Stderr: {e.stderr}", file=sys.stderr)
            return None

    def parse_size(self, size_str):
        """Convertit une chaîne de taille (ex: '100G', '1024M') en octets"""
        size_str = size_str.strip().upper()

        # Correspond au nombre et à l'unité
        match = re.match(r'^(\d+(?:\.\d+)?)\s*([KMGT]?)I?B?$', size_str)
        if not match:
            raise ValueError(f"Format de taille invalide: {size_str}")

        number, unit = match.groups()
        number = float(number)

        multipliers = {
            '': 1,
            'K': 1024,
            'M': 1024**2,
            'G': 1024**3,
            'T': 1024**4,
        }

        return int(number * multipliers.get(unit, 1))

    def format_size(self, bytes_value):
        """Formate les octets en taille lisible"""
        for unit in ['', 'K', 'M', 'G', 'T']:
            if bytes_value < 1024.0:
                return f"{bytes_value:.0f}{unit}iB" if unit else f"{bytes_value:.0f}B"
            bytes_value /= 1024.0
        return f"{bytes_value:.2f}PiB"

    def resource_exists(self, resource_name):
        """Vérifie si une définition de ressource LINSTOR existe"""
        output = self.run_command(['linstor', 'resource-definition', 'list', '--machine-readable'], self.remote_host)
        if output is None:
            return False

        try:
            data = json.loads(output)
            for item in data:
                if isinstance(item, dict) and item.get('name') == resource_name:
                    self.log(f"La définition de ressource '{resource_name}' existe")
                    return True
        except json.JSONDecodeError:
            self.log("Échec de l'analyse de la sortie resource-definition list")

        return False

    def get_resource_size(self, resource_name):
        """Récupère la taille actuelle d'un volume de ressource (en octets)"""
        output = self.run_command(['linstor', 'volume-definition', 'list', '--machine-readable'], self.remote_host)
        if output is None:
            return None

        try:
            data = json.loads(output)
            for item in data:
                if isinstance(item, dict):
                    if item.get('resource_name') == resource_name and item.get('volume_number') == 0:
                        # La taille est en KiB dans LINSTOR
                        size_kib = item.get('size_kib', 0)
                        size_bytes = size_kib * 1024
                        self.log(f"Taille actuelle de '{resource_name}': {self.format_size(size_bytes)}")
                        return size_bytes
        except json.JSONDecodeError:
            self.log("Échec de l'analyse de la sortie volume-definition list")

        return None

    def get_resource_nodes(self, resource_name):
        """Récupère la liste des nœuds où la ressource est déployée"""
        output = self.run_command(['linstor', 'resource', 'list', '--machine-readable'], self.remote_host)
        if output is None:
            return []

        nodes = set()
        try:
            data = json.loads(output)
            for item in data:
                if isinstance(item, dict) and item.get('name') == resource_name:
                    node = item.get('node_name')
                    if node:
                        nodes.add(node)
        except json.JSONDecodeError:
            self.log("Échec de l'analyse de la sortie resource list")

        return list(nodes)

    def create_resource_definition(self, resource_name):
        """Crée une définition de ressource LINSTOR"""
        self.log(f"Création de la définition de ressource '{resource_name}'")
        output = self.run_command(['linstor', 'resource-definition', 'create', resource_name], self.remote_host)
        if output is None:
            return False
        print(f"✓ Définition de ressource '{resource_name}' créée")
        return True

    def create_volume_definition(self, resource_name, size):
        """Crée une définition de volume LINSTOR"""
        self.log(f"Création de la définition de volume pour '{resource_name}' avec taille {size}")
        output = self.run_command(['linstor', 'volume-definition', 'create', resource_name, size], self.remote_host)
        if output is None:
            return False
        print(f"✓ Définition de volume créée pour '{resource_name}' avec taille {size}")
        return True

    def create_resource(self, node, resource_name, storage_pool):
        """Crée une ressource LINSTOR sur un nœud spécifique"""
        self.log(f"Création de la ressource '{resource_name}' sur le nœud '{node}' avec le pool de stockage '{storage_pool}'")
        output = self.run_command([
            'linstor', 'resource', 'create',
            node, resource_name,
            '--storage-pool', storage_pool
        ], self.remote_host)
        if output is None:
            return False
        print(f"✓ Ressource '{resource_name}' créée sur le nœud '{node}'")
        return True

    def resize_volume(self, resource_name, new_size):
        """Redimensionne un volume LINSTOR (uniquement augmentation)"""
        self.log(f"Redimensionnement du volume '{resource_name}' à {new_size}")
        output = self.run_command([
            'linstor', 'volume-definition', 'set-size',
            resource_name, '0', new_size
        ], self.remote_host)
        if output is None:
            return False
        print(f"✓ Volume '{resource_name}' redimensionné à {new_size}")
        return True

    def ensure_resource(self, resource_name, node, size, storage_pool):
        """
        Assure qu'une ressource LINSTOR existe avec la taille correcte

        Args:
            resource_name: Nom de la ressource (ex: 'vm-1000-disk-0')
            node: Nom du nœud cible (ex: 'acemagician')
            size: Taille désirée (ex: '100G')
            storage_pool: Nom du pool de stockage (ex: 'pve-storage')

        Returns:
            True si succès, False sinon
        """
        print(f"\n=== Gestion de la ressource '{resource_name}' sur le nœud '{node}' ===")

        # Vérifie si la définition de ressource existe
        if not self.resource_exists(resource_name):
            print(f"La ressource '{resource_name}' n'existe pas. Création...")

            # Crée la définition de ressource
            if not self.create_resource_definition(resource_name):
                return False

            # Crée la définition de volume
            if not self.create_volume_definition(resource_name, size):
                return False

            # Crée la ressource sur le nœud
            if not self.create_resource(node, resource_name, storage_pool):
                return False

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

        # La ressource existe - vérifier la taille
        print(f"La ressource '{resource_name}' existe déjà. Vérification de la taille...")
        current_size = self.get_resource_size(resource_name)
        desired_size_bytes = self.parse_size(size)

        if current_size is None:
            print(f"[ATTENTION] Impossible de déterminer la taille actuelle de '{resource_name}'")
            return False

        if current_size < desired_size_bytes:
            print(f"Taille actuelle ({self.format_size(current_size)}) inférieure à la taille désirée ({size})")
            print(f"Redimensionnement de la ressource '{resource_name}' à {size}...")
            if not self.resize_volume(resource_name, size):
                return False
            print(f"✓ Ressource '{resource_name}' redimensionnée avec succès")
        elif current_size > desired_size_bytes:
            print(f"[ATTENTION] Taille actuelle ({self.format_size(current_size)}) supérieure à la taille désirée ({size})")
            print(f"[ATTENTION] LINSTOR ne supporte pas la réduction de volumes. Conservation de la taille actuelle.")
        else:
            print(f"✓ Taille correspondante ({self.format_size(current_size)}). Aucune action nécessaire.")

        # Vérifie si la ressource est déployée sur le bon nœud
        deployed_nodes = self.get_resource_nodes(resource_name)
        if node not in deployed_nodes:
            print(f"Ressource '{resource_name}' non déployée sur le nœud '{node}'. Déploiement...")
            if not self.create_resource(node, resource_name, storage_pool):
                return False
            print(f"✓ Ressource '{resource_name}' déployée avec succès sur le nœud '{node}'")
        else:
            self.log(f"Ressource '{resource_name}' déjà déployée sur le nœud '{node}'")

        return True


def parse_terraform_config(terraform_dir):
    """
    Parse les fichiers Terraform pour extraire la configuration des ressources

    Args:
        terraform_dir: Chemin vers le répertoire racine Terraform

    Returns:
        Dictionnaire de configuration des ressources
    """
    config = {}
    terraform_path = Path(terraform_dir)

    # Cherche tous les répertoires pve*
    for pve_dir in terraform_path.glob('pve*'):
        if not pve_dir.is_dir():
            continue

        main_tf = pve_dir / 'main.tf'
        if not main_tf.exists():
            continue

        try:
            with open(main_tf, 'r') as f:
                content = f.read()

            # Extrait VMID
            vmid_match = re.search(r'vmid\s*=\s*(\d+)', content)
            if not vmid_match:
                continue
            vmid = vmid_match.group(1)

            # Extrait target_node
            node_match = re.search(r'target_node\s*=\s*"([^"]+)"', content)
            if not node_match:
                continue
            node = node_match.group(1)

            # Extrait disk size depuis var reference
            size_match = re.search(r'size\s*=\s*var\.(\w+)\.disk_size', content)
            if size_match:
                var_name = size_match.group(1)
                # Cherche la valeur par défaut dans variables.tf ou dans le workflow
                # Pour simplifier, on utilise les valeurs par défaut
                if 'etcd' in var_name.lower():
                    size = '20G'
                else:
                    size = '100G'
            else:
                size = '100G'

            # Extrait storage pool
            storage_match = re.search(r'storage\s*=\s*var\.(\w+)', content)
            if storage_match:
                storage_var = storage_match.group(1)
                if 'etcd' in storage_var.lower():
                    storage_pool = 'local-lvm'
                else:
                    storage_pool = 'linstor_storage'
            else:
                storage_pool = 'linstor_storage'

            resource_name = f"vm-{vmid}-disk-0"
            config[resource_name] = {
                'node': node,
                'size': size,
                'storage_pool': storage_pool
            }

        except Exception as e:
            print(f"[ATTENTION] Erreur lors de la lecture de {main_tf}: {e}", file=sys.stderr)
            continue

    return config


def main():
    parser = argparse.ArgumentParser(
        description='Gestion des ressources LINSTOR pour les VMs Proxmox',
        formatter_class=argparse.RawDescriptionHelpFormatter
    )

    parser.add_argument('--resource', help='Nom de la ressource (ex: vm-1000-disk-0)')
    parser.add_argument('--node', help='Nom du nœud cible (ex: acemagician)')
    parser.add_argument('--size', help='Taille désirée (ex: 100G)')
    parser.add_argument('--storage-pool', help='Nom du pool de stockage (ex: pve-storage)')
    parser.add_argument('--terraform-dir', help='Chemin vers le répertoire Terraform pour lecture automatique de la config')
    parser.add_argument('--all', action='store_true', help='Traiter toutes les ressources configurées')
    parser.add_argument('--remote-host', help='Hôte sur lequel exécuter les commandes LINSTOR via SSH (ex: thinkpad)')
    parser.add_argument('--verbose', '-v', action='store_true', help='Active la sortie détaillée')

    args = parser.parse_args()

    # Utilise --remote-host si spécifié, sinon utilise le premier noeud trouvé dans la config
    manager = LinstorManager(verbose=args.verbose, remote_host=args.remote_host)

    # Mode 1: Lecture automatique depuis Terraform
    if args.terraform_dir:
        print("=== Lecture de la configuration depuis les fichiers Terraform ===")
        config = parse_terraform_config(args.terraform_dir)
        if not config:
            print("[ERREUR] Aucune configuration trouvée dans les fichiers Terraform", file=sys.stderr)
            sys.exit(1)

        print(f"Configuration trouvée pour {len(config)} ressource(s)")
        all_success = True
        for resource_name, res_config in config.items():
            success = manager.ensure_resource(
                resource_name=resource_name,
                node=res_config['node'],
                size=res_config['size'],
                storage_pool=res_config['storage_pool']
            )
            if not success:
                all_success = False

        if all_success:
            print("\n✓ Toutes les ressources sont prêtes")
            sys.exit(0)
        else:
            print("\n✗ Certaines ressources ont échoué", file=sys.stderr)
            sys.exit(1)

    # Mode 2: Traiter toutes les ressources de la configuration par défaut
    elif args.all:
        print("=== Traitement de toutes les ressources configurées ===")
        all_success = True
        for resource_name, config in RESOURCE_CONFIG.items():
            success = manager.ensure_resource(
                resource_name=resource_name,
                node=config['node'],
                size=config['size'],
                storage_pool=config['storage_pool']
            )
            if not success:
                all_success = False

        if all_success:
            print("\n✓ Toutes les ressources sont prêtes")
            sys.exit(0)
        else:
            print("\n✗ Certaines ressources ont échoué", file=sys.stderr)
            sys.exit(1)

    # Mode 3: Traiter une ressource spécifique
    elif args.resource and args.node and args.size and args.storage_pool:
        success = manager.ensure_resource(
            resource_name=args.resource,
            node=args.node,
            size=args.size,
            storage_pool=args.storage_pool
        )

        if success:
            print(f"\n✓ Ressource '{args.resource}' prête sur le nœud '{args.node}'")
            sys.exit(0)
        else:
            print(f"\n✗ Échec de la gestion de la ressource '{args.resource}'", file=sys.stderr)
            sys.exit(1)

    else:
        parser.print_help()
        print("\n[ERREUR] Arguments insuffisants. Utilisez --terraform-dir, --all, ou spécifiez --resource --node --size --storage-pool", file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()