"""RAM Optimizer and Model Pre-warmer for High-Speed Switching."""
import os
import glob
import time
import httpx
import logging
from typing import Dict, List, Any

logger = logging.getLogger("ram_optimizer")

OLLAMA_API_BASE = "http://localhost:11434"
COMFY_API_BASE = "http://127.0.0.1:8188"
COMFY_MODELS_DIR = "/home/drjones/ComfyUI/models"

def get_detailed_meminfo() -> Dict[str, Any]:
    """Parse /proc/meminfo for precise page cache and RAM stats."""
    info = {}
    try:
        with open("/proc/meminfo", "r") as f:
            for line in f:
                parts = line.split(":")
                if len(parts) == 2:
                    key = parts[0].strip()
                    val = parts[1].strip().split()[0]
                    info[key] = int(val) * 1024  # Convert kB to bytes
    except Exception as e:
        logger.error(f"Failed to read /proc/meminfo: {e}")
    
    total = info.get("MemTotal", 0)
    free = info.get("MemFree", 0)
    available = info.get("MemAvailable", 0)
    cached = info.get("Cached", 0) + info.get("Buffers", 0)
    dirty = info.get("Dirty", 0)
    used = total - free - cached
    if used < 0:
        used = total - available

    return {
        "total_bytes": total,
        "total_gb": round(total / (1024**3), 2),
        "used_bytes": used,
        "used_gb": round(used / (1024**3), 2),
        "cached_bytes": cached,
        "cached_gb": round(cached / (1024**3), 2),
        "free_bytes": free,
        "free_gb": round(free / (1024**3), 2),
        "available_bytes": available,
        "available_gb": round(available / (1024**3), 2),
        "dirty_bytes": dirty,
        "dirty_mb": round(dirty / (1024**2), 2),
        "cache_ratio_pct": round((cached / total * 100) if total > 0 else 0, 1),
    }

def warm_file_to_ram(filepath: str, chunk_size: int = 16 * 1024 * 1024) -> Dict[str, Any]:
    """Pre-fault/read file into Linux OS Page Cache at maximum disk read speed."""
    if not os.path.exists(filepath):
        return {"success": False, "error": f"File not found: {filepath}", "duration_ms": 0}
    
    t0 = time.perf_counter()
    file_size = os.path.getsize(filepath)
    bytes_read = 0
    try:
        with open(filepath, "rb") as f:
            # Hint kernel that we will read this sequentially
            try:
                os.posix_fadvise(f.fileno(), 0, file_size, os.POSIX_FADV_WILLNEED)
            except Exception:
                pass
            
            buf = bytearray(chunk_size)
            while True:
                n = f.readinto(buf)
                if not n:
                    break
                bytes_read += n
        
        duration = time.perf_counter() - t0
        duration_ms = round(duration * 1000, 2)
        speed_mb_s = round((bytes_read / (1024**2)) / duration if duration > 0 else 0, 2)
        return {
            "success": True,
            "filepath": filepath,
            "size_bytes": file_size,
            "size_mb": round(file_size / (1024**2), 2),
            "bytes_read": bytes_read,
            "duration_ms": duration_ms,
            "speed_mb_s": speed_mb_s,
        }
    except Exception as e:
        return {"success": False, "error": str(e), "duration_ms": round((time.perf_counter() - t0) * 1000, 2)}

async def warm_ollama_model(model_name: str, keep_alive: str = "5m") -> Dict[str, Any]:
    """Warm an Ollama model into memory and measure time."""
    t0 = time.perf_counter()
    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            resp = await client.post(
                f"{OLLAMA_API_BASE}/api/generate",
                json={"model": model_name, "prompt": "", "keep_alive": keep_alive},
            )
            duration = time.perf_counter() - t0
            if resp.status_code == 200:
                data = resp.json()
                return {
                    "success": True,
                    "model": model_name,
                    "duration_ms": round(duration * 1000, 2),
                    "load_duration_ms": round(data.get("load_duration", 0) / 1e6, 2),
                    "total_duration_ms": round(data.get("total_duration", 0) / 1e6, 2),
                }
            else:
                return {
                    "success": False,
                    "model": model_name,
                    "error": f"HTTP {resp.status_code}: {resp.text}",
                    "duration_ms": round(duration * 1000, 2),
                }
    except Exception as e:
        return {"success": False, "model": model_name, "error": str(e), "duration_ms": round((time.perf_counter() - t0) * 1000, 2)}

def find_comfy_model_files() -> List[Dict[str, Any]]:
    """Discover all model files under ComfyUI models."""
    results = []
    extensions = ("*.safetensors", "*.ckpt", "*.pt", "*.bin")
    if os.path.exists(COMFY_MODELS_DIR):
        for root, _, files in os.walk(COMFY_MODELS_DIR):
            for file in files:
                if any(file.endswith(ext.replace("*", "")) for ext in extensions):
                    full_path = os.path.join(root, file)
                    rel_path = os.path.relpath(full_path, COMFY_MODELS_DIR)
                    size = os.path.getsize(full_path)
                    results.append({
                        "filename": file,
                        "rel_path": rel_path,
                        "full_path": full_path,
                        "size_bytes": size,
                        "size_mb": round(size / (1024**2), 2),
                        "size_gb": round(size / (1024**3), 3),
                    })
    return results

async def warm_all_models() -> Dict[str, Any]:
    """Warm all available Ollama and ComfyUI models into Linux RAM Cache."""
    t0 = time.perf_counter()
    warmed_ollama = []
    warmed_comfy = []
    
    # 1. Ollama models
    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            tags_resp = await client.get(f"{OLLAMA_API_BASE}/api/tags")
            if tags_resp.status_code == 200:
                models = tags_resp.json().get("models", [])
                for m in models:
                    name = m.get("name")
                    res = await warm_ollama_model(name, keep_alive="1m")
                    warmed_ollama.append(res)
    except Exception as e:
        logger.error(f"Error discovering Ollama models: {e}")
    
    # 2. ComfyUI models
    comfy_files = find_comfy_model_files()
    for f in comfy_files:
        res = warm_file_to_ram(f["full_path"])
        warmed_comfy.append(res)
        
    total_duration_ms = round((time.perf_counter() - t0) * 1000, 2)
    meminfo = get_detailed_meminfo()
    
    return {
        "status": "completed",
        "total_duration_ms": total_duration_ms,
        "ollama_models_warmed": warmed_ollama,
        "comfy_files_warmed": warmed_comfy,
        "meminfo_after": meminfo,
    }