#!/usr/bin/env python3
"""
hot_swap_manager.py — Hot Swap System for Ollama Models & Surfaces

Manages dynamic loading/unloading of AI models with memory-aware scheduling.
Integrates with ENE mesh for distributed hot-swap coordination.

Features:
- Memory-pressure-aware model swapping
- Automatic model tier selection (fast <-> capable)
- Zero-downtime swap via overlap window
- ENE gossip for distributed state

Author: OTOM Research Stack
Version: 1.0.0
"""

import json
import logging
import time
import psutil
import requests
import hashlib
import threading
from pathlib import Path
from dataclasses import dataclass, asdict
from typing import Dict, List, Optional, Any, Callable
from datetime import datetime, timedelta
from enum import Enum

logger = logging.getLogger(__name__)


class ModelTier(Enum):
    """Model capability tiers for swap decisions."""
    FAST = "fast"           # Small, quick models (8B)
    BALANCED = "balanced"   # Mid-size (14-32B)
    CAPABLE = "capable"     # Large models (70B+)
    SPECIALIST = "specialist"  # Domain-specific


@dataclass
class ModelSpec:
    """Specification for a hot-swappable model."""
    name: str
    size_gb: float
    vram_gb: float
    tier: ModelTier
    ram_gb: float = 0.0  # System RAM needed if not fitting in VRAM
    load_time_ms: int = 5000
    context_length: int = 32768
    quantization: str = "Q4_K_M"
    
    def fits_in_vram(self, available_vram_gb: float) -> bool:
        return self.vram_gb <= available_vram_gb
    
    def total_memory_needed(self) -> float:
        return max(self.vram_gb, self.ram_gb)


@dataclass
class SwapEvent:
    """Record of a hot swap operation."""
    timestamp: float
    from_model: Optional[str]
    to_model: str
    reason: str
    duration_ms: int
    memory_before_gb: float
    memory_after_gb: float
    success: bool


