import os
import requests
import json
import time
from dotenv import load_dotenv
from pathlib import Path

# Charger les variables d'environnement
load_dotenv()

LAMBDA_API_KEY = os.getenv("LAMBDA_API_KEY")
LAMBDA_API_BASE_URL = "https://cloud.lambda.ai/api/v1"

class LambdaManager:
    """Gestionnaire pour l'API Lambda Cloud On-Demand GPU."""
    
    def __init__(self, api_key=None):
        self.api_key = api_key or LAMBDA_API_KEY
        if not self.api_key:
            print("❌ LAMBDA_API_KEY non configurée dans le .env")
            
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

    def _get(self, endpoint):
        """Effectue une requête GET vers l'API Lambda."""
        try:
            response = requests.get(f"{LAMBDA_API_BASE_URL}/{endpoint}", headers=self.headers)
            response.raise_for_status()
            return response.json().get("data")
        except Exception as e:
            print(f"❌ Erreur GET /{endpoint}: {e}")
            return None

    def _post(self, endpoint, payload):
        """Effectue une requête POST vers l'API Lambda."""
        try:
            response = requests.post(f"{LAMBDA_API_BASE_URL}/{endpoint}", headers=self.headers, json=payload)
            response.raise_for_status()
            return response.json().get("data")
        except Exception as e:
            print(f"❌ Erreur POST /{endpoint}: {e}")
            return None

    def list_instance_types(self):
        """Liste les types d'instances disponibles et leur prix."""
        data = self._get("instance-types")
        if data:
            print("\n🖥️  Types d'instances disponibles :")
            for name, details in data.items():
                price = details.get("instance_type", {}).get("price_cents_per_hour", 0) / 100
                print(f"- {name:20} : ${price:0.2f}/hr (Config: {details.get('instance_type', {}).get('description')})")
        return data

    def launch_instance(self, instance_type_name="gpu_1x_a10", region_name="us-east-1", ssh_key_names=None):
        """Lance une nouvelle instance on-demand."""
        if not ssh_key_names:
            print("⚠️ Aucun nom de clé SSH spécifié. Vérifiez votre compte Lambda.")
            
        payload = {
            "instance_type_name": instance_type_name,
            "region_name": region_name,
            "ssh_key_names": ssh_key_names or []
        }
        
        print(f"🚀 Lancement de l'instance {instance_type_name}...")
        data = self._post("instance-operations/launch", payload)
        if data:
            ids = data.get("instance_ids", [])
            print(f"✅ Instances lancées : {', '.join(ids)}")
            return ids
        return None

    def list_instances(self):
        """Liste les instances actuellement en cours ou en cours de démarrage."""
        data = self._get("instances")
        if data:
            print("\n🛰️  Instances actives :")
            for inst in data:
                ip = inst.get("ip", "N/A")
                status = inst.get("status", "unknown")
                print(f"- ID: {inst.get('id')} | IP: {ip:15} | Status: {status:10} | Type: {inst.get('instance_type', {}).get('name')}")
        else:
            print("ℹ️ Aucune instance active.")
        return data

    def terminate_instances(self, instance_ids):
        """Arrête et supprime les instances spécifiées."""
        if isinstance(instance_ids, str):
            instance_ids = [instance_ids]
            
        payload = {
            "instance_ids": instance_ids
        }
        
        print(f"🛑 Terminaison des instances : {', '.join(instance_ids)}...")
        data = self._post("instance-operations/terminate", payload)
        if data:
            terminated = data.get("terminated_instance_ids", [])
            print(f"✅ Instances terminées : {', '.join(terminated)}")
            return terminated
        return None

def main():
    manager = LambdaManager()
    
    import sys
    if len(sys.argv) < 2:
        print("\n💡 Usage: python lambda_manager.py [list-types|list|launch|terminate]")
        return

    cmd = sys.argv[1].lower()
    
    if cmd == "list-types":
        manager.list_instance_types()
    elif cmd == "list":
        manager.list_instances()
    elif cmd == "launch":
        # Exemple de lancement par défaut
        manager.launch_instance()
    elif cmd == "terminate":
        if len(sys.argv) < 3:
            print("❌ Spécifiez l'ID de l'instance à terminer.")
        else:
            manager.terminate_instances(sys.argv[2:])

if __name__ == "__main__":
    main()