#!/usr/bin/env python3
"""
Mini COT Service — Chain-of-Thought Analysis for SpecMem

Unix socket service using Pythia-410M for codebase analysis.
Two modes:
  - Crawl mode (default): 5% CPU, ~600MB RAM. Reads code_chunks, cross-refs docs,
    stores findings in codebase_training table.
  - Active scoring: Springs to action for COT scoring requests from find_code_pointers.

Resource Management:
  - Shares QQMS throttling with Frankenstein embedding server
  - Respects SPECMEM_CPU_MIN/MAX, SPECMEM_RAM_MIN_MB/MAX_MB
  - Layer-by-layer model loading for snappy startup
  - Pauses crawl during heavy embedding ops (QQMS coordination)
  - CPU core pinning via SPECMEM_CPU_CORES_MIN/MAX

Socket Protocol:
  {"type":"health"} → {"status":"ok","mode":"crawl"}
  {"type":"score","code":"...","query":"..."} → {"score":0.85,"reasoning":"..."}
  {"type":"status"} → {"files_analyzed":150,"bugs_found":3}
  {"type":"pause"} → {"status":"paused"} (QQMS coordination)
  {"type":"resume"} → {"status":"resumed"}

@author hardwicksoftwareservices
@website https://justcalljon.pro

TODO: This file is 1000+ lines - consider splitting into mini_cot_service.py, cot_analyzer.py, cot_crawler.py
"""

import os
import sys
import json
import time
import socket
import signal
import struct
import threading
import traceback
import argparse
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, Dict, Any, List

# ============================================================================
# CPU Thread Governance — matches frankenstein-embeddings.py pattern
# Priority: SPECMEM_CPU_THREADS env → user-config.json resources.cpuCoreMax → default 1
# ============================================================================

PROJECT_PATH = os.environ.get('SPECMEM_PROJECT_PATH', os.getcwd())

def _get_cpu_thread_limit():
    """Get CPU thread limit from env or user-config.json"""
    if os.environ.get('SPECMEM_CPU_THREADS'):
        return int(os.environ['SPECMEM_CPU_THREADS'])
    try:
        config_path = os.path.join(PROJECT_PATH, 'specmem', 'user-config.json')
        if os.path.exists(config_path):
            with open(config_path, 'r') as f:
                config = json.load(f)
                core_max = config.get('resources', {}).get('cpuCoreMax')
                if core_max is not None:
                    return int(core_max)
    except Exception as e:
        print(f"⚠️ Could not read CPU core limit from config: {e}", file=sys.stderr)
    return 1

_CPU_THREAD_LIMIT = _get_cpu_thread_limit()

# Set thread limits at module level BEFORE any model imports
os.environ.setdefault('OMP_NUM_THREADS', str(_CPU_THREAD_LIMIT))
os.environ.setdefault('MKL_NUM_THREADS', str(_CPU_THREAD_LIMIT))
os.environ.setdefault('OPENBLAS_NUM_THREADS', str(_CPU_THREAD_LIMIT))
os.environ.setdefault('NUMEXPR_NUM_THREADS', str(_CPU_THREAD_LIMIT))
print(f"🔒 Mini COT CPU threads: {_CPU_THREAD_LIMIT} (cpucoremax to adjust)", file=sys.stderr)

# ============================================================================
# Resource Configuration — shared with Frankenstein embedding server
# ============================================================================

@dataclass
class MiniCOTResourceConfig:
    """
    Resource limits shared with embedding server via env vars.
    Mini COT uses a SUBSET of the total budget (crawl = 5% CPU target).
    """
    # CPU limits from shared env
    cpu_min: float = float(os.environ.get('SPECMEM_CPU_MIN', '20'))
    cpu_max: float = float(os.environ.get('SPECMEM_CPU_MAX', '40'))

    # RAM limits (MB) from shared env
    ram_min_mb: float = float(os.environ.get('SPECMEM_RAM_MIN_MB', '4000'))
    ram_max_mb: float = float(os.environ.get('SPECMEM_RAM_MAX_MB', '6000'))

    # CPU core limits
    cpu_cores_min: int = int(os.environ.get('SPECMEM_CPU_CORES_MIN', '1'))
    cpu_cores_max: int = int(os.environ.get('SPECMEM_CPU_CORES_MAX', str(os.cpu_count() or 2)))

    # Mini COT's share of the resource budget
    # Crawl mode: use at most 5% of CPU budget, 600MB RAM
    crawl_cpu_percent: float = 5.0
    crawl_ram_cap_mb: float = 600.0

    # Active scoring: temporarily use up to 15% CPU
    active_cpu_percent: float = 15.0
    active_ram_cap_mb: float = 800.0

    # Thread limits for torch
    crawl_threads: int = 1
    active_threads: int = max(1, min(2, _CPU_THREAD_LIMIT))

    # Sleep between crawl chunks (seconds) — pacing for 5% CPU
    crawl_sleep_base: float = 2.0
    crawl_sleep_max: float = 10.0

    def get_crawl_threads(self) -> int:
        """Threads for crawl mode — always minimal"""
        return min(self.crawl_threads, self.cpu_cores_max)

    def get_active_threads(self) -> int:
        """Threads for active scoring — bounded by core limits"""
        return min(self.active_threads, self.cpu_cores_max)


