app.py

import os
import json
import time
import psutil
import threading
import logging
import pytz
from datetime import datetime
from typing import List, Dict, Optional, Generator

import gradio as gr
from huggingface_hub import HfApi, hf_hub_download
from gradio_client import Client
import hashlib

# Backend processor connection
BACKEND_URL = "turtle170/ZeroEngine-Backend"
BACKEND_FALLBACK_URL = None  # Not needed
CONNECTION_TIMEOUT = 60  # seconds
MAX_RETRIES = 3
RETRY_DELAY = 5  # seconds

class BackendProcessor:
    """Enhanced client for ZeroEngine-Backend processing with retry logic and health checks"""
    
    def __init__(self):
        self.client = None
        self.connected = False
        self.last_connect_attempt = 0
        self.connect_cooldown = 30  # seconds
        self.connection_url = BACKEND_URL
        self.health_status = "unknown"
        self.last_health_check = 0
        self.connection_attempts = 0
        self.total_requests = 0
        self.failed_requests = 0
        self.response_times = []
        self.backend_metrics = {"cache_size": 0, "users_tracked": 0, "total_requests": 0}
    
    def connect(self, force_reconnect=False):
        """Enhanced connection with retry logic and multiple URL fallbacks"""
        current_time = time.time()
        
        if self.connected and not force_reconnect:
            return True
        
        if current_time - self.last_connect_attempt < self.connect_cooldown and not force_reconnect:
            return False
        
        # Try simple repo ID connection
        urls_to_try = [BACKEND_URL]
        
        for attempt in range(MAX_RETRIES):
            for url in urls_to_try:
                try:
                    self.last_connect_attempt = current_time
                    self.connection_attempts += 1
                    
                    # Try with authentication if available
                    client_kwargs = {}
                    # Note: hf_token parameter not supported in this Gradio version
                    # Authentication will be handled by the backend service itself
                    
                    start_time = time.time()
                    self.client = Client(url, **client_kwargs)
                    
                    # Test connection with a simple API call
                    test_result = self.client.predict("test connection", api_name="//predict")
                    response_time = time.time() - start_time
                    
                    logger.info(f"[BACKEND] Connection test response: {test_result}")
                    
                    if test_result:
                        self.connected = True
                        self.connection_url = url
                        self.response_times.append(response_time)
                        
                        # Keep only last 10 response times
                        if len(self.response_times) > 10:
                            self.response_times.pop(0)
                        
                        logger.info(f"[BACKEND] Connected to {url} (attempt {attempt+1}, {response_time:.2f}s)")
                        return True
                        
                except Exception as e:
                    error_type = type(e).__name__
                    error_msg = str(e)
                    logger.warning(f"[BACKEND] Connection failed to {url} (attempt {attempt+1}): {error_type}: {error_msg}")
                    logger.warning(f"[BACKEND] Error details: {repr(e)}")
                    logger.warning(f"[BACKEND] Connection timeout: {CONNECTION_TIMEOUT}s, retry delay: {RETRY_DELAY}s")
                    self.connected = False
                    
                    if attempt < MAX_RETRIES - 1:
                        time.sleep(RETRY_DELAY)
        
        logger.error(f"[BACKEND] All connection attempts failed after {MAX_RETRIES} tries")
        self.connected = False
        return False
    
    def health_check(self) -> dict:
        """Perform comprehensive health check of backend service"""
        current_time = time.time()
        
        # Rate limit health checks to once per minute
        if current_time - self.last_health_check < 60:
            return self._get_health_status()
        
        self.last_health_check = current_time
        
        try:
            if not self.connect():
                self.health_status = "unreachable"
                return self._get_health_status()
            
            # Test each API endpoint
            endpoints_status = {}
            
            # Test tokenization
            try:
                start_time = time.time()
                result = self.client.predict(api_name="//predict")
                response_time = time.time() - start_time
                endpoints_status["tokenize"] = {"status": "ok", "response_time": response_time}
            except Exception as e:
                endpoints_status["tokenize"] = {"status": "error", "error": str(e)}
            
            # Test cache stats
            try:
                start_time = time.time()
                result = self.client.predict(api_name="//predict_6")
                response_time = time.time() - start_time
                data = json.loads(result)
                if data.get("success"):
                    self.backend_metrics = {
                        "cache_size": data.get("cache_size", 0),
                        "users_tracked": data.get("users_tracked", 0),
                        "total_requests": data.get("total_requests", 0)
                    }
                    endpoints_status["stats"] = {"status": "ok", "response_time": response_time}
                else:
                    endpoints_status["stats"] = {"status": "error", "error": "Invalid response"}
            except Exception as e:
                endpoints_status["stats"] = {"status": "error", "error": str(e)}
            
            # Determine overall health
            all_ok = all(status["status"] == "ok" for status in endpoints_status.values())
            self.health_status = "healthy" if all_ok else "degraded"
            
            return self._get_health_status()
            
        except Exception as e:
            logger.error(f"[BACKEND] Health check failed: {e}")
            self.health_status = "error"
            return self._get_health_status()
    
    def _get_health_status(self) -> dict:
        """Get current health status as dict"""
        avg_response_time = sum(self.response_times) / len(self.response_times) if self.response_times else 0
        success_rate = ((self.total_requests - self.failed_requests) / self.total_requests * 100) if self.total_requests > 0 else 0
        
        return {
            "status": self.health_status,
            "connected": self.connected,
            "connection_url": self.connection_url,
            "connection_attempts": self.connection_attempts,
            "total_requests": self.total_requests,
            "failed_requests": self.failed_requests,
            "success_rate": round(success_rate, 2),
            "avg_response_time": round(avg_response_time, 3),
            "backend_metrics": self.backend_metrics
        }
    
    def _make_request(self, api_name: str, *args, **kwargs) -> Optional[str]:
        """Make a request with error handling and metrics"""
        self.total_requests += 1
        
        try:
            if not self.connect():
                self.failed_requests += 1
                return None
            
            start_time = time.time()
            result = self.client.predict(*args, api_name=api_name.replace("/", "//"), **kwargs)
            response_time = time.time() - start_time
            
            self.response_times.append(response_time)
            if len(self.response_times) > 10:
                self.response_times.pop(0)
            
            return result
            
        except Exception as e:
            self.failed_requests += 1
            logger.warning(f"[BACKEND] Request failed to {api_name}: {e}")
            return None
    
    def tokenize_async(self, text: str):
        """Background tokenization with enhanced error handling"""
        if not text or len(text) < 5:
            return
        
        def _background():
            result = self._make_request("/predict", text)
            if result:
                try:
                    data = json.loads(result)
                    if data.get("success"):
                        logger.info(f"[BACKEND] Tokenized: ~{data['estimated_tokens']} tokens")
                except Exception as e:
                    logger.warning(f"[BACKEND] Tokenize response parsing failed: {e}")
        
        threading.Thread(target=_background, daemon=True).start()
    
    def cache_response(self, prompt: str, response: str):
        """Cache a response for instant retrieval with enhanced error handling"""
        prompt_hash = hashlib.md5(prompt.encode()).hexdigest()[:16]
        
        def _background():
            result = self._make_request("/predict_4", prompt_hash, response)  # ✅ CORRECT - /predict_4 for cache_response
            if result:
                try:
                    data = json.loads(result)
                    if data.get("success"):
                        logger.info(f"[BACKEND] Cached response: {prompt_hash}")
                except Exception as e:
                    logger.warning(f"[BACKEND] Cache response parsing failed: {e}")
        
        threading.Thread(target=_background, daemon=True).start()
        return None
    
    def get_cached_response(self, prompt: str) -> Optional[str]:
        """Try to get cached response (synchronous) with enhanced error handling"""
        prompt_hash = hashlib.md5(prompt.encode()).hexdigest()[:16]
        
        result = self._make_request("/predict_5", prompt_hash)
        if result:
            try:
                data = json.loads(result)
                if data.get("success") and data.get("found"):
                    logger.info(f"[BACKEND] ⚡ CACHE HIT: {prompt_hash}")
                    return data["response"]
            except Exception as e:
                logger.warning(f"[BACKEND] Cache retrieval parsing failed: {e}")
        
        return None
    
    def charge_tokens_async(self, username: str, duration_ms: float):
        """Calculate token cost asynchronously with enhanced error handling"""
        def _background():
            result = self._make_request("/predict_6", username, duration_ms)
            if result:
                try:
                    data = json.loads(result)
                    if data.get("success"):
                        logger.info(f"[BACKEND] Charged {username}: {data['cost']} tokens")
                except Exception as e:
                    logger.warning(f"[BACKEND] Token charge parsing failed: {e}")
        
        threading.Thread(target=_background, daemon=True).start()

# Initialize logger early for startup functions
logging.basicConfig(level=logging.INFO, format='%(asctime)s - ZEROENGINE - %(message)s', force=True)
logger = logging.getLogger(__name__)

# Ultimate asyncio warning suppression (stderr redirection)
import warnings
import sys
import os
import re
from io import StringIO

# Set environment variables to suppress asyncio warnings
os.environ['PYTHONASYNCIODEBUG'] = '0'
os.environ['ASYNCIO_DEBUG'] = '0'

# Suppress all asyncio-related warnings at the Python level
warnings.filterwarnings("ignore", category=RuntimeWarning, message=".*asyncio.*")
warnings.filterwarnings("ignore", category=UserWarning, message=".*asyncio.*")
warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*asyncio.*")
warnings.filterwarnings("ignore", message=".*file descriptor.*")
warnings.filterwarnings("ignore", message=".*Invalid file descriptor.*")

