model_manager.py

"""
model_manager.py — نظام إدارة النماذج
══════════════════════════════════════════════════════

القرارات المبنية على debugging حقيقي (محدَّث يونيو 2026):
  • Qwen3.5-9B (NEW!) — الأفضل في فئة <10B، 262K context، 201 لغة، عربي أقوى
  • Qwen3.5 MoE (35B-A3B+) ← مرفوض: KV cache reuse bug معروف
  • Qwen3-8B ← fallback موثوق إذا فشل تحميل Qwen3.5-9B
  • Q4_K_M فقط — Q8_0 عنده garbage output bug في llama.cpp
  • لا speculative decoding — net-negative على CPU Q4
  • Thinking mode: OFF by default في Qwen3.5 (على عكس Qwen3) — أحسن لنظامنا
"""

from __future__ import annotations
import asyncio, json, logging, os, time
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any

logger = logging.getLogger("mythical.models")

MODEL_DIR    = Path(os.getenv("MODEL_DIR",    "/data/models"))
CONFIG_PATH  = Path(os.getenv("MODEL_CONFIG", "/data/model_config.json"))
RAM_LIMIT_GB = float(os.getenv("RAM_LIMIT_GB","14.0"))


@dataclass
class ModelDef:
    id:              str
    name:            str
    repo:            str
    glob_primary:    str
    glob_fallback:   str
    size_gb:         float      # حجم الأوزان تقريباً
    ctx_size:        int        # أقصى context آمن
    quality:         int        # جودة نسبية 0-100
    has_vision:      bool
    has_thinking:    bool
    mmproj_repo:     str = ""
    mmproj_glob:     str = "*mmproj*"
    extra_flags:     str = ""
    temp:     float = 0.6
    top_k:    int   = 20
    top_p:    float = 0.95
    min_p:    float = 0.0
    notes:    str   = ""


