#!/usr/bin/env python3
# ComfyUI Memory Manager - Redémarre ComfyUI après inactivité
# Usage: python3 comfyui-memory-manager.py [--interval 60] [--timeout 600]

import requests
import time
import argparse
import logging
import subprocess
import sys
import signal
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 = 600
DEFAULT_RECENT_THRESHOLD = 300
COMFYUI_URL = "http://localhost:8188"
CONTAINER_NAME = "comfyui"

class ComfyUIMemoryManager:
    def __init__(self, interval=DEFAULT_INTERVAL, timeout=DEFAULT_TIMEOUT, recent_threshold=DEFAULT_RECENT_THRESHOLD):
        self.interval = interval
        self.timeout = timeout
        self.recent_threshold = recent_threshold
        self.last_activity = datetime.now()
        self.running = True
        
        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...")
        self.running = False
        sys.exit(0)
    
    def get_stats(self):
        try:
            response = requests.get(f"{COMFYUI_URL}/system_stats", timeout=10)
            if response.status_code == 200:
                return response.json()
        except:
            return None
    
    def get_vram_percent(self):
        stats = self.get_stats()
        if stats and "devices" in stats:
            device = stats["devices"][0]
            vram_total = device.get("vram_total", 0)
            vram_free = device.get("vram_free", 0)
            if vram_total > 0:
                vram_used = vram_total - vram_free
                return (vram_used / vram_total) * 100
        return 0
    
    def get_queue_status(self):
        try:
            response = requests.get(f"{COMFYUI_URL}/queue", timeout=10)
            if response.status_code == 200:
                data = response.json()
                running = len(data.get("queue_running", []))
                pending = len(data.get("queue_pending", []))
                return running, pending
        except:
            pass
        return 0, 0
    
    def is_recent_history(self):
        try:
            response = requests.get(f"{COMFYUI_URL}/history?max_items=1", timeout=10)
            if response.status_code == 200:
                history = response.json()
                if history:
                    last_item = list(history.values())[0]
                    prompt_time = last_item.get("prompt_time", 0)
                    if prompt_time:
                        last_run = datetime.fromtimestamp(prompt_time)
                        age = (datetime.now() - last_run).total_seconds()
                        if age < self.recent_threshold:
                            self.last_activity = last_run
                            return True
        except:
            pass
        return False
    
    def restart_container(self):
        logger.info("🔄 Redémarrage du container ComfyUI...")
        try:
            subprocess.run(["docker", "restart", CONTAINER_NAME], check=True)
            logger.info("✅ Container redémarré")
            time.sleep(10)
            return True
        except Exception as e:
            logger.error(f"❌ Erreur redémarrage: {e}")
            return False
    
    def run(self):
        logger.info(f"🚀 ComfyUI Memory Manager démarré")
        logger.info(f"   Vérification: {self.interval}s | Timeout: {self.timeout}s | Recent: {self.recent_threshold}s")
        
        while self.running:
            try:
                vram_percent = self.get_vram_percent()
                running, pending = self.get_queue_status()
                has_recent_history = self.is_recent_history()
                
                idle_seconds = (datetime.now() - self.last_activity).total_seconds()
                
                if running > 0 or pending > 0:
                    self.last_activity = datetime.now()
                    status = "active"
                elif has_recent_history:
                    status = "recent"
                else:
                    status = "idle"
                
                logger.info(f"VRAM: {vram_percent:.1f}% | Queue: {running}r/{pending}p | Idle: {idle_seconds:.0f}s | Status: {status}")
                
                if status == "idle" and vram_percent < 1 and idle_seconds >= self.timeout:
                    logger.info(f"⏰ Inactivité longue ({idle_seconds:.0f}s >= {self.timeout}s), VRAM libre")
                    
                    stats = self.get_stats()
                    if stats:
                        ram_total = stats["system"]["ram_total"] / (1024**3)
                        ram_free = stats["system"]["ram_free"] / (1024**3)
                        ram_used = ram_total - ram_free
                        logger.info(f"📊 RAM: {ram_used:.1f}GB/{ram_total:.1f}GB")
                        
                        if ram_used > 5:
                            logger.info("📊 RAM utilisée élevée, redémarrage pour libérer...")
                            self.restart_container()
                        else:
                            logger.info("📊 RAM déjà basse, pas de redémarrage nécessaire")
                    
                    self.last_activity = datetime.now()
                
                time.sleep(self.interval)
                
            except Exception as e:
                logger.error(f"Erreur: {e}")
                time.sleep(self.interval)
        
        logger.info("👋 Manager arrêté")
        sys.exit(0)

def main():
    parser = argparse.ArgumentParser(description="ComfyUI Memory Manager")
    parser.add_argument("--interval", type=int, default=DEFAULT_INTERVAL)
    parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT)
    parser.add_argument("--recent", type=int, default=DEFAULT_RECENT_THRESHOLD)
    parser.add_argument("--debug", action="store_true")
    
    args = parser.parse_args()
    
    if args.debug:
        logging.getLogger().setLevel(logging.DEBUG)
    
    manager = ComfyUIMemoryManager(
        interval=args.interval, 
        timeout=args.timeout,
        recent_threshold=args.recent
    )
    manager.run()

if __name__ == "__main__":
    main()