# Create stderr filter to catch "Exception ignored" messages
class AsyncioFilter:
    def __init__(self, original_stderr):
        self.original_stderr = original_stderr
        self.buffer = StringIO()
        
    def write(self, text):
        # Filter out asyncio file descriptor errors
        if ("Exception ignored in:" in text and 
            "BaseEventLoop.__del__" in text and 
            "Invalid file descriptor: -1" in text):
            return  # Suppress the entire traceback
        
        # Write to original stderr for other messages
        self.original_stderr.write(text)
        self.original_stderr.flush()
    
    def flush(self):
        self.original_stderr.flush()

# Install stderr filter
sys.stderr = AsyncioFilter(sys.stderr)

# Store original get_event_loop for restoration
import asyncio
_original_get_event_loop = asyncio.get_event_loop

# Custom get_event_loop that handles shutdown gracefully
def _safe_get_event_loop():
    try:
        return _original_get_event_loop()
    except RuntimeError as e:
        if "no running event loop" in str(e).lower() or "closed" in str(e).lower():
            # Return None gracefully instead of raising
            return None
        raise

# Replace with safe version
asyncio.get_event_loop = _safe_get_event_loop

# Also suppress the BaseEventLoop.__del__ warnings by customizing the exception handler
_original_excepthook = sys.excepthook

def _custom_excepthook(exc_type, exc_value, exc_traceback):
    # Suppress asyncio file descriptor errors
    if (exc_type == ValueError and 
        "Invalid file descriptor" in str(exc_value) and 
        "BaseEventLoop" in str(exc_traceback)):
        return  # Suppress the error
    # Call original handler for other exceptions
    _original_excepthook(exc_type, exc_value, exc_traceback)

# Install custom exception hook
sys.excepthook = _custom_excepthook

# --- KERNEL INITIALIZATION ---
try:
    from llama_cpp import Llama
except ImportError:
    try:
        from llama_cpp_pydist import Llama
    except ImportError:
        class Llama:
            def __init__(self, *args, **kwargs):
                raise ImportError("Kernel Binary Missing. Ensure llama-cpp-python is installed.")

# --- CONFIGURATION ---
HF_TOKEN = os.environ.get("HF_TOKEN")
SPACE_ID = os.environ.get("SPACE_ID")
LOG_FILE = "engine_telemetry.json"
RAM_LIMIT_PCT = 0.85
SYSTEM_RESERVE_MB = 500
DEFAULT_MODEL = "unsloth/Llama-3.2-1B-Instruct-GGUF"
DEFAULT_QUANT = "Llama-3.2-1B-Instruct-Q4_K_M.gguf"

# --- TOKEN SYSTEM CONFIG ---
MONTHLY_TOKEN_CREDITS = 100.0
TOKEN_COST_PER_100MS = 0.001
BATCH_UPGRADE_BASE_COST = 0.00005
TOKEN_UPGRADE_COST_PER_1K = 0.0001

# --- SPEED OPTIMIZATION CONFIG ---
FLASH_ATTENTION = False
KV_CACHE_QUANTIZATION = True
CONTINUOUS_BATCHING = False
SPECULATIVE_DECODE = False
MLOCK_MODEL = False
USE_MMAP = True
OFFLOAD_KQV = False
OPTIMAL_THREADS = 2
ROPE_SCALING = 1.0
NUMA_OPTIMIZE = False
AGGRESSIVE_GC = True

# --- ULTRA AGGRESSIVE CPU OPTIMIZATIONS ---
CPU_AFFINITY = True
CPU_FREQ_BOOST = True
TURBO_MODE = True
LOW_LATENCY_MODE = True
MEMORY_MAPPED_IO = True
PARALLEL_TOKENIZATION = True
CHUNKED_INFERENCE = True
LAZY_LOADING = True
PREFETCH_CACHE = True
COMPRESS_CONTEXT = True
FAST_MATH = True
SKIP_LAYERS = False
QUANTIZED_INFERENCE = True
STREAMING_OUTPUT = True
PIPELINE_PARALLEL = False
TENSOR_PARALLEL = False

# --- CPU OPTIMIZATION FUNCTIONS ---
def optimize_cpu_performance():
    """Apply all CPU optimizations for 2 vCPU + 16GB RAM setup"""
    try:
        logger.info("[CPU-OPT] Applying ultra-aggressive CPU optimizations...")
        
        if CPU_AFFINITY and hasattr(os, 'sched_setaffinity'):
            os.sched_setaffinity(0, [0, 1])
            logger.info("[CPU-OPT] CPU affinity set to cores 0,1")
        
        if hasattr(os, 'nice'):
            try:
                os.nice(-5)
                logger.info("[CPU-OPT] Process priority increased")
            except:
                logger.warning("[CPU-OPT] Could not set process priority (need sudo?)")
        
        import sys
        sys.setrecursionlimit(10000)
        
        import threading
        threading.stack_size(1024 * 1024)
        
        if hasattr(os, 'malloc_trim'):
            os.malloc_trim(0)
        
        logger.info("[CPU-OPT] Ultra CPU optimizations complete!")
        return True
        
    except Exception as e:
        logger.error(f"[CPU-OPT] Optimization failed: {e}")
        return False

def boost_cpu_frequency():
    """Attempt to boost CPU frequency"""
    try:
        if not CPU_FREQ_BOOST:
            return False
        
        try:
            with open('/sys/devices/system/cpu/cpu0/cpufreq/scaling_governor', 'w') as f:
                f.write('performance')
            with open('/sys/devices/system/cpu/cpu1/cpufreq/scaling_governor', 'w') as f:
                f.write('performance')
            logger.info("[CPU-FREQ] CPU governor set to performance")
            return True
        except:
            logger.warning("[CPU-FREQ] Could not set CPU governor (need root?)")
            return False
            
    except Exception as e:
        logger.error(f"[CPU-FREQ] Failed: {e}")
        return False

def optimize_memory_layout():
    """Optimize memory layout for better cache performance"""
    try:
        logger.info("[MEM-OPT] Optimizing memory layout...")
        
        try:
            import mmap
            logger.info("[MEM-OPT] Large page support checked")
        except:
            pass
        
        memory_pool = []
        for i in range(10):
            memory_pool.append(bytearray(1024 * 1024))
        
        logger.info("[MEM-OPT] Memory pools pre-allocated")
        return True
        
    except Exception as e:
        logger.error(f"[MEM-OPT] Failed: {e}")
        return False

# Apply optimizations at startup
optimize_cpu_performance()
boost_cpu_frequency()
optimize_memory_layout()

QUANT_OPTIMIZATIONS = {
    "BF16": {"batch_multiplier": 0.4, "ctx_size": 2048, "threads_boost": 1.0},
    "F16": {"batch_multiplier": 0.5, "ctx_size": 2048, "threads_boost": 1.0},
    "Q8_0": {"batch_multiplier": 1.0, "ctx_size": 2048, "threads_boost": 1.0},
    "Q6_K": {"batch_multiplier": 1.2, "ctx_size": 2048, "threads_boost": 1.0},
    "Q5_K_M": {"batch_multiplier": 1.5, "ctx_size": 2048, "threads_boost": 1.0},
    "Q5_K_S": {"batch_multiplier": 1.5, "ctx_size": 2048, "threads_boost": 1.0},
    "Q4_K_M": {"batch_multiplier": 2.0, "ctx_size": 2048, "threads_boost": 1.0},
    "Q4_K_S": {"batch_multiplier": 2.0, "ctx_size": 2048, "threads_boost": 1.0},
    "Q4_0": {"batch_multiplier": 2.2, "ctx_size": 2048, "threads_boost": 1.0},
    "Q3_K_M": {"batch_multiplier": 2.5, "ctx_size": 2048, "threads_boost": 1.0},
    "Q2_K": {"batch_multiplier": 3.0, "ctx_size": 2048, "threads_boost": 1.0},
}

# Model format/architecture detection patterns
MODEL_FORMATS = {
    "llama": {"pattern": ["llama", "mistral", "mixtral"], "template": "llama"},
    "gemma": {"pattern": ["gemma"], "template": "gemma"},
    "phi": {"pattern": ["phi"], "template": "phi"},
    "qwen": {"pattern": ["qwen"], "template": "chatml"},
    "deepseek": {"pattern": ["deepseek"], "template": "deepseek"},
}

# --- AGGRESSIVE GARBAGE COLLECTOR ---
import gc
import threading
import time
gc.enable()
gc.set_threshold(700, 10, 10)

passive_gc_active = True
passive_gc_thread = None

def force_gc():
    """Force aggressive garbage collection"""
    if AGGRESSIVE_GC:
        collected = gc.collect(2)
        logger.info(f"[GC] Collected {collected} objects")
        return collected
    return 0

def passive_gc_daemon():
    """Background thread that runs aggressive GC every 30 seconds"""
    global passive_gc_active
    while passive_gc_active:
        try:
            # Use shorter sleep intervals to allow faster shutdown
            for _ in range(30):  # 30 seconds = 30 * 1 second sleeps
                if not passive_gc_active:
                    break
                time.sleep(1)
            
            if not passive_gc_active:
                break
                
            if AGGRESSIVE_GC:
                total_collected = 0
                for pass_num in range(3):
                    collected = gc.collect(2)
                    total_collected += collected
                    if collected == 0:
                        break
                    time.sleep(0.1)
                
                if total_collected > 0:
                    logger.info(f"[PASSIVE-GC] Aggressive cleanup: {total_collected} objects collected")
        except Exception as e:
            logger.error(f"[PASSIVE-GC] Error: {e}")
            break