def get_system_ram_mb() -> float:
    try:
        with open('/proc/meminfo', 'r') as f:
            for line in f:
                if line.startswith('MemTotal:'):
                    return int(line.split()[1]) / 1024
    except:
        pass
    return 4000.0


class LowResourceConfig:
    """Auto-detect low-resource environments and adjust accordingly."""
    def __init__(self):
        total_ram = get_system_ram_mb()
        self.is_low_resource = total_ram < 4000
        self.lazy_loading = True  # Always lazy load
        self.disk_cache = self.is_low_resource
        self.aggressive_cleanup = self.is_low_resource
        self.idle_unload_seconds = 300 if self.is_low_resource else 0
        if self.is_low_resource:
            print(f"⚠️ Low resource mode: {total_ram:.0f}MB RAM — aggressive cleanup enabled", file=sys.stderr)


def get_available_ram_mb() -> float:
    try:
        with open('/proc/meminfo', 'r') as f:
            for line in f:
                if line.startswith('MemAvailable:'):
                    return int(line.split()[1]) / 1024
    except:
        pass
    return 1000.0


def get_cpu_percent() -> float:
    """Get current CPU usage (0-100)"""
    try:
        with open('/proc/stat', 'r') as f:
            line = f.readline()
        parts = line.split()
        idle = int(parts[4])
        total = sum(int(p) for p in parts[1:])
        # Need two samples
        time.sleep(0.1)
        with open('/proc/stat', 'r') as f:
            line = f.readline()
        parts2 = line.split()
        idle2 = int(parts2[4])
        total2 = sum(int(p) for p in parts2[1:])
        idle_delta = idle2 - idle
        total_delta = total2 - total
        if total_delta == 0:
            return 0.0
        return (1.0 - idle_delta / total_delta) * 100
    except:
        return 0.0


# ============================================================================
# QQMS Coordination — pause/resume based on embedding server load
# ============================================================================

class QQMSCoordinator:
    """
    Coordinates with the Frankenstein embedding server via shared state.
    When embedding server is under heavy load, Mini COT pauses crawling.
    """
    def __init__(self, resource_config: MiniCOTResourceConfig):
        self.config = resource_config
        self.paused = False
        self.pause_reason = ""
        self._lock = threading.Lock()

    def should_pause(self) -> bool:
        """Check if we should pause based on system resources"""
        cpu = get_cpu_percent()
        ram = get_available_ram_mb()

        # Pause if CPU exceeds our budget
        if cpu > self.config.cpu_max:
            self.pause_reason = f"CPU {cpu:.1f}% > max {self.config.cpu_max}%"
            return True

        # Pause if RAM is getting tight
        if ram < self.config.ram_min_mb * 0.3:
            self.pause_reason = f"RAM {ram:.0f}MB < threshold"
            return True

        return False

    def get_crawl_delay(self) -> float:
        """Dynamic sleep between crawl operations based on system load"""
        cpu = get_cpu_percent()
        if cpu > self.config.cpu_max * 0.8:
            return self.config.crawl_sleep_max
        elif cpu > self.config.cpu_max * 0.5:
            ratio = (cpu - self.config.cpu_max * 0.5) / (self.config.cpu_max * 0.3)
            return self.config.crawl_sleep_base + ratio * (self.config.crawl_sleep_max - self.config.crawl_sleep_base)
        return self.config.crawl_sleep_base

    def pause(self):
        with self._lock:
            self.paused = True

    def resume(self):
        with self._lock:
            self.paused = False
            self.pause_reason = ""

    @property
    def is_paused(self) -> bool:
        with self._lock:
            return self.paused


# ============================================================================
# Model Manager — Layer-by-layer loading for snappy startup
# ============================================================================