# ─────────────────────────────────────────────────────────────────────────────
# CATALOG — محدَّث يونيو 2026، مُختبَر على CPU Basic HF Spaces
# ─────────────────────────────────────────────────────────────────────────────
CATALOG: list[ModelDef] = [

    # ═══════════════════════════════════════════════════════════════════════
    # TIER 1 — CHAMPION: Qwen3.5-9B Q4_K_M  ← الجديد يونيو 2026
    # • أحسن نموذج <10B في يونيو 2026 بإجماع الـbenchmarks
    # • 262K context window (vs 16K لـQwen3-8B)
    # • 201 لغة — عربي أقوى بكتير (MMLU-ProX + WMT24++ على 55 لغة)
    # • Thinking mode متاح لكن OFF by default (ممتاز لنظامنا!)
    # • ⚠️ تحذير: Qwen3.5 MoE (35B-A3B+) عنده KV cache bug — الـ9B dense تمام
    # ═══════════════════════════════════════════════════════════════════════
    ModelDef(
        id           = "qwen35-9b",
        name         = "Qwen3.5 9B Q4_K_M (unsloth)",
        repo         = "unsloth/Qwen3.5-9B-GGUF",
        glob_primary = "Qwen3.5-9B-Q4_K_M.gguf",
        glob_fallback= "Qwen3.5-9B-UD-Q4_K_XL.gguf",
        size_gb      = 5.6,
        ctx_size     = 32768,   # نستخدم 32K من أصل 262K — آمن على 16GB RAM
        quality      = 95,
        has_vision   = False,   # vision ممكنة مستقبلاً بـmmproj منفصل
        has_thinking = True,
        mmproj_repo  = "",
        mmproj_glob  = "",
        # presence_penalty=1.5 موصى بيها من unsloth لتقليل التكرار في Qwen3.5
        extra_flags  = "--jinja --reasoning-format deepseek --presence-penalty 1.5",
        temp=0.7, top_k=20, top_p=0.8, min_p=0.0,
        notes="CHAMPION Jun 2026. 262K ctx, 201 langs, thinking OFF by default. ~5.6GB.",
    ),

    # ═══════════════════════════════════════════════════════════════════════
    # TIER 2 — RELIABLE: Qwen3-8B Q4_K_M  ← الافتراضي السابق، مستقر 100%
    # • مُختبَر في production، مفيش bugs معروفة
    # • fallback آمن لو Qwen3.5-9B اتأخر تحميله
    # ═══════════════════════════════════════════════════════════════════════
    ModelDef(
        id           = "qwen3-8b",
        name         = "Qwen3 8B Q4_K_M (unsloth)",
        repo         = "unsloth/Qwen3-8B-GGUF",
        glob_primary = "Qwen3-8B-Q4_K_M.gguf",
        glob_fallback= "Qwen3-8B-Q4_K_S.gguf",
        size_gb      = 5.1,
        ctx_size     = 16384,
        quality      = 88,
        has_vision   = False,
        has_thinking = True,
        mmproj_repo  = "",
        mmproj_glob  = "",
        extra_flags  = "--jinja --reasoning-format deepseek",
        temp=0.6, top_k=20, top_p=0.95, min_p=0.0,
        notes="Reliable fallback. Production-tested. 5.1GB. Thinking ON by default.",
    ),

    # ═══════════════════════════════════════════════════════════════════════
    # TIER 3 — FAST: Qwen3-4B Q4_K_M  ← للمهام البسيطة السريعة
    # ═══════════════════════════════════════════════════════════════════════
    ModelDef(
        id           = "qwen3-4b",
        name         = "Qwen3 4B Q4_K_M (unsloth)",
        repo         = "unsloth/Qwen3-4B-GGUF",
        glob_primary = "Qwen3-4B-Q4_K_M.gguf",
        glob_fallback= "Qwen3-4B-Q4_K_S.gguf",
        size_gb      = 2.5,
        ctx_size     = 16384,
        quality      = 76,
        has_vision   = False,
        has_thinking = True,
        mmproj_repo  = "",
        mmproj_glob  = "",
        extra_flags  = "--jinja --reasoning-format deepseek",
        temp=0.6, top_k=20, top_p=0.95, min_p=0.0,
        notes="Emergency fallback. 2.5GB. Fast on CPU. Simple tasks only.",
    ),

    # ═══════════════════════════════════════════════════════════════════════
    # TIER 4 — HEAVY: Qwen3-14B  ← للأجهزة اللي عندها >12GB متاحة
    # ملاحظة: Qwen3.5-9B (quality=95) أحسن منه (94) مع نص حجمه!
    # ═══════════════════════════════════════════════════════════════════════
    ModelDef(
        id           = "qwen3-14b",
        name         = "Qwen3 14B Q4_K_M (unsloth)",
        repo         = "unsloth/Qwen3-14B-GGUF",
        glob_primary = "Qwen3-14B-Q4_K_M.gguf",
        glob_fallback= "*Q4_K_M*",
        size_gb      = 9.5,
        ctx_size     = 12288,
        quality      = 94,
        has_vision   = False,
        has_thinking = True,
        mmproj_repo  = "",
        mmproj_glob  = "",
        extra_flags  = "--jinja --reasoning-format deepseek",
        temp=0.6, top_k=20, top_p=0.95, min_p=0.0,
        notes="Heavy. 9.5GB. Slower than Qwen3.5-9B but higher param count.",
    ),
]

CATALOG_BY_ID = {m.id: m for m in CATALOG}


@dataclass
class ModelConfig:
    active_id:    str
    active_path:  str
    mmproj_path:  str
    pending_id:   str = ""
    pending_path: str = ""
    last_check:   float = 0.0
    version:      int = 2

    def save(self):
        CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
        CONFIG_PATH.write_text(json.dumps(asdict(self), indent=2))

    @classmethod
    def load(cls) -> "ModelConfig":
        if not CONFIG_PATH.exists():
            cfg = cls(active_id="qwen35-9b", active_path="", mmproj_path="",
                      last_check=time.time())   # prevents immediate upgrade check on fresh install
            cfg.save()
            return cfg
        try:
            d = json.loads(CONFIG_PATH.read_text())
            # Migration: تنظيف حقول قديمة
            valid = {k for k in cls.__dataclass_fields__}
            return cls(**{k: v for k, v in d.items() if k in valid})
        except Exception as e:
            logger.warning(f"[models] Config load error ({e}), using default.")
            return cls(active_id="qwen35-9b", active_path="", mmproj_path="")