class HotSwapManager:
    """
    Manages hot-swapping of Ollama models based on system resources.
    
    Architecture:
    - Monitor: Track memory/VRAM pressure
    - Policy: Decide when to swap tiers
    - Executor: Perform atomic swap with overlap
    - Gossip: Sync state across ENE mesh
    """
    
    OLLAMA_HOST = "http://localhost:11435"
    
    # Model registry with specs — auto-populated from Ollama at init
    _BASE_REGISTRY = {
        # Small, fast models (< 5GB)
        "deepseek-r1:8b": ModelSpec(
            name="deepseek-r1:8b",
            size_gb=4.0,
            vram_gb=4.5,
            ram_gb=0.5,
            tier=ModelTier.FAST,
            load_time_ms=2000,
            context_length=32768
        ),
        "llama3.1:8b": ModelSpec(
            name="llama3.1:8b",
            size_gb=4.7,
            vram_gb=5.0,
            ram_gb=0.5,
            tier=ModelTier.FAST,
            load_time_ms=2000,
            context_length=32768
        ),
        "gemma3:1b": ModelSpec(
            name="gemma3:1b",
            size_gb=0.8,
            vram_gb=1.0,
            ram_gb=0.2,
            tier=ModelTier.FAST,
            load_time_ms=500,
            context_length=8192
        ),
        "smollm2:135m": ModelSpec(
            name="smollm2:135m",
            size_gb=0.3,
            vram_gb=0.5,
            ram_gb=0.1,
            tier=ModelTier.FAST,
            load_time_ms=300,
            context_length=2048
        ),
        "mixtral:8x7b": ModelSpec(
            name="mixtral:8x7b",
            size_gb=4.2,
            vram_gb=4.8,
            ram_gb=0.5,
            tier=ModelTier.FAST,
            load_time_ms=3000,
            context_length=32768
        ),
        
        # Balanced models (5-15GB)
        "qwen2.5-coder:14b": ModelSpec(
            name="qwen2.5-coder:14b",
            size_gb=9.0,
            vram_gb=9.5,
            ram_gb=1.0,
            tier=ModelTier.BALANCED,
            load_time_ms=4000,
            context_length=32768
        ),
        "phi4:14b": ModelSpec(
            name="phi4:14b",
            size_gb=9.0,
            vram_gb=9.5,
            ram_gb=1.0,
            tier=ModelTier.BALANCED,
            load_time_ms=4000,
            context_length=32768
        ),
        "dagbs/qwen2.5-coder-14b-instruct-abliterated:latest": ModelSpec(
            name="qwen2.5-coder-14b-abliterated",
            size_gb=9.0,
            vram_gb=9.5,
            ram_gb=1.0,
            tier=ModelTier.BALANCED,
            load_time_ms=4000,
            context_length=32768
        ),
        
        # Mid-large models (15-25GB)
        "deepseek-r1:32b": ModelSpec(
            name="deepseek-r1:32b",
            size_gb=19.0,
            vram_gb=20.0,
            ram_gb=2.0,
            tier=ModelTier.BALANCED,
            load_time_ms=8000,
            context_length=32768
        ),
        "richardyoung/bfs-prover-v2-32b:q4_k_m": ModelSpec(
            name="bfs-prover-v2-32b",
            size_gb=18.0,
            vram_gb=19.0,
            ram_gb=1.0,
            tier=ModelTier.SPECIALIST,
            load_time_ms=8000,
            context_length=32768
        ),
        "charaf/Huihui-Qwen3.6-27B-abliterated-mlx-nvfp4:instruct": ModelSpec(
            name="huihui-qwen3.6-27b",
            size_gb=16.0,
            vram_gb=17.0,
            ram_gb=1.0,
            tier=ModelTier.SPECIALIST,
            load_time_ms=7000,
            context_length=32768
        ),
        
        # Specialist models
        "hf.co/TrevorJS/gemma-4-E2B-it-uncensored-GGUF:latest": ModelSpec(
            name="gemma-4-e2b",
            size_gb=3.0,
            vram_gb=3.5,
            ram_gb=0.5,
            tier=ModelTier.SPECIALIST,
            load_time_ms=1500,
            context_length=8192
        ),
    }
    
    MODEL_REGISTRY: Dict[str, ModelSpec] = {}
    
    def _discover_models(self):
        """Populate MODEL_REGISTRY from actually-installed Ollama models."""
        try:
            resp = requests.get(f"{self.OLLAMA_HOST}/api/tags", timeout=10)
            data = resp.json()
            installed = {m["name"] for m in data.get("models", [])}
        except Exception as e:
            print(f"[WARN] Could not query Ollama models: {e}")
            installed = set(self._BASE_REGISTRY.keys())
        
        for name, spec in self._BASE_REGISTRY.items():
            if name in installed:
                self.MODEL_REGISTRY[name] = spec
        
        # Add any installed models we don't have specs for with safe defaults
        for model_name in installed:
            if model_name not in self.MODEL_REGISTRY:
                # Infer size from name
                size_gb = 5.0  # default guess
                if ":" in model_name:
                    tag = model_name.split(":")[-1]
                    if any(x in tag for x in ["135m", "1b", "2b", "3b"]):
                        size_gb = 1.0
                    elif any(x in tag for x in ["7b", "8b"]):
                        size_gb = 4.5
                    elif any(x in tag for x in ["14b", "13b"]):
                        size_gb = 9.0
                    elif any(x in tag for x in ["27b", "32b"]):
                        size_gb = 18.0
                    elif any(x in tag for x in ["70b"]):
                        size_gb = 40.0
                
                tier = ModelTier.BALANCED if size_gb > 8 else ModelTier.FAST
                
                self.MODEL_REGISTRY[model_name] = ModelSpec(
                    name=model_name,
                    size_gb=size_gb,
                    vram_gb=size_gb + 1.0,
                    ram_gb=1.0,
                    tier=tier,
                    load_time_ms=int(size_gb * 500),
                    context_length=32768
                )
        
        print(f"[DISCOVER] {len(self.MODEL_REGISTRY)} models registered from Ollama")
    
    def __init__(self, 
                 memory_threshold_gb: float = 4.0,
                 vram_threshold_gb: float = 2.0,
                 swap_overlap_ms: int = 500):
        """
        Initialize hot swap manager.
        
        Args:
            memory_threshold_gb: Free memory threshold for downsizing
            vram_threshold_gb: Free VRAM threshold for downsizing
            swap_overlap_ms: Overlap window for zero-downtime swap
        """
        self.memory_threshold_gb = memory_threshold_gb
        self.vram_threshold_gb = vram_threshold_gb
        self.swap_overlap_ms = swap_overlap_ms
        
        self.current_model: Optional[str] = None
        self.swap_history: List[SwapEvent] = []
        self.callbacks: List[Callable[[SwapEvent], None]] = []
        self._lock = threading.Lock()
        
        # Discover available models
        self._discover_models()
        
        # Load current state
        self._refresh_state()
    
    def _refresh_state(self):
        """Refresh current model state from Ollama."""
        try:
            resp = requests.get(f"{self.OLLAMA_HOST}/api/ps", timeout=5)
            data = resp.json()
            models = data.get("models", [])
            if models:
                self.current_model = models[0].get("name")
            else:
                self.current_model = None
        except Exception as e:
            print(f"[WARN] Could not refresh Ollama state: {e}")
            self.current_model = None
    
    def get_system_resources(self) -> Dict[str, float]:
        """Get current system resource status."""
        mem = psutil.virtual_memory()
        
        resources = {
            "total_ram_gb": mem.total / (1024**3),
            "available_ram_gb": mem.available / (1024**3),
            "used_ram_gb": mem.used / (1024**3),
            "ram_percent": mem.percent,
        }
        
        # Get VRAM from Ollama if available
        try:
            resp = requests.get(f"{self.OLLAMA_HOST}/api/tags", timeout=5)
            data = resp.json()
            # Ollama doesn't expose VRAM directly, we infer from loaded model
            if self.current_model and self.current_model in self.MODEL_REGISTRY:
                spec = self.MODEL_REGISTRY[self.current_model]
                resources["used_vram_gb"] = spec.vram_gb
                resources["available_vram_gb"] = 12.0 - spec.vram_gb  # RTX 4070 SUPER
            else:
                resources["used_vram_gb"] = 0.0
                resources["available_vram_gb"] = 12.0
        except (requests.RequestException, ValueError, KeyError) as exc:
            logger.warning("Failed to query Ollama for VRAM info: %s", exc)
            resources["used_vram_gb"] = 0.0
            resources["available_vram_gb"] = 0.0
        
        return resources
    
    def should_swap_down(self) -> bool:
        """Check if we should swap to a smaller model."""
        resources = self.get_system_resources()
        
        # Check RAM pressure
        if resources["available_ram_gb"] < self.memory_threshold_gb:
            return True
        
        # Check VRAM pressure
        if resources["available_vram_gb"] < self.vram_threshold_gb:
            return True
        
        # Check if swap is high
        swap = psutil.swap_memory()
        if swap.percent > 50:
            return True
        
        return False
    
    def can_swap_up(self, target_model: str) -> bool:
        """Check if we can swap to a larger model."""
        if target_model not in self.MODEL_REGISTRY:
            return False
        
        spec = self.MODEL_REGISTRY[target_model]
        resources = self.get_system_resources()
        
        # Check if enough free memory
        total_needed = spec.total_memory_needed()
        if resources["available_ram_gb"] < total_needed + self.memory_threshold_gb:
            return False
        
        # Check VRAM
        if not spec.fits_in_vram(resources["available_vram_gb"] + resources["used_vram_gb"]):
            # Would need system RAM fallback
            if resources["available_ram_gb"] < spec.ram_gb + 2.0:  # 2GB buffer
                return False
        
        return True
    
    def get_swap_candidate(self, direction: str = "auto") -> Optional[str]:
        """
        Get best model for hot swap.
        
        Args:
            direction: "up", "down", or "auto" based on resources
        """
        if direction == "auto":
            if self.should_swap_down():
                direction = "down"
            else:
                direction = "up"
        
        current_tier = ModelTier.FAST
        if self.current_model and self.current_model in self.MODEL_REGISTRY:
            current_tier = self.MODEL_REGISTRY[self.current_model].tier
        
        if direction == "down":
            # Find smaller model in same or lower tier
            candidates = [
                (name, spec) for name, spec in self.MODEL_REGISTRY.items()
                if spec.tier.value <= current_tier.value
                and spec.size_gb < (self.MODEL_REGISTRY.get(self.current_model, spec).size_gb * 0.8 if self.current_model else float('inf'))
            ]
            # Sort by size ascending
            candidates.sort(key=lambda x: x[1].size_gb)
            if candidates:
                return candidates[0][0]
            
            # Fallback to smallest
            smallest = min(self.MODEL_REGISTRY.items(), key=lambda x: x[1].size_gb)
            return smallest[0]
        
        else:  # direction == "up"
            # Find larger model we can fit
            candidates = [
                (name, spec) for name, spec in self.MODEL_REGISTRY.items()
                if spec.tier.value >= current_tier.value
                and self.can_swap_up(name)
            ]
            # Sort by capability descending
            candidates.sort(key=lambda x: x[1].size_gb, reverse=True)
            if candidates:
                return candidates[0][0]
        
        return None
    
    def perform_swap(self, target_model: str, reason: str = "manual") -> SwapEvent:
        """
        Perform hot swap to target model.
        
        Uses overlap window for minimal downtime:
        1. Start loading target model
        2. Wait for overlap window
        3. Unload current model
        """
        start_time = time.time()
        resources_before = self.get_system_resources()
        
        print(f"\n[🔄 HOT SWAP] {self.current_model} → {target_model}")
        print(f"   Reason: {reason}")
        print(f"   Resources: {resources_before['available_ram_gb']:.1f}GB RAM free, "
              f"{resources_before['available_vram_gb']:.1f}GB VRAM free")
        
        success = False
        
        try:
            # Step 1: Load target model (Ollama handles the overlap internally)
            print(f"   [1/3] Loading {target_model}...")
            load_start = time.time()
            
            # Generate a simple request to force model load
            resp = requests.post(
                f"{self.OLLAMA_HOST}/api/generate",
                json={"model": target_model, "prompt": "Hi", "stream": False},
                timeout=120
            )
            
            if resp.status_code != 200:
                raise RuntimeError(f"Failed to load model: {resp.text}")
            
            load_duration = (time.time() - load_start) * 1000
            print(f"   ✓ Loaded in {load_duration:.0f}ms")
            
            # Step 2: Verify target is active
            print(f"   [2/3] Verifying target model...")
            time.sleep(0.5)  # Brief overlap
            self._refresh_state()
            
            if self.current_model != target_model:
                raise RuntimeError(f"Model swap verification failed: got {self.current_model}")
            
            print(f"   ✓ {target_model} is active")
            
            # Step 3: Unload old model if different
            if self.current_model and self.current_model != target_model:
                print(f"   [3/3] Unloading previous model...")
                # Ollama auto-unloads, but we can force it
                requests.post(
                    f"{self.OLLAMA_HOST}/api/generate",
                    json={"model": target_model, "prompt": "", "keep_alive": 0},
                    timeout=10
                )
            else:
                print(f"   [3/3] Previous model auto-unloaded")
            
            success = True
            print(f"   ✓ Swap complete")
            
        except Exception as e:
            print(f"   ✗ Swap failed: {e}")
            success = False
        
        duration_ms = int((time.time() - start_time) * 1000)
        resources_after = self.get_system_resources()
        
        event = SwapEvent(
            timestamp=start_time,
            from_model=self.current_model,
            to_model=target_model,
            reason=reason,
            duration_ms=duration_ms,
            memory_before_gb=resources_before["available_ram_gb"],
            memory_after_gb=resources_after["available_ram_gb"],
            success=success
        )
        
        with self._lock:
            self.swap_history.append(event)
        
        # Notify callbacks
        for callback in self.callbacks:
            try:
                callback(event)
            except Exception as exc:
                logger.error("Swap event callback failed: %s", exc)
        
        return event
    
    def auto_swap(self) -> Optional[SwapEvent]:
        """Automatically swap based on current resource pressure."""
        if self.should_swap_down():
            candidate = self.get_swap_candidate("down")
            if candidate and candidate != self.current_model:
                return self.perform_swap(candidate, "memory_pressure")
        
        # Check if we can upgrade
        elif self.current_model:
            current_spec = self.MODEL_REGISTRY.get(self.current_model)
            if current_spec and current_spec.tier == ModelTier.FAST:
                # We're on fast model, try to upgrade if resources allow
                candidate = self.get_swap_candidate("up")
                if candidate and candidate != self.current_model:
                    return self.perform_swap(candidate, "resources_available")
        
        return None
    
    def swap_to_tier(self, tier: ModelTier) -> Optional[SwapEvent]:
        """Swap to a model of specific capability tier."""
        candidates = [
            name for name, spec in self.MODEL_REGISTRY.items()
            if spec.tier == tier and self.can_swap_up(name)
        ]
        
        if not candidates:
            print(f"[WARN] No models available for tier {tier.value}")
            return None
        
        # Pick largest that fits
        best = max(candidates, key=lambda n: self.MODEL_REGISTRY[n].size_gb)
        return self.perform_swap(best, f"tier_request:{tier.value}")
    
    def get_status(self) -> Dict[str, Any]:
        """Get current hot swap manager status."""
        resources = self.get_system_resources()
        
        return {
            "current_model": self.current_model,
            "resources": resources,
            "should_swap_down": self.should_swap_down(),
            "swap_history_count": len(self.swap_history),
            "recent_swaps": [asdict(e) for e in self.swap_history[-5:]],
            "available_models": list(self.MODEL_REGISTRY.keys()),
            "tiers": [t.value for t in ModelTier]
        }
    
    def register_callback(self, callback: Callable[[SwapEvent], None]):
        """Register callback for swap events."""
        self.callbacks.append(callback)
    
    def save_state(self, path: Optional[Path] = None):
        """Save swap manager state."""
        if path is None:
            path = Path("/home/allaun/Documents/Research Stack/data/hot_swap_state.json")
        
        state = {
            "current_model": self.current_model,
            "swap_history": [asdict(e) for e in self.swap_history],
            "saved_at": datetime.now().isoformat()
        }
        
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            json.dump(state, f, indent=2)
        
        return path