class ModelManager:
    """
    Dual-mode model manager:

    1. ONNX SCORING (instant) — loads first, handles find_code_pointers scoring
       - Quantized ONNX model (394MB int8) via onnxruntime
       - Forward pass + cosine similarity = ~50-100ms per score
       - Works on i3 laptops, containers, everywhere

    2. PYTORCH ANALYSIS (crawl) — loads lazily when crawl engine needs it
       - Full Pythia-410M for text generation
       - Analyzes code chunks, identifies bugs, writes descriptions
       - Background crawl, can take time — that's fine

    Scoring is always instant. Analysis loads when there's RAM and work to do.
    """
    def __init__(self, model_name: str, device: str, resource_config: MiniCOTResourceConfig):
        self.model_name = model_name
        self.device = device
        self.config = resource_config

        # ONNX scoring (fast path)
        self.onnx_session = None
        self.tokenizer = None
        self.onnx_output_names = None
        self.onnx_input_names = None
        self.is_loaded = False      # ONNX scoring ready
        self.is_loading = False
        self.load_progress = 0.0

        # PyTorch analysis (crawl path)
        self.torch_model = None
        self.torch_loaded = False
        self.torch_loading = False

        self._lock = threading.Lock()
        self._torch_lock = threading.Lock()

    def ensure_loaded(self):
        """Load ONNX scorer if not loaded. Thread-safe."""
        if self.is_loaded:
            return
        with self._lock:
            if self.is_loaded:
                return
            self._load_onnx()

    def ensure_torch_loaded(self):
        """Load PyTorch model for text generation (crawl analysis). Thread-safe."""
        # Needs tokenizer from ONNX load first
        self.ensure_loaded()
        if self.torch_loaded:
            return
        with self._torch_lock:
            if self.torch_loaded:
                return
            self._load_torch()

    def _find_model_path(self):
        """Find the pre-packed ONNX model directory"""
        pre_packed = os.environ.get('SPECMEM_MODEL_CACHE')
        if pre_packed and os.path.isdir(pre_packed):
            return pre_packed
        script_dir = os.path.dirname(os.path.abspath(__file__))
        candidate = os.path.join(script_dir, 'models', 'pythia-410m-onnx-quant')
        if os.path.isdir(candidate):
            return candidate
        candidate = '/app/models/pythia-410m-onnx-quant'
        if os.path.isdir(candidate):
            return candidate
        return None

    def _load_onnx(self):
        """Load ONNX quantized model for instant scoring"""
        self.is_loading = True
        try:
            import onnxruntime as ort
            import numpy as np

            model_dir = self._find_model_path()
            if not model_dir:
                raise FileNotFoundError("Pre-packed Pythia ONNX model not found")

            onnx_path = os.path.join(model_dir, 'model_quantized.onnx')
            if not os.path.exists(onnx_path):
                raise FileNotFoundError(f"ONNX model not found at {onnx_path}")

            print(f"⚡ Mini COT: Loading ONNX scorer...", file=sys.stderr)
            self.load_progress = 0.1

            from transformers import AutoTokenizer
            self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
            if self.tokenizer.pad_token is None:
                self.tokenizer.pad_token = self.tokenizer.eos_token
            self.load_progress = 0.3
            print(f"  Tokenizer loaded", file=sys.stderr)

            # ONNX session — optimized for CPU
            sess_opts = ort.SessionOptions()
            sess_opts.inter_op_num_threads = 1
            sess_opts.intra_op_num_threads = max(1, self.config.get_active_threads())
            sess_opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
            sess_opts.enable_cpu_mem_arena = False
            sess_opts.enable_mem_pattern = True

            self.onnx_session = ort.InferenceSession(
                onnx_path, sess_options=sess_opts, providers=['CPUExecutionProvider']
            )
            self.onnx_output_names = [o.name for o in self.onnx_session.get_outputs()]
            self.onnx_input_names = [i.name for i in self.onnx_session.get_inputs()]
            self.load_progress = 0.8

            # Warmup forward pass
            self._get_hidden_state("warmup")

            self.load_progress = 1.0
            self.is_loaded = True
            print(f"⚡ ONNX scorer ready — instant scoring enabled", file=sys.stderr)
            print(f"   Threads: {sess_opts.intra_op_num_threads}, RAM: {get_available_ram_mb():.0f}MB", file=sys.stderr)

        except Exception as e:
            print(f"❌ ONNX scoring load failed: {e}", file=sys.stderr)
            traceback.print_exc(file=sys.stderr)
            raise
        finally:
            self.is_loading = False

    def _load_torch(self):
        """Load model for text generation (crawl analysis) — prefers local ONNX via optimum"""
        self.torch_loading = True
        try:
            import torch
            torch.set_num_threads(self.config.get_crawl_threads())

            available_ram = get_available_ram_mb()
            if available_ram < self.config.crawl_ram_cap_mb * 1.5:
                print(f"⚠️ Not enough RAM for generation model ({available_ram:.0f}MB), skipping", file=sys.stderr)
                return

            model_dir = self._find_model_path()
            onnx_path = os.path.join(model_dir, 'model_quantized.onnx') if model_dir else None

            if model_dir and onnx_path and os.path.exists(onnx_path):
                # Use pre-packed ONNX model via optimum (no network required)
                print(f"🧠 Loading ONNX generation model from {model_dir}...", file=sys.stderr)
                from optimum.onnxruntime import ORTModelForCausalLM
                self.torch_model = ORTModelForCausalLM.from_pretrained(
                    model_dir,
                    file_name="model_quantized.onnx"
                )
            else:
                # No local model found — refuse to download from the internet
                search_paths = [
                    os.environ.get('SPECMEM_MODEL_CACHE', '(not set)'),
                    os.path.join(os.path.dirname(__file__), '..', 'models', 'pythia-410m-onnx-quant'),
                    '/app/models/pythia-onnx-quant',
                ]
                raise RuntimeError(
                    f"Local ONNX model not found. Searched:\n"
                    + "\n".join(f"  - {p}" for p in search_paths)
                    + "\n\nRun `specmem init` to download models via Git LFS release tarball."
                    + "\nSpecMem will NOT download models from the internet at runtime."
                )

            self.torch_loaded = True
            print(f"🧠 Generation model loaded for crawl analysis", file=sys.stderr)
            print(f"   Threads: {torch.get_num_threads()}, RAM: {get_available_ram_mb():.0f}MB", file=sys.stderr)

        except Exception as e:
            print(f"⚠️ Generation model load failed: {e} — crawl analysis disabled", file=sys.stderr)
        finally:
            self.torch_loading = False

    def _get_hidden_state(self, text: str, max_length: int = 512):
        """
        Single forward pass → last-token logits as semantic fingerprint.
        The logit distribution encodes what the model predicts comes next,
        which is a rich representation of the input's meaning.
        Returns: L2-normalized numpy array
        """
        import numpy as np

        inputs = self.tokenizer(
            text, return_tensors="np", truncation=True,
            max_length=max_length, padding=False
        )

        seq_len = inputs['input_ids'].shape[1]

        # Build feed dict — handle all required inputs including KV cache
        feed = {}
        for name in self.onnx_input_names:
            if name == 'input_ids':
                feed[name] = inputs['input_ids'].astype(np.int64)
            elif name == 'attention_mask':
                feed[name] = inputs['attention_mask'].astype(np.int64)
            elif name == 'position_ids':
                # Position IDs: [0, 1, 2, ..., seq_len-1]
                feed[name] = np.arange(seq_len, dtype=np.int64).reshape(1, -1)
            elif 'past_key_values' in name:
                # Empty KV cache for initial pass
                # Shape: (batch=1, num_heads=16, past_seq_len=0, head_dim=64)
                # Pythia-410M: 16 heads, hidden_size=1024 → head_dim=64
                feed[name] = np.zeros((1, 16, 0, 64), dtype=np.float32)

        outputs = self.onnx_session.run(self.onnx_output_names, feed)
        # First output is logits: (batch, seq_len, vocab_size)
        logits = outputs[0][0, -1, :]
        norm = np.linalg.norm(logits)
        if norm > 0:
            logits = logits / norm
        return logits

    def score_similarity(self, query: str, code: str) -> float:
        """
        Score code relevance via ONNX hidden-state cosine similarity.
        Two forward passes (~50-100ms total on CPU). Instant.
        """
        self.ensure_loaded()
        import numpy as np

        query_vec = self._get_hidden_state(f"Query: {query}")
        code_vec = self._get_hidden_state(f"Code: {code[:1500]}")

        similarity = float(np.dot(query_vec, code_vec))
        return max(0.0, min(1.0, (similarity + 1.0) / 2.0))

    def generate(self, prompt: str, max_tokens: int = 256) -> str:
        """Generate text using PyTorch model (for crawl analysis)"""
        self.ensure_torch_loaded()
        if not self.torch_loaded:
            return ""

        import torch
        old_threads = torch.get_num_threads()
        torch.set_num_threads(self.config.get_active_threads())

        try:
            inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
            with torch.no_grad():
                outputs = self.torch_model.generate(
                    **inputs,
                    max_new_tokens=max_tokens,
                    do_sample=True,
                    temperature=0.3,
                    top_p=0.9,
                    repetition_penalty=1.1
                )
            response = self.tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
            return response.strip()
        finally:
            torch.set_num_threads(old_threads)

    def unload(self):
        """Unload models to free RAM"""
        self.onnx_session = None
        self.torch_model = None
        self.tokenizer = None
        self.is_loaded = False
        self.torch_loaded = False
        self.is_loaded = False
        self.load_progress = 0.0
        import gc
        gc.collect()


