#!/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 os
from datetime import datetime

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
        self.unload_timer = None
        
        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 propre...")
        self.running = False
        sys.exit(0)
    
    def get_vram_usage(self):
        try:
            response = requests.get(f"{COMFYUI_URL}/system_stats", timeout=10)
            if response.status_code == 200:
                data = response.json()
                devices = data.get("devices", [])
                if devices:
                    vram_total = devices[0].get("vram_total", 0)
                    vram_free = devices[0].get("vram_free", 0)
                    vram_used = vram_total - vram_free
                    return vram_used, vram_total
        except Exception as e:
            logger.warning(f"Erreur获取VRAM: {e}")
        return None, None
    
    def check_recent_history(self):
        try:
            response = requests.get(f"{COMFYU_URL}/history?max_items=1", timeout=10)
            if response.status_code == 200:
                history = response.json()
                if history:
                    last_item = list(history.values())[0]
                    if "last_node" in last_item:
                        prompt_time = last_item.get("prompt_time", 0)
                        if prompt_time:
                            last_run = datetime.fromtimestamp(prompt_time)
                            self.last_activity = last_run
                            return True
        except Exception as e:
            pass
        return False
    
    def unload_models(self):
        logger.info("🔄 Déchargement des modèles ComfyUI...")
        try:
            response = requests.post(
                f"{COMFYUI_URL}/internal/model/unload",
                timeout=30
            )
            if response.status_code in [200, 404]:
                logger.info("✅ Déchargement envoyé")
                return True
        except Exception as e:
            logger.warning(f"Erreur déchargement: {e}")
        
        try:
            workflow = {
                "3": {
                    "inputs": {"value": "unload"},
                    "class_type": "UnloadAllModelsNode"
                }
            }
            response = requests.post(
                f"{COMFYUI_URL}/prompt",
                json={"prompt": workflow},
                timeout=30
            )
            if response.status_code == 200:
                logger.info("✅ Workflow déchargement envoyé")
                return True
        except Exception as e:
            logger.error(f"Erreur workflow: {e}")
        return False
    
    def run(self):
        logger.info(f"🚀 ComfyUI Unload Monitor démarré")
        logger.info(f"   Vérification: {self.interval}s | Timeout: {self.timeout}s")
        
        while self.running:
            try:
                vram_used, vram_total = self.get_vram_usage()
                
                if vram_used is not None and vram_total is not None:
                    vram_gb = vram_used / (1024**3)
                    vram_total_gb = vram_total / (1024**3)
                    vram_percent = (vram_used / vram_total) * 100 if vram_total > 0 else 0
                    
                    self.check_recent_history()
                    idle_seconds = (datetime.now() - self.last_activity).total_seconds()
                    
                    logger.info(f"VRAM: {vram_gb:.2f}GB/{vram_total_gb:.2f}GB ({vram_percent:.1f}%) | Idle: {idle_seconds:.0f}s")
                    
                    if vram_percent > 5 and idle_seconds >= self.timeout:
                        logger.info(f"⏰ Inactivité ({idle_seconds:.0f}s >= {self.timeout}s) - Déchargement...")
                        self.unload_models()
                        self.last_activity = datetime.now()
                
                time.sleep(self.interval)
                
            except Exception as e:
                logger.error(f"Erreur: {e}")
                time.sleep(self.interval)
        
        logger.info("👋 Monitor arrêté")
        sys.exit(0)

def main():
    parser = argparse.ArgumentParser(description="ComfyUI Model Unload Monitor")
    parser.add_argument("--interval", type=int, default=DEFAULT_INTERVAL)
    parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT)
    parser.add_argument("--debug", action="store_true")
    
    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()