class ModelManager:

    UPDATE_INTERVAL_H = float(os.getenv("UPDATE_CHECK_HOURS", "168"))

    def __init__(self):
        self.cfg = ModelConfig.load()
        self._downloading = False
        self._bg_task: asyncio.Task | None = None

    # ── Selection ────────────────────────────────────────────────────────
    def select_best(self, ram_gb: float | None = None) -> ModelDef:
        limit = ram_gb or RAM_LIMIT_GB
        candidates = sorted(
            [m for m in CATALOG if m.size_gb <= limit],
            key=lambda m: m.quality,
            reverse=True,
        )
        if not candidates:
            return min(CATALOG, key=lambda m: m.size_gb)
        return candidates[0]

    # ── Discovery ────────────────────────────────────────────────────────
    def find_local(self, m: ModelDef) -> tuple[str, str]:
        """يرجع (model_path, mmproj_path) من الديسك."""
        def _find(glob: str, min_mb: int) -> str:
            for p in sorted(MODEL_DIR.glob(glob),
                            key=lambda x: x.stat().st_size, reverse=True):
                if p.stat().st_size > min_mb * 1024 * 1024:
                    return str(p)
            return ""

        model = _find(m.glob_primary, 500) or _find(m.glob_fallback, 500)
        mmproj = _find(m.mmproj_glob, 50) if m.has_vision else ""
        return model, mmproj

    # ── Download ─────────────────────────────────────────────────────────
    async def download(self, m: ModelDef, force=False) -> tuple[str, str]:
        if not force:
            existing, existing_mm = self.find_local(m)
            if existing:
                return existing, existing_mm

        self._downloading = True
        logger.info(f"[models] Downloading {m.name} from {m.repo}...")

        loop = asyncio.get_running_loop()

        def _dl(repo, glob):
            import subprocess
            # Build env — HF_TOKEN optional for public repos
            env = {**os.environ, "HF_HUB_ENABLE_HF_TRANSFER": "1"}
            token = os.getenv("HF_TOKEN", "").strip()
            if token:
                env["HUGGING_FACE_HUB_TOKEN"] = token
                env["HF_TOKEN"] = token
            for g in [glob, m.glob_fallback if glob == m.glob_primary else ""]:
                if not g:
                    continue
                cmd = ["huggingface-cli", "download", repo,
                       "--include", g, "--local-dir", str(MODEL_DIR)]
                if not token:
                    # Public access — verbose to see any auth errors
                    cmd.append("--quiet")
                else:
                    cmd.append("--quiet")
                logger.info(f"[models] Running: {' '.join(cmd[:5])}...")
                r = subprocess.run(cmd, env=env, capture_output=True, timeout=600)
                if r.returncode == 0:
                    logger.info(f"[models] Download succeeded for glob: {g}")
                    return ""
                else:
                    err = r.stderr.decode("utf-8", errors="replace")[:300]
                    logger.warning(f"[models] Download failed (glob={g}): {err}")
            return ""

        TIMEOUT = float(os.getenv("DOWNLOAD_TIMEOUT", "30.0")) * 20  # max 10 min
        try:
            # [CRIT-09 fix] timeout prevents infinite hang on HF outage
            await asyncio.wait_for(
                loop.run_in_executor(None, _dl, m.repo, m.glob_primary),
                timeout=TIMEOUT)
            if m.has_vision and m.mmproj_repo:
                await asyncio.wait_for(
                    loop.run_in_executor(None, _dl, m.mmproj_repo, m.mmproj_glob),
                    timeout=TIMEOUT)
            model, mmproj = self.find_local(m)
            logger.info(f"[models] Done: {model}")
            return model, mmproj
        except asyncio.TimeoutError:
            logger.error(f"[models] Download timed out after {TIMEOUT:.0f}s")
            return "", ""
        finally:
            self._downloading = False

    # ── Apply pending upgrade ────────────────────────────────────────────
    def apply_pending(self) -> bool:
        if not self.cfg.pending_id or not self.cfg.pending_path:
            return False
        if not Path(self.cfg.pending_path).exists():
            self.cfg.pending_id = self.cfg.pending_path = ""
            self.cfg.save()
            return False
        old = self.cfg.active_id
        self.cfg.active_id   = self.cfg.pending_id
        self.cfg.active_path = self.cfg.pending_path
        self.cfg.pending_id  = self.cfg.pending_path = ""
        self.cfg.save()
        logger.info(f"[models] Hot-swap: {old} → {self.cfg.active_id}")
        return True

    # ── Auto-update background loop ──────────────────────────────────────
    def start_updater(self):
        self._bg_task = asyncio.create_task(self._update_loop(), name="model_updater")

    async def _update_loop(self):
        await asyncio.sleep(7200)   # wait 2h after startup before first upgrade check
        while True:
            try:
                if (time.time() - self.cfg.last_check) / 3600 >= self.UPDATE_INTERVAL_H:
                    await self._smart_upgrade()
            except Exception as e:
                logger.error(f"[models] Update loop error: {e}")
            await asyncio.sleep(3600)

    async def _smart_upgrade(self):
        current = CATALOG_BY_ID.get(self.cfg.active_id)
        if not current:
            return
        better = [
            m for m in CATALOG
            if m.quality > current.quality
            and m.size_gb <= RAM_LIMIT_GB
            and m.id != current.id
        ]
        if not better:
            self.cfg.last_check = time.time()
            self.cfg.save()
            logger.info(f"[models] Already on best model for this hardware.")
            return
        best = max(better, key=lambda m: m.quality)
        logger.info(f"[models] Better model found: {best.name}. Downloading in background...")

        async def _bg():
            path, mmproj = await self.download(best)
            if path:
                self.cfg.pending_id   = best.id
                self.cfg.pending_path = path
                self.cfg.last_check   = time.time()
                self.cfg.save()
                logger.info(f"[models] Upgrade ready: {best.name}. Activates on restart.")
        asyncio.create_task(_bg())

    def stop(self):
        if self._bg_task:
            self._bg_task.cancel()

    # ── Build llama-server command ───────────────────────────────────────
    def build_cmd(self, host="127.0.0.1", port="8080",
                  ctx_size: str | None = None) -> list[str]:
        m = CATALOG_BY_ID.get(self.cfg.active_id)
        if not m:
            raise ValueError(f"Unknown model id: {self.cfg.active_id}")

        ctx = ctx_size or str(m.ctx_size)

        cmd = [
            "llama-server",
            "--model",      self.cfg.active_path,
            "--host",       host,
            "--port",       port,
            "--ctx-size",   ctx,
            "-t",           "2",
            "-tb",          "2",
            "-np",          "2",
            "--cont-batching",
            "--poll",       "100",
            "--prio",       "3",
            "--prio-batch", "2",
            "--cpu-range",  "0-1",
            "--cpu-strict", "1",
            "-b",           "512",
            "-ub",          "128",
            "-ctk",         "q4_0",
            "-ctv",         "q4_0",
            "-fa",          "on",
            "--slot-save-path", "/data/slot_cache",
            "--mlock",
            "--no-mmap",
            "--timeout",    "600",   # must be > THINKING_TIMEOUT (300s)
            "--log-disable",
            # NO speculative decoding — net-negative on CPU Q4 quantized models
    # Qwen3 sampling params (official recommended)
            "--temp",       str(m.temp),
            "--top-k",      str(m.top_k),
            "--top-p",      str(m.top_p),
            "--min-p",      str(m.min_p),
        ]

        # Vision
        if self.cfg.mmproj_path and m.has_vision:
            cmd += ["--mmproj", self.cfg.mmproj_path, "--no-mmproj-offload"]

        # Model-specific flags (jinja, reasoning-format, etc.)
        if m.extra_flags:
            import shlex
            cmd += shlex.split(m.extra_flags)

        return cmd

    # ── Status & Catalog ─────────────────────────────────────────────────
    def status(self) -> dict:
        m = CATALOG_BY_ID.get(self.cfg.active_id, {})
        pm = CATALOG_BY_ID.get(self.cfg.pending_id, {}) if self.cfg.pending_id else {}
        return {
            "active": {
                "id":        self.cfg.active_id,
                "name":      m.name if m else "unknown",
                "path":      self.cfg.active_path,
                "mmproj":    self.cfg.mmproj_path,
                "quality":   m.quality if m else 0,
                "thinking":  m.has_thinking if m else False,
                "vision":    m.has_vision if m else False,
                "ctx_size":  m.ctx_size if m else 0,
                "size_gb":   m.size_gb if m else 0,
            },
            "pending":       {"id": self.cfg.pending_id, "name": pm.name if pm else None} if self.cfg.pending_id else None,
            "is_downloading":self._downloading,
            "last_check_h":  round((time.time() - self.cfg.last_check)/3600, 1),
            "update_every_h":self.UPDATE_INTERVAL_H,
        }

    def catalog(self) -> list[dict]:
        result = []
        for m in CATALOG:
            local, _ = self.find_local(m)
            result.append({
                "id":        m.id,
                "name":      m.name,
                "quality":   m.quality,
                "size_gb":   m.size_gb,
                "thinking":  m.has_thinking,
                "vision":    m.has_vision,
                "on_disk":   bool(local),
                "active":    m.id == self.cfg.active_id,
                "pending":   m.id == self.cfg.pending_id,
                "notes":     m.notes,
            })
        return result