def main():
    """Demonstrate hot swap manager."""
    print("=" * 70)
    print("HOT SWAP MANAGER FOR OLLAMA MODELS")
    print("Memory-aware dynamic model switching")
    print("=" * 70)
    
    manager = HotSwapManager(
        memory_threshold_gb=4.0,
        vram_threshold_gb=1.0
    )
    
    # Show current status
    print("\n📊 CURRENT STATUS")
    print("-" * 50)
    status = manager.get_status()
    print(f"Current Model: {status['current_model'] or 'None'}")
    print(f"Available RAM: {status['resources']['available_ram_gb']:.1f} GB")
    print(f"Available VRAM: {status['resources']['available_vram_gb']:.1f} GB")
    print(f"Swap Down Recommended: {status['should_swap_down']}")
    
    # Show available models
    print("\n📋 REGISTERED MODELS")
    print("-" * 50)
    for name, spec in manager.MODEL_REGISTRY.items():
        print(f"  {name}")
        print(f"    Size: {spec.size_gb:.1f}GB | VRAM: {spec.vram_gb:.1f}GB | Tier: {spec.tier.value}")
    
    # Check if we need to swap
    if status['should_swap_down']:
        print("\n⚠️  MEMORY PRESSURE DETECTED")
        print("-" * 50)
        candidate = manager.get_swap_candidate("down")
        print(f"Recommended swap: {status['current_model']} → {candidate}")
        
        # Perform swap
        event = manager.perform_swap(candidate, "memory_pressure")
        print(f"\nSwap result: {'✓ Success' if event.success else '✗ Failed'}")
        print(f"Duration: {event.duration_ms}ms")
    
    elif status['current_model'] and manager.MODEL_REGISTRY.get(status['current_model'], ModelTier.FAST).tier == ModelTier.FAST:
        print("\n✅ RESOURCES AVAILABLE - CAN UPGRADE")
        print("-" * 50)
        candidate = manager.get_swap_candidate("up")
        if candidate:
            print(f"Recommended upgrade: {status['current_model']} → {candidate}")
            
            # Ask user
            response = input("\nPerform upgrade? (y/n): ").lower()
            if response == 'y':
                event = manager.perform_swap(candidate, "resources_available")
                print(f"Swap result: {'✓ Success' if event.success else '✗ Failed'}")
    
    else:
        print("\n✅ CURRENT STATE OPTIMAL")
        print("-" * 50)
        print("No swap needed at this time.")
    
    # Save state
    path = manager.save_state()
    print(f"\nState saved to: {path}")
    
    print("\n" + "=" * 70)
    print("HOT SWAP MANAGER READY")
    print("=" * 70)
    print("\nUsage:")
    print("  manager = HotSwapManager()")
    print("  manager.perform_swap('deepseek-r1:8b', 'manual')")
    print("  manager.auto_swap()  # Auto-detect and swap")
    print("  manager.swap_to_tier(ModelTier.FAST)  # Tier-based swap")
    
    return manager


if __name__ == "__main__":
    main()