# ============================================================================
# Database Interface — reads code_chunks, writes codebase_training
# ============================================================================

class DatabaseInterface:
    """PostgreSQL interface for reading chunks and writing training data"""

    def __init__(self):
        self.conn = None
        self.project_path = os.environ.get('SPECMEM_PROJECT_PATH', os.getcwd())
        self._connect()

    def _connect(self):
        try:
            import psycopg2
            # Use same connection params as SpecMem
            db_url = os.environ.get('DATABASE_URL') or os.environ.get('SPECMEM_DATABASE_URL')
            if db_url:
                self.conn = psycopg2.connect(db_url)
            else:
                self.conn = psycopg2.connect(
                    host=os.environ.get('SPECMEM_DB_HOST', 'localhost'),
                    port=int(os.environ.get('SPECMEM_DB_PORT', '5432')),
                    database=os.environ.get('SPECMEM_DB_NAME', 'specmem'),
                    user=os.environ.get('SPECMEM_DB_USER', 'specmem'),
                    password=os.environ.get('SPECMEM_DB_PASSWORD', 'specmem')
                )
            self.conn.autocommit = True

            # Set search path to project schema
            # Priority: env var → schema file (written by container entrypoint) → auto-derive from path
            schema = os.environ.get('SPECMEM_SCHEMA') or os.environ.get('SPECMEM_DB_SCHEMA')
            if not schema:
                # Container mode: read schema from file written by entrypoint
                schema_file = os.path.join(
                    os.environ.get('SPECMEM_SOCKET_DIR', ''),
                    'project-schema'
                )
                if os.path.exists(schema_file):
                    try:
                        schema = open(schema_file).read().strip()
                    except:
                        pass
            if not schema:
                # Auto-derive from project path (matches Node.js projectNamespacing.js)
                project_path = os.environ.get('SPECMEM_PROJECT_PATH', '/')
                if project_path and project_path != '/':
                    import re
                    dirname = os.path.basename(project_path).lower()
                    dirname = re.sub(r'[^a-z0-9_]', '_', dirname)
                    dirname = re.sub(r'_+', '_', dirname).strip('_') or 'default'
                    schema = f'specmem_{dirname}'
                else:
                    schema = 'public'
            with self.conn.cursor() as cur:
                cur.execute(f"SET search_path TO {schema}, public")
            print(f"   Schema: {schema}", file=sys.stderr)

            print(f"✅ Database connected (project: {self.project_path})", file=sys.stderr)
        except Exception as e:
            print(f"⚠️ Database connection failed: {e} - crawl mode disabled", file=sys.stderr)
            self.conn = None

    def get_unanalyzed_chunks(self, limit: int = 10) -> List[Dict]:
        """Get code_chunks that don't have codebase_training entries yet"""
        if not self.conn:
            return []
        try:
            with self.conn.cursor() as cur:
                cur.execute("""
                    SELECT cc.id, cc.file_path, cc.content, cc.language,
                           cc.start_line, cc.end_line, cc.chunk_type
                    FROM code_chunks cc
                    LEFT JOIN codebase_training ct ON ct.chunk_id = cc.id
                        AND ct.project_path = %s
                    WHERE cc.project_path = %s
                      AND ct.id IS NULL
                      AND cc.content IS NOT NULL
                      AND LENGTH(cc.content) > 10
                    ORDER BY cc.file_path, cc.chunk_index
                    LIMIT %s
                """, (self.project_path, self.project_path, limit))
                cols = [desc[0] for desc in cur.description]
                return [dict(zip(cols, row)) for row in cur.fetchall()]
        except Exception as e:
            print(f"⚠️ Failed to get unanalyzed chunks: {e}", file=sys.stderr)
            return []

    def get_doc_context(self, language: str) -> str:
        """Get relevant doc content from codebase_files for cross-referencing"""
        if not self.conn:
            return ""
        try:
            with self.conn.cursor() as cur:
                # Get doc files that might be relevant (READMEs, docs)
                cur.execute("""
                    SELECT file_path, LEFT(content, 2000) as content
                    FROM codebase_files
                    WHERE project_path = %s
                      AND (file_path LIKE '%%.md' OR file_path LIKE '%%.txt' OR file_path LIKE '%%.rst')
                      AND content IS NOT NULL
                    ORDER BY
                        CASE WHEN file_path ILIKE '%%readme%%' THEN 0
                             WHEN file_path ILIKE '%%doc%%' THEN 1
                             ELSE 2
                        END
                    LIMIT 3
                """, (self.project_path,))
                docs = cur.fetchall()
                if docs:
                    return "\n---\n".join([f"[{d[0]}]\n{d[1]}" for d in docs])
            return ""
        except Exception as e:
            return ""

    def store_training_entry(self, entry: Dict):
        """Store a codebase_training entry"""
        if not self.conn:
            return
        try:
            with self.conn.cursor() as cur:
                cur.execute("""
                    INSERT INTO codebase_training
                        (file_path, chunk_id, analysis_type, severity, title,
                         description, confidence, line_start, line_end,
                         language, doc_source, project_path)
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                """, (
                    entry['file_path'], entry.get('chunk_id'),
                    entry['analysis_type'], entry.get('severity', 'info'),
                    entry['title'], entry.get('description', ''),
                    entry.get('confidence', 0.5),
                    entry.get('line_start'), entry.get('line_end'),
                    entry.get('language'), entry.get('doc_source'),
                    self.project_path
                ))
        except Exception as e:
            print(f"⚠️ Failed to store training entry: {e}", file=sys.stderr)

    def get_training_stats(self) -> Dict:
        """Get stats for status reporting"""
        if not self.conn:
            return {"files_analyzed": 0, "bugs_found": 0, "total_entries": 0}
        try:
            with self.conn.cursor() as cur:
                cur.execute("""
                    SELECT
                        COUNT(DISTINCT file_path) as files_analyzed,
                        COUNT(*) FILTER (WHERE analysis_type = 'bug') as bugs_found,
                        COUNT(*) FILTER (WHERE analysis_type = 'security') as security_issues,
                        COUNT(*) as total_entries
                    FROM codebase_training
                    WHERE project_path = %s
                """, (self.project_path,))
                row = cur.fetchone()
                return {
                    "files_analyzed": row[0] or 0,
                    "bugs_found": row[1] or 0,
                    "security_issues": row[2] or 0,
                    "total_entries": row[3] or 0
                }
        except:
            return {"files_analyzed": 0, "bugs_found": 0, "total_entries": 0}

    def prune_if_needed(self):
        """Prune training entries if exceeding 10x codebase_files count"""
        if not self.conn:
            return
        try:
            with self.conn.cursor() as cur:
                cur.execute("""
                    SELECT COUNT(*) FROM codebase_files WHERE project_path = %s
                """, (self.project_path,))
                file_count = cur.fetchone()[0] or 0
                max_entries = file_count * 10

                cur.execute("""
                    SELECT COUNT(*) FROM codebase_training WHERE project_path = %s
                """, (self.project_path,))
                training_count = cur.fetchone()[0] or 0

                if training_count > max_entries and max_entries > 0:
                    excess = training_count - max_entries
                    cur.execute("""
                        DELETE FROM codebase_training
                        WHERE id IN (
                            SELECT id FROM codebase_training
                            WHERE project_path = %s
                            ORDER BY created_at ASC
                            LIMIT %s
                        )
                    """, (self.project_path, excess))
                    print(f"🧹 Pruned {excess} oldest training entries (limit: {max_entries})", file=sys.stderr)
        except Exception as e:
            print(f"⚠️ Prune failed: {e}", file=sys.stderr)