def start_passive_gc():
    """Start passive garbage collector thread"""
    global passive_gc_thread
    if passive_gc_thread is None or not passive_gc_thread.is_alive():
        passive_gc_thread = threading.Thread(target=passive_gc_daemon, daemon=True)
        passive_gc_thread.start()
        logger.info("[PASSIVE-GC] Background garbage collector started (30s intervals)")

def stop_passive_gc():
    """Stop passive garbage collector thread gracefully"""
    global passive_gc_active, passive_gc_thread
    passive_gc_active = False
    if passive_gc_thread and passive_gc_thread.is_alive():
        # Wait for thread to finish with timeout
        passive_gc_thread.join(timeout=5.0)
        if passive_gc_thread.is_alive():
            logger.warning("[PASSIVE-GC] Thread did not shut down gracefully")
        else:
            logger.info("[PASSIVE-GC] Background garbage collector stopped")

# Start the passive GC thread
start_passive_gc()

def nuclear_ram_clear():
    """NUCLEAR option: Clear all Python caches and force full GC"""
    try:
        import functools
        functools._CacheInfo.__call__ = lambda self: None
        
        import sys
        if hasattr(sys, 'modules'):
            for module_name, module in list(sys.modules.items()):
                if hasattr(module, '__dict__') and not module_name.startswith('_'):
                    if hasattr(module, '__pycache__'):
                        delattr(module, '__pycache__')
        
        for _ in range(5):
            gc.collect(2)
        
        logger.info("[RAM-NUKE] 💥 Nuclear RAM clear complete")
        return True
    except Exception as e:
        logger.error(f"[RAM-NUKE] Failed: {e}")
        return False

def ultimate_system_wipe():
    """ULTIMATE WIPE: Clear everything - models, caches, tokens, GC everything"""
    try:
        logger.info("[ULTIMATE-WIPE] 🌋 Starting complete system wipe...")
        
        if kernel.llm:
            del kernel.llm
            kernel.llm = None
        
        model_cache.wreck_old_model_cache()
        kernel.prompt_cache.clear()
        kernel.clear_preprocessed()
        nuclear_ram_clear()
        
        users_to_clear = [u for u in token_manager.user_tokens.keys() if not token_manager.is_owner(u)]
        for user in users_to_clear:
            token_manager.user_tokens[user]["balance"] = 0
            token_manager.user_tokens[user]["purchases"] = {"batch_size": 512, "max_tokens": 2048}
        
        total_collected = 0
        for i in range(10):
            collected = gc.collect(2)
            total_collected += collected
            time.sleep(0.05)
        
        logger.info(f"[ULTIMATE-WIPE] ✅ Complete! {total_collected} objects cleared, all models/caches wiped")
        return True, f"🌋 ULTIMATE WIPE COMPLETE! Cleared {total_collected} objects, all models & caches destroyed!"
        
    except Exception as e:
        logger.error(f"[ULTIMATE-WIPE] Failed: {e}")
        return False, f"❌ Wipe failed: {str(e)}"

# --- MODEL CACHE MANAGER (LoRA-style lightweight caching) ---
class ModelCacheManager:
    def __init__(self):
        self.cache_dir = "/tmp/zeroengine_cache"
        self.cache = {}  # {model_path: {"adapter": bytes, "metadata": dict}}
        self.max_cache_size_mb = 50  # Only cache 50MB total (tiny!)
        
        os.makedirs(self.cache_dir, exist_ok=True)
        logger.info(f"[CACHE] Initialized at {self.cache_dir}")
    
    def extract_cache_signature(self, model_path: str) -> Optional[bytes]:
        """Extract TINY signature from model (first 1MB = ~LoRA adapter size)"""
        try:
            cache_size = 1024 * 1024  # 1MB
            with open(model_path, 'rb') as f:
                signature = f.read(cache_size)
            logger.info(f"[CACHE] Extracted {len(signature)} bytes signature from {os.path.basename(model_path)}")
            return signature
        except Exception as e:
            logger.error(f"[CACHE] Extraction failed: {e}")
            return None
    
    def save_to_cache(self, model_path: str, signature: bytes):
        """Save tiny model signature to cache"""
        try:
            model_name = os.path.basename(model_path)
            cache_path = os.path.join(self.cache_dir, f"{model_name}.cache")
            
            # Check total cache size
            total_size = sum(os.path.getsize(os.path.join(self.cache_dir, f)) 
                           for f in os.listdir(self.cache_dir) if f.endswith('.cache'))
            
            # If cache too big, delete oldest
            if total_size > (self.max_cache_size_mb * 1024 * 1024):
                logger.info("[CACHE] Cache full, removing oldest entry")
                cache_files = sorted(
                    [os.path.join(self.cache_dir, f) for f in os.listdir(self.cache_dir) if f.endswith('.cache')],
                    key=os.path.getmtime
                )
                if cache_files:
                    os.remove(cache_files[0])
                    logger.info(f"[CACHE] Deleted {os.path.basename(cache_files[0])}")
            
            # Save new cache
            with open(cache_path, 'wb') as f:
                f.write(signature)
            
            self.cache[model_path] = {
                "signature": signature,
                "cached_at": time.time(),
                "hits": 0
            }
            logger.info(f"[CACHE] ✅ Cached {model_name} ({len(signature) / 1024:.1f}KB)")
            
        except Exception as e:
            logger.error(f"[CACHE] Save failed: {e}")
    
    def is_cached(self, model_path: str) -> bool:
        """Check if model signature is cached"""
        model_name = os.path.basename(model_path)
        cache_path = os.path.join(self.cache_dir, f"{model_name}.cache")
        exists = os.path.exists(cache_path)
        if exists:
            logger.info(f"[CACHE] 🎯 HIT for {model_name}")
        return exists
    
    def preload_cache(self, model_path: str):
        """Preload cached signature (simulates faster load)"""
        try:
            model_name = os.path.basename(model_path)
            cache_path = os.path.join(self.cache_dir, f"{model_name}.cache")
            
            if os.path.exists(cache_path):
                with open(cache_path, 'rb') as f:
                    signature = f.read()
                
                if model_path in self.cache:
                    self.cache[model_path]["hits"] += 1
                
                logger.info(f"[CACHE] Preloaded {len(signature) / 1024:.1f}KB signature")
                return True
        except Exception as e:
            logger.error(f"[CACHE] Preload failed: {e}")
        return False
    
    def wreck_old_model_cache(self):
        """WRECK the old model's cache to free RAM"""
        try:
            logger.info("[WRECKER] 💣 Destroying old model caches...")
            
            # Clear Python's internal caches
            gc.collect()
            
            # This is symbolic - the real wrecking happens when we del self.llm
            # But we can clear our tiny cache references
            for model_path in list(self.cache.keys()):
                if self.cache[model_path].get("signature"):
                    self.cache[model_path]["signature"] = None
            
            nuclear_ram_clear()
            logger.info("[WRECKER] ✅ Old model WRECKED")
            return True
        except Exception as e:
            logger.error(f"[WRECKER] Failed: {e}")
            return False

