#!/usr/bin/env python3
# ComfyUI Model Unloader - Décharge les modèles après inactivité
# Usage: python3 comfyui-unload-monitor.py [--interval 60] [--timeout 300]

import requests
import time
import argparse
import logging
import sys
import signal
import subprocess
from datetime import datetime, timedelta

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

DEFAULT_INTERVAL = 60
DEFAULT_TIMEOUT = 300
COMFYUI_URL = "http://localhost:8188"

class ComfyUIUnloadMonitor:
    def __init__(self, interval=DEFAULT_INTERVAL, timeout=DEFAULT_TIMEOUT):
        self.interval = interval
        self.timeout = timeout
        self.last_activity = datetime.now()
        self.running = True
        self.models_loaded = False
        
        signal.signal(signal.SIGINT, self.signal_handler)
        signal.signal(signal.SIGTERM, self.signal_handler)
    
    def signal_handler(self, signum, frame):
        logger.info("Signal reçu, arrêt du monitor...")
        self.running = False
    
    def check_comfyui_activity(self):
        try:
            response = requests.get(f"{COMFYUI_URL}/system_stats", timeout=5)
            if response.status_code == 200:
                stats = response.json()
                return stats
            return None
        except requests.exceptions.RequestException as e:
            logger.warning(f"ComfyUI non accessible: {e}")
            return None
    
    def check_queue_and_history(self):
        try:
            queue_response = requests.get(f"{COMFYUI_URL}/queue", timeout=5)
            history_response = requests.get(f"{COMFYUI_URL}/history?max_items=1", timeout=5)
            
            if queue_response.status_code == 200 and history_response.status_code == 200:
                queue_data = queue_response.json()
                history_data = history_response.json()
                
                running = queue_data.get("queue_running", [])
                pending = queue_data.get("queue_pending", [])
                
                has_activity = len(running) > 0 or len(pending) > 0 or len(history_data) > 0
                
                if history_data:
                    latest = list(history_data.values())[0]
                    if "nodes" in latest:
                        self.last_activity = datetime.now()
                        logger.debug(f"Activité détectée: {len(running)} en cours, {len(pending)} en attente")
                
                return has_activity
            return False
        except Exception as e:
            logger.debug(f"Erreur check_queue: {e}")
            return False
    
    def get_memory_usage(self):
        stats = self.check_comfyui_activity()
        if stats and "memory" in stats:
            vram_used = stats["memory"].get("vram_used", 0)
            ram_used = stats["memory"].get("ram_used", 0)
            total_vram = stats["memory"].get("vram_total", 0)
            
            if total_vram > 0:
                vram_percent = (vram_used / total_vram) * 100
                return vram_used, ram_used, vram_percent
        return None, None, None
    
    def unload_all_models(self):
        logger.info("🔄 Déchargement de tous les modèles ComfyUI...")
        
        workflow = {
            "3": {
                "inputs": {"value": "unload_trigger"},
                "class_type": "UnloadAllModelsNode"
            }
        }
        
        try:
            response = requests.post(
                f"{COMFYUI_URL}/prompt",
                json={"prompt": workflow},
                timeout=30
            )
            
            if response.status_code == 200:
                logger.info("✅ Modèles déchargés avec succès")
                self.models_loaded = False
                return True
            else:
                logger.error(f"❌ Erreur déchargement: {response.status_code}")
                return False
                
        except Exception as e:
            logger.error(f"❌ Exception lors du déchargement: {e}")
            return False
    
    def soft_cleanup(self):
        logger.info("🧹 Nettoyage soft (gc.collect + empty cache)...")
        try:
            requests.post(f"{COMFYUI_URL}/internal/cleanup", timeout=10)
        except:
            pass
        
        vram, ram, percent = self.get_memory_usage()
        if vram is not None:
            logger.info(f"📊 VRAM après cleanup: {vram/1024**3:.2f}GB ({percent:.1f}%)")
    
    def run(self):
        logger.info(f"🚀 Démarrage ComfyUI Unload Monitor")
        logger.info(f"   Intervalle de vérification: {self.interval}s")
        logger.info(f"   Timeout d'inactivité: {self.timeout}s")
        logger.info(f"   URL ComfyUI: {COMFYUI_URL}")
        
        while self.running:
            try:
                vram, ram, vram_percent = self.get_memory_usage()
                
                if vram is not None and vram_percent > 1:
                    self.models_loaded = True
                    idle_time = (datetime.now() - self.last_activity).total_seconds()
                    
                    logger.debug(f"📊 VRAM: {vram/1024**3:.2f}GB ({vram_percent:.1f}%), "
                               f"RAM: {ram/1024**3:.2f}GB, "
                               f"Inactivité: {idle_time:.0f}s")
                    
                    if idle_time >= self.timeout:
                        logger.info(f"⏰ Inactivité de {idle_time:.0f}s >= {self.timeout}s")
                        self.unload_all_models()
                        time.sleep(5)
                        self.soft_cleanup()
                        self.last_activity = datetime.now()
                else:
                    if self.models_loaded:
                        logger.info("📊 Modèles semblent déchargés (VRAM < 1%)")
                        self.models_loaded = False
                
                time.sleep(self.interval)
                
            except Exception as e:
                logger.error(f"Erreur dans la boucle principale: {e}")
                time.sleep(self.interval)
        
        logger.info("👋 Monitor arrêté")

def main():
    parser = argparse.ArgumentParser(description="ComfyUI Model Unload Monitor")
    parser.add_argument("--interval", type=int, default=DEFAULT_INTERVAL,
                       help=f"Intervalle de vérification en secondes (défaut: {DEFAULT_INTERVAL})")
    parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT,
                       help=f"Timeout d'inactivité en secondes (défaut: {DEFAULT_TIMEOUT})")
    parser.add_argument("--debug", action="store_true", help="Mode debug")
    
    args = parser.parse_args()
    
    if args.debug:
        logging.getLogger().setLevel(logging.DEBUG)
    
    monitor = ComfyUIUnloadMonitor(interval=args.interval, timeout=args.timeout)
    monitor.run()

if __name__ == "__main__":
    main()