# ============================================================================
# Crawl Engine — background analysis loop
# ============================================================================

class CrawlEngine:
    """
    Background crawl that analyzes code chunks using Pythia.
    Runs at 5% CPU with QQMS-coordinated pacing.
    """
    def __init__(self, model: ModelManager, db: DatabaseInterface,
                 qqms: QQMSCoordinator, resource_config: MiniCOTResourceConfig):
        self.model = model
        self.db = db
        self.qqms = qqms
        self.config = resource_config
        self.running = False
        self.thread = None
        self.chunks_analyzed = 0

    def start(self):
        if self.running:
            return
        self.running = True
        self.thread = threading.Thread(target=self._crawl_loop, daemon=True, name="minicot-crawl")
        self.thread.start()
        print("🐛 Crawl engine started", file=sys.stderr)

    def stop(self):
        self.running = False
        if self.thread:
            self.thread.join(timeout=5)

    def _crawl_loop(self):
        """Main crawl loop — processes unanalyzed chunks with pacing"""
        # PRIME: Pre-load model after short delay so scoring requests are instant
        # Without this, model only loads on first score request (slow first call)
        time.sleep(15)  # Let other services start first
        if self.running and not self.model.is_loaded:
            available = get_available_ram_mb()
            if available > self.config.crawl_ram_cap_mb * 1.5:
                print("🔥 PRIMING: Pre-loading model for instant scoring...", file=sys.stderr)
                try:
                    self.model.ensure_loaded()
                    print("✅ Model primed and ready for scoring", file=sys.stderr)
                except Exception as e:
                    print(f"⚠️ Model priming failed: {e} (will lazy-load on first request)", file=sys.stderr)

        idle_cycles = 0

        while self.running:
            try:
                # Check QQMS coordination — pause if embedding server is busy
                if self.qqms.is_paused or self.qqms.should_pause():
                    time.sleep(self.config.crawl_sleep_max)
                    continue

                # Get batch of unanalyzed chunks
                chunks = self.db.get_unanalyzed_chunks(limit=5)
                if not chunks:
                    idle_cycles += 1
                    # Longer sleep when idle
                    time.sleep(min(30, self.config.crawl_sleep_base * idle_cycles))
                    if idle_cycles > 10:
                        idle_cycles = 10  # Cap idle backoff
                    continue

                idle_cycles = 0

                # Ensure model is loaded (lazy)
                if not self.model.is_loaded:
                    # Check RAM before loading
                    available = get_available_ram_mb()
                    if available < self.config.crawl_ram_cap_mb * 1.5:
                        print(f"⏳ Waiting for RAM ({available:.0f}MB available, need ~{self.config.crawl_ram_cap_mb:.0f}MB)", file=sys.stderr)
                        time.sleep(10)
                        continue
                    self.model.ensure_loaded()

                for chunk in chunks:
                    if not self.running or self.qqms.is_paused:
                        break

                    self._analyze_chunk(chunk)
                    self.chunks_analyzed += 1

                    # QQMS pacing — dynamic sleep based on CPU load
                    delay = self.qqms.get_crawl_delay()
                    time.sleep(delay)

                # Periodic pruning
                if self.chunks_analyzed % 50 == 0 and self.chunks_analyzed > 0:
                    self.db.prune_if_needed()

            except Exception as e:
                print(f"⚠️ Crawl error: {e}", file=sys.stderr)
                time.sleep(self.config.crawl_sleep_max)

    def _analyze_chunk(self, chunk: Dict):
        """Analyze a single code chunk with Pythia"""
        try:
            content = chunk['content'][:2000]  # Truncate for model context
            language = chunk.get('language', 'unknown')

            # Get doc context for cross-referencing
            doc_context = self.db.get_doc_context(language)

            prompt = f"""<|system|>
You are a code analyzer. Analyze the following code and identify any bugs, patterns, quality issues, or security concerns. Be concise.
{f'Reference docs: {doc_context[:500]}' if doc_context else ''}
</s>
<|user|>
Language: {language}
File: {chunk['file_path']} (lines {chunk.get('start_line', '?')}-{chunk.get('end_line', '?')})

```
{content}
```

List issues found (type: bug/pattern/quality/security/suggestion). For each: title, severity (critical/warning/info/suggestion), confidence (0-1), description.
If no issues, say "No issues found."
</s>
<|assistant|>"""

            response = self.model.generate(prompt, max_tokens=200)

            # Parse response into training entries
            entries = self._parse_analysis(response, chunk)
            for entry in entries:
                self.db.store_training_entry(entry)

        except Exception as e:
            print(f"⚠️ Chunk analysis failed ({chunk['file_path']}): {e}", file=sys.stderr)

    def _parse_analysis(self, response: str, chunk: Dict) -> List[Dict]:
        """Parse Pythia response into structured training entries"""
        entries = []

        if "no issues" in response.lower():
            return entries

        # Simple heuristic parsing — Pythia output is unpredictable
        # Look for common patterns
        analysis_types = {
            'bug': 'bug', 'error': 'bug', 'fix': 'bug',
            'pattern': 'pattern', 'antipattern': 'pattern',
            'quality': 'quality', 'style': 'quality', 'readability': 'quality',
            'security': 'security', 'vulnerability': 'security', 'injection': 'security',
            'suggestion': 'suggestion', 'improvement': 'suggestion', 'recommend': 'suggestion'
        }

        severities = {
            'critical': 'critical', 'high': 'critical',
            'warning': 'warning', 'medium': 'warning',
            'info': 'info', 'low': 'info',
            'suggestion': 'suggestion'
        }

        # Split by numbered items or lines
        lines = [l.strip() for l in response.split('\n') if l.strip()]

        current_type = 'suggestion'
        current_severity = 'info'

        for line in lines[:5]:  # Max 5 entries per chunk
            # Detect analysis type
            line_lower = line.lower()
            for keyword, atype in analysis_types.items():
                if keyword in line_lower:
                    current_type = atype
                    break

            for keyword, sev in severities.items():
                if keyword in line_lower:
                    current_severity = sev
                    break

            if len(line) > 10:  # Skip very short lines
                entries.append({
                    'file_path': chunk['file_path'],
                    'chunk_id': chunk['id'],
                    'analysis_type': current_type,
                    'severity': current_severity,
                    'title': line[:200],
                    'description': line,
                    'confidence': 0.5,  # Pythia = moderate confidence
                    'line_start': chunk.get('start_line'),
                    'line_end': chunk.get('end_line'),
                    'language': chunk.get('language'),
                    'doc_source': None
                })

        return entries