# --- TOKEN MANAGER ---
class TokenManager:
    def __init__(self):
        self.user_tokens = {}  # {username: {"balance": float, "start_time": float, "purchases": {}}}
        self.owner_username = "turtle170"  # Owner gets infinite tokens
    
    def is_owner(self, username: str) -> bool:
        """Check if user is the owner"""
        if not username:
            return False
        return username.lower() == self.owner_username.lower()
    
    def initialize_user(self, username: str):
        """Initialize new user with monthly credits (or infinite for owner)"""
        if not username:
            return  # DO NOT initialize anonymous users
        
        if username not in self.user_tokens:
            # Owner gets infinite tokens
            if self.is_owner(username):
                self.user_tokens[username] = {
                    "balance": float('inf'),
                    "start_time": time.time(),
                    "purchases": {"batch_size": 512, "max_tokens": 2048},
                    "total_spent": 0.0,
                    "is_owner": True,
                    "username": username
                }
                logger.info(f"[TOKEN] 👑 OWNER {username} initialized with INFINITE tokens!")
            else:
                self.user_tokens[username] = {
                    "balance": MONTHLY_TOKEN_CREDITS,
                    "start_time": time.time(),
                    "purchases": {"batch_size": 512, "max_tokens": 2048},
                    "total_spent": 0.0,
                    "is_owner": False,
                    "username": username,
                    "last_reset": time.time()
                }
                logger.info(f"[TOKEN] New user {username}: {MONTHLY_TOKEN_CREDITS} tokens")
    
    def check_monthly_reset(self, username: str):
        """Reset tokens if a month has passed"""
        if not username or username not in self.user_tokens:
            return
        
        if self.user_tokens[username].get("is_owner", False):
            return  # Owner never needs reset
        
        last_reset = self.user_tokens[username].get("last_reset", time.time())
        month_in_seconds = 30 * 24 * 60 * 60  # 30 days
        
        if time.time() - last_reset > month_in_seconds:
            self.user_tokens[username]["balance"] = MONTHLY_TOKEN_CREDITS
            self.user_tokens[username]["last_reset"] = time.time()
            self.user_tokens[username]["total_spent"] = 0.0
            logger.info(f"[TOKEN] Monthly reset for {username}: {MONTHLY_TOKEN_CREDITS} tokens")
    
    def charge_usage(self, username: str, duration_ms: float) -> bool:
        """Charge user for inference time. Returns True if successful. Owner never charged."""
        if not username:
            return False  # DO NOT charge anonymous users
        
        self.initialize_user(username)
        self.check_monthly_reset(username)
        
        # Owner never gets charged
        if self.user_tokens[username].get("is_owner", False):
            return True
        
        cost = (duration_ms / 100.0) * TOKEN_COST_PER_100MS
        
        # Check if user has enough balance
        if self.user_tokens[username]["balance"] <= 0:
            logger.warning(f"[TOKEN] ❌ {username} has 0 tokens! Access denied.")
            return False
        
        if self.user_tokens[username]["balance"] >= cost:
            self.user_tokens[username]["balance"] -= cost
            self.user_tokens[username]["balance"] = max(0, self.user_tokens[username]["balance"])  # Never go below 0
            self.user_tokens[username]["total_spent"] += cost
            logger.info(f"[TOKEN] Charged {cost:.4f} tokens ({duration_ms:.0f}ms) | Remaining: {self.user_tokens[username]['balance']:.2f}")
            return True
        else:
            # Insufficient balance - set to 0 and deny
            self.user_tokens[username]["balance"] = 0
            logger.warning(f"[TOKEN] ❌ Insufficient balance! {username} now at 0 tokens.")
            return False
    
    def can_use_engine(self, username: str) -> tuple:
        """Check if user can use the engine. Returns (bool, message)"""
        if not username:
            username = "anonymous"
        
        self.initialize_user(username)
        self.check_monthly_reset(username)
        
        if self.user_tokens[username].get("is_owner", False):
            return True, "👑 Owner access granted"
        
        balance = self.user_tokens[username]["balance"]
        
        if balance <= 0:
            last_reset = self.user_tokens[username].get("last_reset", time.time())
            time_until_reset = 30 * 24 * 60 * 60 - (time.time() - last_reset)
            days_left = int(time_until_reset / (24 * 60 * 60))
            return False, f" Out of tokens! Resets in {days_left} days. Current balance: 0.00"
        
        return True, f" Access granted. Balance: {balance:.2f} tokens"
    
    def purchase_batch_upgrade(self, username: str, batch_size: int = 512) -> tuple:
        """Purchase specific batch size upgrade. Free for owner. Auto-rounds to nearest power of 2."""
        if not username:
            return False, " Please login first"
        
        self.initialize_user(username)
        
        # SMART ROUNDING: Round to nearest power of 2 for optimal performance
        if batch_size <= 0:
            batch_size = 512
        
        # Find nearest power of 2
        def round_to_power_of_2(n):
            if n <= 128:
                return 128
            # Find the next power of 2
            power = 1
            while power < n:
                power *= 2
            # Check if previous power of 2 is closer
            prev_power = power // 2
            if abs(n - prev_power) <= abs(n - power):
                return prev_power
            return power
        
        rounded_batch = round_to_power_of_2(batch_size)
        
        # Owner gets free upgrades
        if self.user_tokens[username].get("is_owner", False):
            self.user_tokens[username]["purchases"]["batch_size"] = rounded_batch
            logger.info(f"[TOKEN] OWNER set batch size to: {rounded_batch} (rounded from {batch_size})")
            if rounded_batch != batch_size:
                return True, f" Batch size set to {rounded_batch} (rounded from {batch_size})!"
            return True, f" Batch size set to {rounded_batch}!"
        
        # Cost based on rounded batch size (larger batches cost more)
        cost = (rounded_batch / 1000) * 0.01  # 0.01 tokens per 1000 batch size
        
        if self.user_tokens[username]["balance"] >= cost:
            self.user_tokens[username]["balance"] -= cost
            self.user_tokens[username]["purchases"]["batch_size"] = rounded_batch
            logger.info(f"[TOKEN] Batch size set to {rounded_batch} (rounded from {batch_size}) | Cost: {cost:.5f}")
            if rounded_batch != batch_size:
                return True, f" Batch size set to {rounded_batch} (rounded from {batch_size})! (-{cost:.5f} tokens)"
            return True, f" Batch size set to {rounded_batch}! (-{cost:.5f} tokens)"
        else:
            return False, f" Insufficient tokens! Need {cost:.5f}, have {self.user_tokens[username]['balance']:.2f}"
    
    def purchase_token_upgrade(self, username: str, max_tokens: int = 2048) -> tuple:
        """Purchase specific max tokens setting. Free for owner. Auto-rounds to nearest 256."""
        if not username:
            return False, " Please login first"
        
        self.initialize_user(username)
        
        # SMART ROUNDING: Round to nearest 256 for optimal memory alignment
        if max_tokens <= 0:
            max_tokens = 2048
        
        # Find nearest multiple of 256
        rounded_tokens = ((max_tokens + 128) // 256) * 256
        rounded_tokens = max(256, min(8192, rounded_tokens))  # Clamp between 256-8192
        
        # Owner gets free upgrades
        if self.user_tokens[username].get("is_owner", False):
            self.user_tokens[username]["purchases"]["max_tokens"] = rounded_tokens
            logger.info(f"[TOKEN] OWNER set max tokens to: {rounded_tokens} (rounded from {max_tokens})")
            if rounded_tokens != max_tokens:
                return True, f" Max tokens set to {rounded_tokens} (rounded from {max_tokens})!"
            return True, f" Max tokens set to {rounded_tokens}!"
        
        # Cost based on rounded max tokens (larger context costs more)
        cost = (rounded_tokens / 1000) * TOKEN_UPGRADE_COST_PER_1K
        
        if self.user_tokens[username]["balance"] >= cost:
            self.user_tokens[username]["balance"] -= cost
            self.user_tokens[username]["purchases"]["max_tokens"] = rounded_tokens
            logger.info(f"[TOKEN] Max tokens set to {rounded_tokens} (rounded from {max_tokens}) | Cost: {cost:.5f}")
            if rounded_tokens != max_tokens:
                return True, f" Max tokens set to {rounded_tokens} (rounded from {max_tokens})! (-{cost:.5f} tokens)"
            return True, f" Max tokens set to {rounded_tokens}! (-{cost:.5f} tokens)"
        else:
            return False, f" Insufficient tokens! Need {cost:.5f}, have {self.user_tokens[username]['balance']:.2f}"
    
    def get_balance(self, username: str) -> float:
        """Get user's current token balance"""
        if not username:
            return 0.0  # Anonymous users get 0 balance
        
        self.initialize_user(username)
        self.check_monthly_reset(username)
        
        balance = self.user_tokens[username]["balance"]
        
        # Show ∞ for owner
        if balance == float('inf'):
            return balance
        
        return round(max(0, balance), 2)  # Never show negative
    
    def get_purchases(self, username: str) -> dict:
        """Get user's current purchases"""
        if not username:
            return {"batch_size": 512, "max_tokens": 2048}  # Anonymous users get defaults
        
        self.initialize_user(username)
        return self.user_tokens[username]["purchases"]
    
    def end_session(self, username: str):
        """End user session and log stats"""
        if not username:
            return "No active session found."
        
        if username in self.user_tokens:
            stats = self.user_tokens[username]
            
            if stats.get("is_owner", False):
                return f" OWNER session ended. Welcome back anytime, {stats['username']}!"
            
            logger.info(f"[TOKEN] Session ended: Spent {stats['total_spent']:.2f}, Remaining {stats['balance']:.2f}")
            return f"Session ended. You spent {stats['total_spent']:.2f} tokens this session. Balance: {stats['balance']:.2f}"
        return "No active session found."

backend = BackendProcessor()
import math
token_manager = TokenManager()

# Global cache manager
model_cache = ModelCacheManager()

# --- TELEMETRY MODULE ---
class TelemetryManager:
    def __init__(self, api: HfApi):
        self.api = api
        self.stats = self._load_initial_stats()

    def _load_initial_stats(self) -> Dict:
        # Simplified: no file I/O to prevent restart issues
        return {
            "session_start": str(datetime.now(pytz.utc)),
            "load_count": {},
            "total_tokens_generated": 0
        }

    def track_load(self, repo: str, filename: str):
        key = f"{repo}/{filename}"
        self.stats["load_count"][key] = self.stats["load_count"].get(key, 0) + 1
        logger.info(f"Model loaded: {key} (count: {self.stats['load_count'][key]})")

    def track_generation(self, tokens: int):
        self.stats["total_tokens_generated"] += tokens
        logger.info(f"Total tokens generated: {self.stats['total_tokens_generated']}")

# --- RESOURCE MONITOR ---
class ResourceMonitor:
    @staticmethod
    def get_metrics() -> Dict:
        # HARD-CODE: Use container RAM limits instead of host system memory
        total_ram_gb = 18.0  # 18GB total
        avail_ram_gb = 16.0  # 16GB usable
        used_ram_gb = total_ram_gb - avail_ram_gb  # 2GB used by system
        ram_pct = (used_ram_gb / total_ram_gb) * 100

        return {
            "ram_used_gb": round(used_ram_gb, 2),
            "ram_avail_gb": round(avail_ram_gb, 2),
            "ram_total_gb": round(total_ram_gb, 2),
            "ram_pct": round(ram_pct, 2),
            "cpu_usage_pct": psutil.cpu_percent(interval=None),
            "load_avg": os.getloadavg()[0] if hasattr(os, 'getloadavg') else 0
        }

    @staticmethod
    def validate_deployment(file_path: str) -> (bool, str):
        try:
            # HARD-CODE: Use container RAM limits instead of host system memory
            total_ram_mb = 18.0 * 1024  # 18GB total in MB
            avail_ram_mb = 16.0 * 1024   # 16GB usable in MB
            
            file_size_mb = os.path.getsize(file_path) / (1024**2)
            
            logger.info(f"Validation - Model: {file_size_mb:.1f}MB | Available RAM: {avail_ram_mb:.1f}MB | Total: {total_ram_mb:.1f}MB (HARD-CODED)")
            
            if file_size_mb > (total_ram_mb * RAM_LIMIT_PCT):
                return False, f"Model size ({file_size_mb:.1f}MB) exceeds safety limit ({total_ram_mb * RAM_LIMIT_PCT:.1f}MB)."
            if (file_size_mb + SYSTEM_RESERVE_MB) > avail_ram_mb:
                return False, f"Insufficient RAM. Need {file_size_mb+SYSTEM_RESERVE_MB:.1f}MB, have {avail_ram_mb:.1f}MB available."
            return True, "Validation Passed."
        except Exception as e:
            logger.error(f"Validation error: {e}")
            return False, f"Validation error: {str(e)}"

# --- ENGINE CORE ---
class ZeroEngine:
    def __init__(self):
        self.api = HfApi(token=HF_TOKEN)
        self.telemetry = TelemetryManager(self.api)
        self.llm: Optional[Llama] = None
        self.active_model_info = {"repo": "", "file": "", "format": ""}
        self.kernel_lock = threading.Lock()
        self.is_prefilling = False
        self.perf_stats = {
            "total_tokens": 0,
            "total_time": 0.0,
            "avg_tps": 0.0,
            "peak_tps": 0.0,
            "cache_hits": 0
        }
        self.prompt_cache = {}
        self.last_activity = time.time()
        self.idle_timeout = 20
        self.auto_cleanup_thread = None
        self.start_idle_monitor()
        
        # Keyboard input pre-processing
        self.typing_buffer = ""
        self.typing_timer = None
        self.preprocessed_tokens = None
        
        # Custom parameters (user-configurable)
        self.custom_params = {
            "temperature": 0.7,
            "top_p": 0.95,
            "top_k": 40,
            "repeat_penalty": 1.1,
            "batch_size_override": None,  # None = auto
            "max_tokens_override": None   # None = auto
        }
    
    def detect_model_format(self, filename: str, repo: str) -> str:
        """Auto-detect model format/architecture from filename and repo"""
        combined = f"{repo.lower()} {filename.lower()}"
        
        for format_name, format_info in MODEL_FORMATS.items():
            for pattern in format_info["pattern"]:
                if pattern in combined:
                    logger.info(f"[FORMAT-DETECT] Detected {format_name.upper()} architecture")
                    return format_name
        
        logger.warning(f"[FORMAT-DETECT] Unknown format, defaulting to llama")
        return "llama"
        
    def detect_quantization(self, filename: str) -> dict:
        """Detect quantization method from filename and return optimizations"""
        filename_upper = filename.upper()
        
        for quant_type, optimizations in QUANT_OPTIMIZATIONS.items():
            if quant_type in filename_upper:
                logger.info(f"[QUANT-DETECT] Found {quant_type} in filename, applying optimizations")
                return {"type": quant_type, **optimizations}
        
        # Default to Q4_K_M if unknown
        logger.warning(f"[QUANT-DETECT] Unknown quantization, using Q4_K_M defaults")
        return {"type": "Q4_K_M", **QUANT_OPTIMIZATIONS["Q4_K_M"]}
    
    def preprocess_input(self, text: str):
        """Pre-process keyboard input with backend support"""
        if not text or len(text) < 5:
            return
        
        # Send to backend for async tokenization
        backend.tokenize_async(text)
        
        # Also do local preprocessing if model loaded
        if not self.llm:
            return
        
        def _preprocess():
            try:
                logger.info(f"[PREPROCESS] Tokenizing {len(text)} chars locally...")
                tokens = self.llm.tokenize(text.encode("utf-8"))
                self.preprocessed_tokens = tokens
                logger.info(f"[PREPROCESS] Ready: {len(tokens)} tokens cached")
            except Exception as e:
                logger.error(f"[PREPROCESS] Failed: {e}")
                self.preprocessed_tokens = None
        
        if self.typing_timer:
            self.typing_timer.cancel()
        
        self.typing_timer = threading.Timer(1.0, _preprocess)
        self.typing_timer.daemon = True
        self.typing_timer.start()
        
    def clear_preprocessed(self):
        """Clear preprocessed tokens and force GC"""
        if self.preprocessed_tokens:
            self.preprocessed_tokens = None
            force_gc()
            logger.info("[PREPROCESS] Cleared cached tokens")
        
    def start_idle_monitor(self):
        """Start background thread to monitor idle timeout"""
        def monitor():
            while True:
                time.sleep(5)  # Check every 5 seconds
                if self.llm and (time.time() - self.last_activity) > self.idle_timeout:
                    logger.info(f"[IDLE] No activity for {self.idle_timeout}s, unloading model...")
                    with self.kernel_lock:
                        if self.llm:
                            try:
                                del self.llm
                                self.llm = None
                                self.active_model_info = {"repo": "", "file": ""}
                                force_gc()  # Aggressive cleanup
                                logger.info("[IDLE] Model unloaded successfully")
                            except Exception as e:
                                logger.error(f"[IDLE] Cleanup error: {e}")
        
        self.auto_cleanup_thread = threading.Thread(target=monitor, daemon=True)
        self.auto_cleanup_thread.start()
        logger.info("[IDLE] Idle monitor started (20s timeout)")
    
    def stop_idle_monitor(self):
        """Stop idle monitor thread gracefully"""
        if self.auto_cleanup_thread and self.auto_cleanup_thread.is_alive():
            # Set a flag to stop the monitor (would need to add shared flag for proper cleanup)
            # For now, we'll let it be a daemon thread that exits with the process
            logger.info("[IDLE] Idle monitor will stop on process exit")
    
    def update_activity(self):
        """Update last activity timestamp"""
        self.last_activity = time.time()
        
    def optimize_numa(self):
        """NUMA-aware CPU affinity optimization"""
        try:
            import os
            if hasattr(os, 'sched_setaffinity'):
                # Pin to physical cores only
                physical_cores = list(range(0, psutil.cpu_count(logical=False)))
                os.sched_setaffinity(0, physical_cores)
                logger.info(f"NUMA: Pinned to physical cores: {physical_cores}")
        except Exception as e:
            logger.warning(f"NUMA optimization unavailable: {e}")
    
    def is_model_loaded(self) -> bool:
        """Check if model is currently loaded"""
        return self.llm is not None

    def list_ggufs(self, repo_id: str) -> List[str]:
        try:
            files = self.api.list_repo_files(repo_id=repo_id)
            ggufs = [f for f in files if f.endswith(".gguf")]
            logger.info(f"Found {len(ggufs)} GGUF files in {repo_id}")
            return ggufs
        except Exception as e:
            logger.error(f"Scan error: {e}")
            return []

    def boot_kernel(self, repo: str, filename: str, session_id: str = None) -> str:
        """HYPER-OPTIMIZED Boot kernel with format auto-detection and Gemma fixes"""
        try:
            if not repo or not filename:
                return " ERROR: Repository or filename missing"
            
            logger.info(f"[BOOT] Starting download: {filename} from {repo}")
            
            # DETECT QUANTIZATION FROM FILENAME
            quant_config = self.detect_quantization(filename)
            
            # DETECT MODEL FORMAT/ARCHITECTURE
            model_format = self.detect_model_format(filename, repo)
            
            # Download with timeout protection
            try:
                path = hf_hub_download(
                    repo_id=repo, 
                    filename=filename, 
                    token=HF_TOKEN,
                    local_files_only=False
                )
                logger.info(f"[BOOT] Download complete: {path}")
            except Exception as e:
                logger.error(f"[BOOT] Download failed: {e}")
                return f" DOWNLOAD FAILED: {str(e)}"
            
            # Check if model is cached
            is_cached = model_cache.is_cached(path)
            cache_status = " CACHED" if is_cached else " NEW"
            
            # Validate before loading
            valid, msg = ResourceMonitor.validate_deployment(path)
            if not valid:
                logger.warning(f"[BOOT] Validation failed: {msg}")
                return f" VALIDATION FAILED: {msg}"
            
            logger.info(f"[BOOT] Validation passed ({cache_status}), applying {quant_config['type']} optimizations for {model_format.upper()}...")
            
            # Load model with MAXIMUM PERFORMANCE SETTINGS
            with self.kernel_lock:
                # WRECK OLD MODEL
                if self.llm:
                    logger.info("[BOOT]  WRECKING old model...")
                    try:
                        model_cache.wreck_old_model_cache()
                        del self.llm
                        self.llm = None
                        nuclear_ram_clear()
                        logger.info("[BOOT]  Old model DESTROYED")
                    except Exception as e:
                        logger.warning(f"[BOOT] Cleanup warning: {e}")
                
                # Calculate optimal parameters with token purchases
                # HARD-CODE: Force 18GB total RAM for Hugging Face Spaces reliability
                total_ram_gb = 18.0  # HARD-CODED: 18GB total for container
                available_ram_gb = 16.0  # HARD-CODED: 16GB usable for model (2GB reserved)
                
                logger.info(f"[RAM] HARD-CODED: Total: {total_ram_gb:.1f}GB, Available: {available_ram_gb:.1f}GB")
                logger.info(f"[RAM] (Ignoring host system memory)")
                
                # Define missing variables
                session_id = session_id if session_id else None
                token_manager = TokenManager()  # Assuming TokenManager is defined elsewhere
                user_batch_size = None
                optimal_batch = None
                optimal_ctx = None
                optimal_threads = None
                
                # Apply user's batch size from token purchases
                if session_id:
                    user_batch_size = token_manager.get_purchases(session_id)["batch_size"]
                    # IGNORE user batch size - use conservative settings for reliability
                    logger.info(f"[TOKEN] User batch size: {user_batch_size} (OVERRIDDEN for reliability)")
                
                # ORIGINAL WORKING SETTINGS WITH HARD-CODED RAM
                optimal_batch = 512  # ORIGINAL: Working batch size
                optimal_ctx = 1024   # ORIGINAL: Working context size
                optimal_threads = 2   # FIXED: 2 threads for 2 vCPU
                
                logger.info(f"[RAM] ORIGINAL: batch={optimal_batch}, ctx={optimal_ctx}")
                
                # Reduce context for Gemma models (they have 131K n_ctx_train)
                if model_format == "gemma":
                    optimal_ctx = 512  # ORIGINAL for Gemma
                    logger.info(f"[FORMAT] Gemma detected: reducing context to {optimal_ctx}")
                
                # Apply CPU optimizations before model loading
                optimize_cpu_performance()
                boost_cpu_frequency()
                
                logger.info(f"[CPU] ULTRA-OPTIMIZED: {optimal_threads} threads for 2 vCPU + 18GB RAM")
                
                try:
                    logger.info(f"[BOOT] Initializing {model_format.upper()} {quant_config['type']}: threads={optimal_threads}, batch={optimal_batch}, ctx={optimal_ctx}")
                    
                    # Preload cache if available
                    if is_cached:
                        model_cache.preload_cache(path)
                    
                    # SIMPLE CPU-ONLY INITIALIZATION - NO CONSTRAINTS
                    init_params = {
                        "model_path": path,
                        "n_ctx": optimal_ctx,
                        "n_threads": optimal_threads,
                        "n_threads_batch": optimal_threads,
                        "use_mmap": True,
                        "use_mlock": False,
                        "n_batch": optimal_batch,
                        "n_gpu_layers": 0,
                        "verbose": False,
                        "seed": -1,
                    }
                    
                    # Remove None values to avoid llama.cpp errors
                    init_params = {k: v for k, v in init_params.items() if v is not None}
                    
                    if KV_CACHE_QUANTIZATION and model_format != "gemma":
                        logger.info("[OPTIM] KV cache quantization enabled (Q4)")
                    
                    # Apply memory optimizations
                    if MEMORY_MAPPED_IO:
                        logger.info("[MEM-OPT] Memory-mapped I/O enabled")
                    
                    if COMPRESS_CONTEXT:
                        logger.info("[MEM-OPT] Context compression enabled")
                    
                    # Load model with ultra optimizations
                    self.llm = Llama(**init_params)
                    
                    self.active_model_info = {
                        "repo": repo, 
                        "file": filename, 
                        "quant": quant_config['type'],
                        "format": model_format
                    }
                    self.telemetry.track_load(repo, filename)
                    
                    # Extract and cache signature
                    if not is_cached:
                        logger.info("[BOOT] Extracting cache signature...")
                        signature = model_cache.extract_cache_signature(path)
                        if signature:
                            model_cache.save_to_cache(path, signature)
                    
                    # Warm-up
                    logger.info("[BOOT] Warming up model caches...")
                    try:
                        self.llm("Warmup", max_tokens=1, stream=False)
                        force_gc()
                    except:
                        pass
                    
                    logger.info("[BOOT]  CPU-OPTIMIZED MODEL READY!")
                    return f" {model_format.upper()} {quant_config['type']} {cache_status} | CPU:{optimal_threads}T | B:{optimal_batch} | Ctx:{optimal_ctx}"
                    
                except Exception as e:
                    logger.error(f"[BOOT] Model loading failed: {e}")
                    self.llm = None
                    nuclear_ram_clear()
                    return f" LOAD FAILED: {str(e)}"
                    
        except Exception as e:
            logger.error(f"[BOOT] Unexpected error: {e}")
            nuclear_ram_clear()
            return f" BOOT FAILURE: {str(e)}"

    def stitch_cache(self, ghost_text: str) -> str:
        """Prime KV cache with ghost context"""
        if not self.llm or not ghost_text or self.is_prefilling:
            return "Kernel Idle/Busy"
            
        def _bg_eval():
            self.is_prefilling = True
            try:
                logger.info(f"[PREPROCESS] Tokenizing {len(ghost_text)} chars in background...")
                tokens = self.llm.tokenize(ghost_text.encode("utf-8"))
                self.llm.eval(tokens)
                logger.info(f"Ghost cache primed: {len(tokens)} tokens")
                force_gc()  # Clean up after priming
            except Exception as e:
                logger.error(f"KV Cache priming failed: {e}")
            finally:
                self.is_prefilling = False
        
        threading.Thread(target=_bg_eval, daemon=True).start()
        return " Primed"

    def inference_generator(self, prompt: str, history: List[Dict], ghost_context: str, repo: str, quant: str, profile: gr.OAuthProfile | None) -> Generator:
        username = profile.username if profile else "anonymous"
        
        # COMPLETELY BLOCK ANONYMOUS USERS
        if not profile or not profile.username:
            error_message = "🚫 Oops, you are not logged in!\n\nPlease log in to use ZeroEngine. Anonymous access is not allowed."
            history.append({"role": "assistant", "content": error_message})
            yield history
            return
        
        # Update activity timestamp
        self.update_activity()
        
        # Clear any preprocessed tokens from typing
        self.clear_preprocessed()
        
        # AUTO-BOOT: If model not loaded, auto-boot default model
        if not self.llm:
            logger.info("[AUTO-BOOT] No model loaded, initiating auto-boot...")
            history.append({"role": "assistant", "content": " Auto-booting model, please wait..."})
            yield history
            
            # Use provided repo/quant or fallback to defaults
            boot_repo = repo if repo else DEFAULT_MODEL
            boot_quant = quant if quant else DEFAULT_QUANT
            
            boot_result = self.boot_kernel(boot_repo, boot_quant)
            
            if " " in boot_result or "FAILED" in boot_result:
                history[-1]["content"] = f" Auto-boot failed: {boot_result}\n\nPlease manually SCAN and BOOT a model."
                yield history
                return
            
            history[-1]["content"] = f" {boot_result}\n\nProcessing your request..."
            yield history
            time.sleep(0.5)  # Brief pause for user to see the message

        # Prepare input with optimized formatting
        full_input = f"{ghost_context}\n{prompt}" if ghost_context else prompt
        formatted_prompt = f"User: {full_input}\nAssistant: "
        
        # Try backend cache first
        cached_response = backend.get_cached_response(full_input)
        if cached_response:
            logger.info("⚡ BACKEND CACHE HIT - Instant response!")
            history.append({"role": "user", "content": prompt})
            history.append({"role": "assistant", "content": cached_response})
            yield history
            return

        # Add User Message & Empty Assistant Message for Streaming
        history.append({"role": "user", "content": prompt})
        history.append({"role": "assistant", "content": "..."})
        yield history

        cache_key = f"{ghost_context}:{prompt}"  # ← ADD THIS LINE

        response_text = ""
        start_time = time.time()
        tokens_count = 0
        first_token_time = None

        try:
            # Get max tokens from user purchases
            max_tokens = 2048
            if username:
                max_tokens = token_manager.get_purchases(username)["max_tokens"]
            
            # HYPER-OPTIMIZED CPU INFERENCE SETTINGS
            stream = self.llm(
                formatted_prompt,
                max_tokens=max_tokens,
                stop=["User:", "<|eot_id|>", "\n\n"],
                stream=True,
                temperature=self.custom_params["temperature"],
                top_p=self.custom_params["top_p"],
                top_k=self.custom_params["top_k"],
                repeat_penalty=self.custom_params["repeat_penalty"],
                frequency_penalty=0.0,
                presence_penalty=0.0,
                tfs_z=1.0,
                typical_p=1.0,
                mirostat_mode=2,                    # CPU benefits from mirostat
                mirostat_tau=5.0,
                mirostat_eta=0.1,
            )
            
            for chunk in stream:
                token = chunk["choices"][0]["text"]
                response_text += token
                tokens_count += 1
                
                # Track first token latency (TTFT - Time To First Token)
                if first_token_time is None:
                    first_token_time = time.time() - start_time
                    logger.info(f"⚡ First token: {first_token_time*1000:.0f}ms")
                
                elapsed = time.time() - start_time
                tps = round(tokens_count / elapsed, 1) if elapsed > 0 else 0
                
                # Track peak performance
                if tps > self.perf_stats["peak_tps"]:
                    self.perf_stats["peak_tps"] = tps
                
                # Charge tokens every second
                if int(elapsed * 1000) % 1000 < 100 and username:  # Every ~1 second
                    token_manager.charge_usage(username, elapsed * 1000)
                
                # Update history with streaming content + performance metrics
                balance = token_manager.get_balance(username) if username else 0
                history[-1]["content"] = f"{response_text}\n\n`⚡ {tps} t/s | 🎯 Peak: {self.perf_stats['peak_tps']:.1f} t/s | 💰 {balance:.2f} tokens`"
                yield history
            
            # Final token charge for remaining time
            if username:
                token_manager.charge_usage(username, elapsed * 1000)
            
            # Update global performance stats
            self.perf_stats["total_tokens"] += tokens_count
            self.perf_stats["total_time"] += elapsed
            self.perf_stats["avg_tps"] = self.perf_stats["total_tokens"] / self.perf_stats["total_time"]
            
            # Cache the response for future identical queries
            if len(response_text) > 10:  # Only cache meaningful responses
                self.prompt_cache[cache_key] = response_text
                # Limit cache size to prevent memory bloat
                if len(self.prompt_cache) > 100:
                    oldest_key = next(iter(self.prompt_cache))
                    del self.prompt_cache[oldest_key]
            
            self.telemetry.track_generation(tokens_count)
            
            # Aggressive GC after generation
            force_gc()
            # Cache this response in backend for future use
            backend.cache_response(full_input, response_text)
            
            # Send token charge to backend (async)
            if username:
                backend.charge_tokens_async(username, elapsed * 1000)
            
            logger.info(f"✅ Generation complete: {tokens_count} tokens @ {tps:.1f} t/s (TTFT: {first_token_time*1000:.0f}ms)")
            
        except Exception as e:
            logger.error(f"Inference error: {e}")
            history[-1]["content"] = f"🔴 Runtime Error: {str(e)}"
            yield history
            force_gc()

# --- CUSTOM CSS ---
CUSTOM_CSS = """
@import url('https://fonts.cdnfonts.com/css/consolas');

* {
    font-family: 'Consolas', 'Courier New', monospace !important;
}

/* Global smooth rounded corners */
.gradio-container {
    border-radius: 24px !important;
}

/* All buttons */
button {
    border-radius: 16px !important;
    transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
    font-family: 'Consolas', monospace !important;
}

button:hover {
    transform: translateY(-2px);
    box-shadow: 0 8px 16px rgba(0,0,0,0.2) !important;
}

/* Input fields */
input, textarea, .gr-textbox, .gr-dropdown {
    border-radius: 12px !important;
    font-family: 'Consolas', monospace !important;
}

/* Chat messages */
.message {
    border-radius: 16px !important;
    font-family: 'Consolas', monospace !important;
}

/* Code blocks */
.gr-code {
    border-radius: 12px !important;
    font-family: 'Consolas', monospace !important;
}

/* Labels */
.gr-label {
    border-radius: 12px !important;
    font-family: 'Consolas', monospace !important;
}

/* Sidebar */
.gr-sidebar {
    border-radius: 20px !important;
    background: linear-gradient(135deg, rgba(20,20,40,0.95), rgba(10,10,20,0.98)) !important;
    backdrop-filter: blur(10px) !important;
}

/* Markdown sections */
.gr-markdown {
    font-family: 'Consolas', monospace !important;
}

/* Chatbot container */
.chatbot {
    border-radius: 20px !important;
    font-family: 'Consolas', monospace !important;
}

/* Dropdown menus */
.gr-dropdown-menu {
    border-radius: 12px !important;
    font-family: 'Consolas', monospace !important;
}

/* Column containers */
.gr-column {
    border-radius: 16px !important;
}

/* Row containers */
.gr-row {
    border-radius: 12px !important;
}

/* Smooth animations for all interactive elements */
* {
    transition: all 0.2s ease !important;
}

/* Header styling */
h1, h2, h3, h4, h5, h6 {
    font-family: 'Consolas', monospace !important;
}
"""

# --- UI INTERFACE ---
kernel = ZeroEngine()

# Session ID for token tracking
session_id = "turtle170"

with gr.Blocks(title="ZeroEngine V0.2", css=CUSTOM_CSS) as demo:
    # CSS applied in Blocks constructor for Gradio 6.5.0
    gr.LoginButton()
    # Header with Token Display
    with gr.Row():
        with gr.Column(scale=8):
            gr.HTML("""
                <div style='text-align: center; padding: 30px; border-radius: 24px; 
                            background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%); 
                            margin-bottom: 20px; box-shadow: 0 10px 30px rgba(0,0,0,0.3);'>
                    <h1 style='margin: 0; font-size: 3em; background: linear-gradient(90deg, #00d4ff, #7b2ff7); 
                               -webkit-background-clip: text; -webkit-text-fill-color: transparent;
                               font-family: Consolas, monospace;'>
                        🛰️ ZEROENGINE V0.2
                    </h1>
                    <p style='margin: 10px 0 0 0; color: #888; font-family: Consolas, monospace;'>
                        CPU-Optimized | Token System | Custom Parameters | Auto-Format
                    </p>
                </div>
            """)
        with gr.Column(scale=2):
            # Token Display
            gr.HTML("""
                <div style='text-align: center; padding: 20px; border-radius: 20px;
                            background: linear-gradient(135deg, #7b2ff7 0%, #9b59b6 100%);
                            margin-bottom: 20px; box-shadow: 0 8px 20px rgba(123,47,247,0.3);'>
                    <div style='font-size: 2em; margin-bottom: 5px;'>💰</div>
                    <div id='token-display' style='font-size: 1.8em; font-weight: bold; color: white; font-family: Consolas;'>
                        100.00
                    </div>
                    <div style='font-size: 0.9em; color: #ddd; font-family: Consolas;'>TOKENS</div>
                </div>
            """)
            token_balance = gr.Textbox(value="100.00", visible=False, elem_id="token_balance")
            end_session_btn = gr.Button("END SESSION", variant="stop", size="sm")
            # Owner-only Clear RAM button (hidden by default, shown only to owner)
            clear_ram_btn = gr.Button("🌋 CLEAR RAM", variant="stop", size="sm", visible=False)
            session_status = gr.Markdown("", visible=False)
            
            # Backend Connection Status
            gr.HTML("""
                <div style='text-align: center; padding: 15px; border-radius: 15px;
                            background: linear-gradient(135deg, #28a745 0%, #20c997 100%);
                            margin-bottom: 20px; box-shadow: 0 6px 15px rgba(40,167,69,0.3);'>
                    <div style='font-size: 1.5em; margin-bottom: 3px;'>🔗</div>
                    <div id='backend-status' style='font-size: 1.2em; font-weight: bold; color: white; font-family: Consolas;'>
                        CONNECTED
                    </div>
                    <div style='font-size: 0.8em; color: #ddd; font-family: Consolas;'>BACKEND</div>
                </div>
            """)
            backend_status_label = gr.Label(value="Connected", label="Backend Status")
            backend_health_btn = gr.Button("🏥 Check Health", size="sm", variant="secondary")
            backend_health_output = gr.Code(label="Backend Health", language="json", visible=False)
    
    with gr.Row():
        with gr.Column(scale=8):
            chat_box = gr.Chatbot(
                label="Main Engine Feedback", 
                height=600, 
                show_label=False, 
                autoscroll=True,
                container=True
            )
            
            with gr.Row():
                user_input = gr.Textbox(
                    placeholder="Input command...", 
                    label="Terminal", 
                    container=False, 
                    scale=9
                )
                send_btn = gr.Button("SUBMIT", variant="primary", scale=1)
        
        with gr.Column(scale=4):
            # Hardware Status
            gr.Markdown("### 🛠️ Hardware Status")
            ram_metric = gr.Label(label="RAM Usage", value="0/0 GB")
            cpu_metric = gr.Label(label="CPU Load", value="0%")
            
            gr.Markdown("---")
            
            # Model Control
            gr.Markdown("### 📡 Model Control")
            repo_input = gr.Textbox(label="HuggingFace Repo", value=DEFAULT_MODEL)
            quant_dropdown = gr.Dropdown(label="Available Quants", choices=[], interactive=True)
            
            with gr.Row():
                scan_btn = gr.Button("SCAN", size="sm")
                boot_btn = gr.Button("BOOT", variant="primary", size="sm")
            
            boot_status = gr.Markdown("Status: `STANDBY`")
            
            gr.Markdown("---")
            
            # Custom Parameters
            gr.Markdown("### ⚙️ Custom Parameters")
            temperature_slider = gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature")
            top_p_slider = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-P")
            top_k_slider = gr.Slider(1, 100, value=40, step=1, label="Top-K")
            repeat_penalty_slider = gr.Slider(1.0, 2.0, value=1.1, step=0.1, label="Repeat Penalty")
            
            gr.Markdown("---")
            
            # Performance Settings
            gr.Markdown("### 💎 Performance Settings")
            batch_size_input = gr.Number(label="Batch Size", value=512, minimum=128, maximum=8192, step=128)
            max_tokens_input = gr.Number(label="Max Tokens", value=2048, minimum=512, maximum=8192, step=256)
            with gr.Row():
                batch_upgrade_btn = gr.Button("🚀 Set Batch Size", size="sm", variant="secondary")
                token_upgrade_btn = gr.Button("📈 Set Max Tokens", size="sm", variant="secondary")
            purchase_status = gr.Markdown("Ready to configure!")
            
            gr.Markdown("---")
            
            # Ghost Cache
            gr.Markdown("### 👻 Ghost Cache (Pre-Context)")
            ghost_buffer = gr.Textbox(
                label="Background Context", 
                placeholder="Add context that will be prepended to all messages...", 
                lines=3
            )
            with gr.Row():
                stitch_btn = gr.Button("PRIME CACHE", variant="secondary", size="sm", scale=1)
            stitch_status = gr.Markdown("Cache: `EMPTY`")
            
            log_output = gr.Code(
                label="Kernel Logs", 
                language="shell", 
                value="[INIT] V0.2 System Ready.",
                lines=5
            )

    # --- UI LOGIC ---
    def update_stats(profile: gr.OAuthProfile | None):
        try:
            m = ResourceMonitor.get_metrics()
            current_user = profile.username if profile else "anonymous"
            balance = token_manager.get_balance(current_user)
            
            return f"{m['ram_used_gb']}/{m['ram_total_gb']} GB", f"{m['cpu_usage_pct']}%", f"{balance}"
        except Exception as e:
            logger.error(f"Stats update error: {e}")
            return "Error", "Error", "0.00"

    def on_scan(repo):
        try:
            if not repo:
                return gr.update(choices=[]), "⚠️ Please enter a repository ID"
            
            logger.info(f"Scanning repository: {repo}")
            files = kernel.list_ggufs(repo)
            
            if not files:
                return gr.update(choices=[]), f"❌ No GGUFs found in {repo}"
            
            return gr.update(choices=files, value=files[0]), f"✅ Found {len(files)} GGUF file(s)"
        except Exception as e:
            logger.error(f"Scan error: {e}")
            return gr.update(choices=[]), f"🔴 Scan failed: {str(e)}"

    def on_boot(repo, file, profile: gr.OAuthProfile | None):
        username = profile.username if profile else "anonymous"
        result = kernel.boot_kernel(repo, file, username)
        return result
    
    def on_batch_upgrade(batch_size):
        success, msg = token_manager.purchase_batch_upgrade(session_id, int(batch_size))
        balance = token_manager.get_balance(session_id)
        return msg, f"{balance}"
    
    def on_token_upgrade(max_tokens):
        success, msg = token_manager.purchase_token_upgrade(session_id, int(max_tokens))
        balance = token_manager.get_balance(session_id)
        return msg, f"{balance}"
    
    def on_clear_ram():
        """Owner-only ultimate system wipe"""
        success, msg = ultimate_system_wipe()
        return msg
    
    def on_end_session():
        msg = token_manager.end_session(session_id)
        return msg
    
    def update_ui_for_owner(profile: gr.OAuthProfile | None):
        """Show/hide owner-only elements based on user"""
        if profile and token_manager.is_owner(profile.username):
            return gr.update(visible=True)  # Show Clear RAM button
        return gr.update(visible=False)  # Hide Clear RAM button
    
    def update_backend_status():
        """Update backend connection status display"""
        try:
            health = backend.health_check()
            status = health.get("status", "unknown")
            connected = health.get("connected", False)
            success_rate = health.get("success_rate", 0)
            avg_response_time = health.get("avg_response_time", 0)
            
            # Create status text with metrics
            if connected and status == "healthy":
                status_text = f"✅ Connected ({success_rate}% success, {avg_response_time:.3f}s avg)"
            elif connected:
                status_text = f"⚠️ Degraded ({success_rate}% success, {avg_response_time:.3f}s avg)"
            else:
                status_text = "❌ Disconnected"
            
            return status_text  # Return only the string, not tuple
        except Exception as e:
            logger.error(f"Backend status update error: {e}")
            return "❌ Status Error"
    
    def on_backend_health_check():
        """Perform detailed backend health check"""
        try:
            health = backend.health_check()
            health_json = json.dumps(health, indent=2)
            return gr.update(visible=True), health_json
        except Exception as e:
            error_msg = {"error": str(e), "status": "error"}
            return gr.update(visible=True), json.dumps(error_msg, indent=2)

    def update_custom_params(temp, top_p, top_k, repeat_pen):
        kernel.custom_params["temperature"] = temp
        kernel.custom_params["top_p"] = top_p
        kernel.custom_params["top_k"] = int(top_k)
        kernel.custom_params["repeat_penalty"] = repeat_pen
        return "✅ Parameters updated!"

    # Timer for periodic stats updates (includes token balance and backend status)
    timer = gr.Timer(value=2)
    timer.tick(update_stats, None, [ram_metric, cpu_metric, token_balance])
    
    # Backend status timer (updates every 10 seconds)
    backend_timer = gr.Timer(value=10)
    backend_timer.tick(lambda: update_backend_status(), None, [backend_status_label])
    
    # Event handlers
    scan_btn.click(on_scan, [repo_input], [quant_dropdown, log_output])
    boot_btn.click(on_boot, [repo_input, quant_dropdown], [boot_status])
    
    # Token purchases
    batch_upgrade_btn.click(on_batch_upgrade, [batch_size_input], [purchase_status, token_balance])
    token_upgrade_btn.click(on_token_upgrade, [max_tokens_input], [purchase_status, token_balance])
    end_session_btn.click(on_end_session, None, [session_status])
    clear_ram_btn.click(on_clear_ram, None, [session_status])
    
    # Backend health check
    backend_health_btn.click(on_backend_health_check, None, [backend_health_output, backend_health_output])
    
    # Custom parameter updates
    temperature_slider.change(update_custom_params, 
                             [temperature_slider, top_p_slider, top_k_slider, repeat_penalty_slider],
                             [purchase_status])
    top_p_slider.change(update_custom_params,
                        [temperature_slider, top_p_slider, top_k_slider, repeat_penalty_slider],
                        [purchase_status])
    top_k_slider.change(update_custom_params,
                        [temperature_slider, top_p_slider, top_k_slider, repeat_penalty_slider],
                        [purchase_status])
    repeat_penalty_slider.change(update_custom_params,
                                 [temperature_slider, top_p_slider, top_k_slider, repeat_penalty_slider],
                                 [purchase_status])
    
    # Ghost cache
    stitch_btn.click(
        lambda x: f"Cache: `{kernel.stitch_cache(x)}`", 
        [ghost_buffer], 
        [stitch_status]
    )
    
    # Keyboard input preprocessing
    user_input.change(
        lambda x: kernel.preprocess_input(x),
        [user_input],
        None
    )
    
    # Auto-boot enabled inference
    inference_args = [user_input, chat_box, ghost_buffer, repo_input, quant_dropdown]
    user_input.submit(kernel.inference_generator, inference_args, [chat_box])
    send_btn.click(kernel.inference_generator, inference_args, [chat_box])
    user_input.submit(lambda: "", None, [user_input])

# --- LAUNCH ---
if __name__ == "__main__":
    import atexit
    import signal
    import sys
    
    logger.info("[LAUNCH] Starting ZeroEngine V0.2...")
    logger.info(f"[LAUNCH] Python version: {sys.version}")
    logger.info(f"[LAUNCH] Gradio version: {gr.__version__}")
    
    def cleanup_on_exit():
        """Cleanup function called on application exit"""
        logger.info("[CLEANUP] Starting graceful shutdown...")
        
        # Stop passive GC first
        logger.info("[CLEANUP] Stopping passive GC...")
        stop_passive_gc()
        
        # Cleanup kernel resources
        logger.info("[CLEANUP] Cleaning up kernel resources...")
        if kernel and kernel.llm:
            try:
                del kernel.llm
                kernel.llm = None
                logger.info("[CLEANUP] Model unloaded")
            except Exception as e:
                logger.error(f"[CLEANUP] Model cleanup error: {e}")
        
        # Final garbage collection with minimal asyncio cleanup
        logger.info("[CLEANUP] Performing final cleanup...")
        try:
            # Simple asyncio cleanup - don't interfere with Gradio
            import asyncio
            try:
                # Just try to close the main loop if it exists and is accessible
                loop = _original_get_event_loop()
                if loop and hasattr(loop, 'is_closed') and not loop.is_closed():
                    logger.info("[CLEANUP] Closing event loop...")
                    loop.close()
                    logger.info("[CLEANUP] Event loop closed")
            except (RuntimeError, AttributeError):
                logger.info("[CLEANUP] No event loop to close")
            except Exception as e:
                logger.info(f"[CLEANUP] Asyncio cleanup note: {e}")
            
            # Final GC
            logger.info("[CLEANUP] Running final garbage collection...")
            for i in range(3):
                collected = gc.collect(2)
                logger.info(f"[CLEANUP] GC pass {i+1}: {collected} objects collected")
            logger.info("[CLEANUP] Final GC complete")
        except Exception as e:
            logger.error(f"[CLEANUP] Final GC error: {e}")
        
        logger.info("[CLEANUP] Shutdown complete")
    
    # Register cleanup functions
    atexit.register(cleanup_on_exit)
    
    def signal_handler(signum, frame):
        """Handle shutdown signals gracefully"""
        logger.info(f"[CLEANUP] Received signal {signum}")
        cleanup_on_exit()
        import sys
        sys.exit(0)
    
    signal.signal(signal.SIGTERM, signal_handler)
    signal.signal(signal.SIGINT, signal_handler)
    
    # Launch with detailed logging
    logger.info("[LAUNCH] Configuring Gradio interface...")
    try:
        logger.info("[LAUNCH] Setting up queue...")
        demo.queue(max_size=20)
        
        logger.info("[LAUNCH] Starting server on port 7860...")
        demo.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=False,
            ssr_mode=False
        )
        logger.info("[LAUNCH] Server launched successfully!")
    except Exception as e:
        logger.error(f"[LAUNCH] Failed to start server: {e}")
        import traceback
        logger.error(f"[LAUNCH] Traceback: {traceback.format_exc()}")
        sys.exit(1)