# ─────────────────────────────────────────────────────────────────────────────
# CLI — يستخدمه startup.sh
# ─────────────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import sys
    logging.basicConfig(level=logging.INFO)
    mgr = ModelManager()

    cmd = sys.argv[1] if len(sys.argv) > 1 else "status"

    if cmd == "status":
        print(json.dumps(mgr.status(), indent=2))

    elif cmd == "apply-pending":
        print("applied" if mgr.apply_pending() else "no-pending")

    elif cmd == "get-cmd":
        try:
            parts = mgr.build_cmd()
            print(" ".join(parts))
        except ValueError as e:
            print(f"ERROR: {e}", file=sys.stderr); sys.exit(1)

    elif cmd == "ensure-downloaded":
        async def _run():
            token = os.getenv("HF_TOKEN", "").strip()
            print(f"[model_manager] HF_TOKEN: {'SET (len=' + str(len(token)) + ')' if token else 'NOT SET — using anonymous access'}", file=sys.stderr)
            mgr.apply_pending()
            m = CATALOG_BY_ID.get(mgr.cfg.active_id) or mgr.select_best()
            mgr.cfg.active_id = m.id
            print(f"[model_manager] Downloading: {m.name} ({m.size_gb}GB) from {m.repo}", file=sys.stderr)
            path, mmproj = await mgr.download(m)
            if not path:
                print(f"[model_manager] Primary download failed, trying smallest fallback...", file=sys.stderr)
                fb = min(CATALOG, key=lambda x: x.size_gb)
                print(f"[model_manager] Fallback: {fb.name}", file=sys.stderr)
                path, mmproj = await mgr.download(fb)
                if path:
                    mgr.cfg.active_id = fb.id
            if not path:
                print(f"[model_manager] ALL downloads failed!", file=sys.stderr)
            else:
                print(f"[model_manager] SUCCESS: {path}", file=sys.stderr)
            mgr.cfg.active_path = path
            mgr.cfg.mmproj_path = mmproj
            mgr.cfg.save()
            print(json.dumps({
                "model_id":   mgr.cfg.active_id,
                "model_path": mgr.cfg.active_path,
                "mmproj_path":mgr.cfg.mmproj_path,
                "ok":         bool(path),
            }))
        asyncio.run(_run())

    else:
        print(f"Unknown: {cmd}", file=sys.stderr); sys.exit(1)