# ============================================================================
# Active Scoring — COT scoring for find_code_pointers
# ============================================================================

class ActiveScorer:
    """Handles real-time scoring requests from find_code_pointers"""

    def __init__(self, model: ModelManager, resource_config: MiniCOTResourceConfig):
        self.model = model
        self.config = resource_config

    def score(self, code: str, query: str) -> Dict:
        """Score code relevance via ONNX hidden-state cosine similarity (~50-100ms)"""
        try:
            score = self.model.score_similarity(query, code)

            return {
                "score": round(score, 3),
                "reasoning": f"Cosine similarity: {score:.3f}",
                "model": self.model.model_name
            }
        except Exception as e:
            return {
                "score": 0.5,
                "reasoning": f"Scoring failed: {str(e)}",
                "model": self.model.model_name
            }


# ============================================================================
# Unix Socket Server
# ============================================================================

class MiniCOTServer:
    """
    Unix socket server for Mini COT service.
    Handles health, score, status, pause/resume requests.
    """
    def __init__(self, socket_path: str, model: ModelManager, db: DatabaseInterface,
                 crawl: CrawlEngine, scorer: ActiveScorer, qqms: QQMSCoordinator,
                 low_resource_config: LowResourceConfig = None):
        self.socket_path = socket_path
        self.model = model
        self.db = db
        self.crawl = crawl
        self.scorer = scorer
        self.qqms = qqms
        self.server = None
        self.running = False
        self.last_request_time = time.time()
        self.low_resource_config = low_resource_config or LowResourceConfig()
        self._idle_monitor_thread = None

    def _start_idle_monitor(self):
        """Monitor idle time and unload models when not in use (low-resource mode)."""
        idle_seconds = self.low_resource_config.idle_unload_seconds
        if idle_seconds <= 0:
            return  # Disabled

        def _monitor():
            while self.running:
                time.sleep(30)  # Check every 30s
                idle_time = time.time() - self.last_request_time
                if idle_time > idle_seconds and self.model.is_loaded:
                    print(f"💤 Idle for {idle_time:.0f}s — unloading ONNX scorer to save RAM", file=sys.stderr)
                    self.model.unload()
                    import gc
                    gc.collect()
                    print(f"✅ Model unloaded — will reload on next score request", file=sys.stderr)

        self._idle_monitor_thread = threading.Thread(target=_monitor, daemon=True, name="minicot-idle")
        self._idle_monitor_thread.start()
        print(f"🔧 Idle monitor: unload after {idle_seconds}s idle", file=sys.stderr)

    def start(self):
        """Start the Unix socket server"""
        # Clean up stale socket
        if os.path.exists(self.socket_path):
            os.unlink(self.socket_path)

        # Ensure socket directory exists
        os.makedirs(os.path.dirname(self.socket_path), exist_ok=True)

        self.server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.server.bind(self.socket_path)
        self.server.listen(5)
        self.server.settimeout(1.0)  # Allow graceful shutdown
        self.running = True

        # Make socket accessible
        os.chmod(self.socket_path, 0o777)

        # Start idle monitor (low-resource mode only)
        self._start_idle_monitor()

        print(f"🔌 Mini COT server listening on {self.socket_path}", file=sys.stderr)

        while self.running:
            try:
                conn, _ = self.server.accept()
                # Handle each connection in a thread (bounded)
                threading.Thread(
                    target=self._handle_connection,
                    args=(conn,),
                    daemon=True,
                    name="minicot-conn"
                ).start()
            except socket.timeout:
                continue
            except Exception as e:
                if self.running:
                    print(f"⚠️ Accept error: {e}", file=sys.stderr)

    def stop(self):
        """Stop the server"""
        self.running = False
        if self.server:
            self.server.close()
        if os.path.exists(self.socket_path):
            os.unlink(self.socket_path)

    def _handle_connection(self, conn: socket.socket):
        """Handle a single client connection"""
        try:
            conn.settimeout(30.0)
            data = b""
            while True:
                chunk = conn.recv(4096)
                if not chunk:
                    break
                data += chunk
                # Try to parse JSON
                try:
                    request = json.loads(data.decode('utf-8'))
                    break
                except json.JSONDecodeError:
                    if len(data) > 65536:
                        break
                    continue

            if not data:
                return

            request = json.loads(data.decode('utf-8'))
            response = self._handle_request(request)
            conn.sendall(json.dumps(response).encode('utf-8'))

        except Exception as e:
            try:
                error_resp = {"error": str(e)}
                conn.sendall(json.dumps(error_resp).encode('utf-8'))
            except:
                pass
        finally:
            conn.close()

    def _handle_request(self, request: Dict) -> Dict:
        """Route request to appropriate handler"""
        req_type = request.get("type", "")

        if req_type == "health":
            mode = "paused" if self.qqms.is_paused else ("loading" if self.model.is_loading else "crawl")
            return {
                "status": "ok",
                "mode": mode,
                "onnx_scorer": self.model.is_loaded,
                "torch_analyzer": self.model.torch_loaded,
                "model_loaded": self.model.is_loaded,  # compat
                "load_progress": self.model.load_progress,
                "chunks_analyzed": self.crawl.chunks_analyzed,
                "ram_available_mb": round(get_available_ram_mb()),
                "cpu_percent": round(get_cpu_percent(), 1)
            }

        elif req_type == "score":
            self.last_request_time = time.time()  # Keep-alive for idle monitor
            code = request.get("code", "")
            query = request.get("query", "")
            if not code or not query:
                return {"error": "Missing 'code' or 'query'"}
            return self.scorer.score(code, query)

        elif req_type == "status":
            stats = self.db.get_training_stats()
            stats["crawl_running"] = self.crawl.running
            stats["model_loaded"] = self.model.is_loaded
            stats["paused"] = self.qqms.is_paused
            stats["pause_reason"] = self.qqms.pause_reason
            stats["ram_available_mb"] = round(get_available_ram_mb())
            return stats

        elif req_type == "pause":
            self.qqms.pause()
            return {"status": "paused"}

        elif req_type == "resume":
            self.qqms.resume()
            return {"status": "resumed"}

        else:
            return {"error": f"Unknown request type: {req_type}"}


