#!/usr/bin/env python3
# Kreuzberg Memory Manager - Redémarre le container après inactivité
# Usage: python3 kreuzberg-memory-manager.py [--interval 60] [--timeout 600] [--max-memory 2]

import requests
import time
import argparse
import logging
import subprocess
import sys

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

DEFAULT_INTERVAL = 60
DEFAULT_TIMEOUT = 600
DEFAULT_MAX_MEMORY = 2.0
KREUZBERG_API = "http://localhost:8300"
CONTAINER_NAME = "kreuzberg"

class KreuzbergMemoryManager:
    def __init__(self, interval=DEFAULT_INTERVAL, timeout=DEFAULT_TIMEOUT, max_memory_gb=DEFAULT_MAX_MEMORY):
        self.interval = interval
        self.timeout = timeout
        self.max_memory_gb = max_memory_gb
        self.last_activity = time.time()
        self.running = True
        self.container_id = None
        
        import signal
        signal.signal(signal.SIGINT, self.signal_handler)
        signal.signal(signal.SIGTERM, self.signal_handler)
        
        self.find_container()
    
    def signal_handler(self, signum, frame):
        logger.info("Signal reçu, arrêt du manager...")
        self.running = False
    
    def find_container(self):
        try:
            result = subprocess.run(
                ["docker", "ps", "-q", "--filter", f"name={CONTAINER_NAME}"],
                capture_output=True, text=True
            )
            if result.stdout.strip():
                self.container_id = result.stdout.strip()
                logger.info(f"Container trouvé: {self.container_id[:12]}")
            else:
                logger.warning(f"Container {CONTAINER_NAME} non trouvé")
        except Exception as e:
            logger.error(f"Erreur recherche container: {e}")
    
    def get_container_memory(self):
        if not self.container_id:
            return None
        try:
            result = subprocess.run(
                ["docker", "stats", "--no-stream", "--format", "{{.MemUsage}}", self.container_id],
                capture_output=True, text=True
            )
            if result.stdout.strip():
                mem_str = result.stdout.strip()
                return self.parse_memory(mem_str)
        except Exception as e:
            logger.debug(f"Erreur stats: {e}")
        return None
    
    def parse_memory(self, mem_str):
        try:
            parts = mem_str.split("/")
            if len(parts) >= 1:
                used = parts[0].strip()
                if "GiB" in used:
                    return float(used.replace("GiB", "")) * 1024
                elif "MiB" in used:
                    return float(used.replace("MiB", ""))
                elif "GB" in used:
                    return float(used.replace("GB", "")) * 1024
                elif "MB" in used:
                    return float(used.replace("MB", ""))
        except:
            pass
        return None
    
    def check_api_activity(self):
        try:
            response = requests.get(f"{KREUZBERG_API}/health", timeout=5)
            if response.status_code in [200, 404]:
                self.last_activity = time.time()
                return True
        except:
            pass
        return False
    
    def restart_container(self):
        logger.info(f"🔄 Redémarrage du container {CONTAINER_NAME}...")
        try:
            subprocess.run(["docker", "restart", CONTAINER_NAME], check=True)
            logger.info("✅ Container redémarré")
            time.sleep(5)
            self.find_container()
            return True
        except Exception as e:
            logger.error(f"❌ Erreur redémarrage: {e}")
            return False
    
    def run(self):
        logger.info(f"🚀 Démarrage Kreuzberg Memory Manager")
        logger.info(f"   Intervalle: {self.interval}s")
        logger.info(f"   Timeout inactivité: {self.timeout}s")
        logger.info(f"   Seuil mémoire: {self.max_memory_gb}GB")
        
        idle_time = 0
        
        while self.running:
            try:
                memory_mb = self.get_container_memory()
                
                if memory_mb is not None:
                    memory_gb = memory_mb / 1024
                    
                    is_active = self.check_api_activity()
                    
                    if is_active:
                        idle_time = 0
                        self.last_activity = time.time()
                        logger.debug(f"📊 Mémoire: {memory_gb:.2f}GB, API active")
                    else:
                        idle_time += self.interval
                        logger.debug(f"📊 Mémoire: {memory_gb:.2f}GB, Inactif: {idle_time}s")
                    
                    if memory_gb > self.max_memory_gb:
                        logger.warning(f"⚠️ Mémoire ({memory_gb:.2f}GB) > seuil ({self.max_memory_gb}GB)")
                        if idle_time >= self.timeout:
                            logger.info(f"⏰ Seuil timeout atteint ({idle_time}s >= {self.timeout}s)")
                            self.restart_container()
                            idle_time = 0
                
                time.sleep(self.interval)
                
            except Exception as e:
                logger.error(f"Erreur: {e}")
                time.sleep(self.interval)
        
        logger.info("👋 Manager arrêté")

def main():
    parser = argparse.ArgumentParser(description="Kreuzberg Memory Manager")
    parser.add_argument("--interval", type=int, default=DEFAULT_INTERVAL,
                       help=f"Intervalle (défaut: {DEFAULT_INTERVAL}s)")
    parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT,
                       help=f"Timeout inactivité (défaut: {DEFAULT_TIMEOUT}s)")
    parser.add_argument("--max-memory", type=float, default=DEFAULT_MAX_MEMORY,
                       help=f"Seuil mémoire GB (défaut: {DEFAULT_MAX_MEMORY}GB)")
    parser.add_argument("--debug", action="store_true", help="Mode debug")
    
    args = parser.parse_args()
    
    if args.debug:
        logging.getLogger().setLevel(logging.DEBUG)
    
    manager = KreuzbergMemoryManager(
        interval=args.interval,
        timeout=args.timeout,
        max_memory_gb=args.max_memory
    )
    manager.run()

if __name__ == "__main__":
    main()