# ============================================================================
# Main Entry Point
# ============================================================================

def main():
    parser = argparse.ArgumentParser(description="Mini COT Service for SpecMem")
    parser.add_argument('--socket', required=True, help='Unix socket path')
    parser.add_argument('--model', default='EleutherAI/pythia-410m', help='Model name')
    parser.add_argument('--device', default='cpu', help='Device (cpu/cuda)')
    args = parser.parse_args()

    print(f"🚀 Mini COT Service starting...", file=sys.stderr)
    print(f"   Socket: {args.socket}", file=sys.stderr)
    print(f"   Model: {args.model}", file=sys.stderr)
    print(f"   Device: {args.device}", file=sys.stderr)

    # Initialize resource config
    resource_config = MiniCOTResourceConfig()
    low_resource_config = LowResourceConfig()
    print(f"   CPU budget: {resource_config.crawl_cpu_percent}% crawl / {resource_config.active_cpu_percent}% active", file=sys.stderr)
    print(f"   RAM cap: {resource_config.crawl_ram_cap_mb}MB crawl / {resource_config.active_ram_cap_mb}MB active", file=sys.stderr)
    print(f"   CPU cores: {resource_config.cpu_cores_min}-{resource_config.cpu_cores_max}", file=sys.stderr)
    print(f"   Threads: {resource_config.crawl_threads} crawl / {resource_config.active_threads} active", file=sys.stderr)
    if low_resource_config.idle_unload_seconds > 0:
        print(f"   Idle unload: {low_resource_config.idle_unload_seconds}s", file=sys.stderr)

    # Initialize components
    qqms = QQMSCoordinator(resource_config)
    model = ModelManager(args.model, args.device, resource_config)
    db = DatabaseInterface()
    crawl = CrawlEngine(model, db, qqms, resource_config)
    scorer = ActiveScorer(model, resource_config)
    server = MiniCOTServer(args.socket, model, db, crawl, scorer, qqms, low_resource_config)

    # Signal handlers for clean shutdown
    def shutdown(signum, frame):
        print(f"\n🛑 Shutting down Mini COT (signal {signum})...", file=sys.stderr)
        server.running = False
        crawl.stop()
        server.stop()
        model.unload()
        sys.exit(0)

    signal.signal(signal.SIGTERM, shutdown)
    signal.signal(signal.SIGINT, shutdown)

    # Start crawl engine (background thread)
    crawl.start()

    # Start socket server (blocks)
    try:
        server.start()
    except KeyboardInterrupt:
        shutdown(signal.SIGINT, None)
    except Exception as e:
        print(f"❌ Server error: {e}", file=sys.stderr)
        traceback.print_exc(file=sys.stderr)
        shutdown(signal.SIGTERM, None)


if __name__ == '__main__':
    main()