"""Ensure the oMLX server is up for a run.

oMLX runs as a managed background service (auto-starts on login) and loads models
on-demand with LRU eviction, so there's usually nothing to do but health-check. For the
nightly launchd job we still `omlx start` defensively and warm the active model so the
first agent turn isn't paying cold-load latency.
"""

from __future__ import annotations

import os
import shutil
import subprocess
import time

import httpx
from loguru import logger

from config.settings import get_settings

_OMLX_CANDIDATES = ["omlx", os.path.expanduser("~/.omlx/bin/omlx")]


def _omlx_bin() -> str | None:
    for cand in _OMLX_CANDIDATES:
        path = shutil.which(cand) if os.path.sep not in cand else (cand if os.path.exists(cand) else None)
        if path:
            return path
    return None


def server_up() -> bool:
    s = get_settings()
    try:
        r = httpx.get(
            f"{s.llm_base_url}/models",
            headers={"Authorization": f"Bearer {s.llm_api_key}"},
            timeout=5,
        )
        return r.status_code == 200
    except Exception:
        return False


def start_server() -> None:
    if server_up():
        logger.info("oMLX already up at {}", get_settings().llm_base_url)
        return
    bin_ = _omlx_bin()
    if not bin_:
        logger.warning("omlx CLI not found; cannot auto-start. Start the oMLX app manually.")
        return
    proc = subprocess.run([bin_, "start"], capture_output=True, text=True, timeout=120)
    logger.info("omlx start -> {} {}", proc.returncode, (proc.stdout + proc.stderr).strip()[:300])


def wait_healthy(timeout_s: int = 120) -> bool:
    deadline = time.time() + timeout_s
    while time.time() < deadline:
        if server_up():
            return True
        time.sleep(3)
    return False


def warm_model() -> None:
    """Fire a tiny completion so the active model is resident before the first agent turn."""
    from llm.client import build_llm

    try:
        build_llm().invoke("ok")
        logger.info("warmed model")
    except Exception as e:
        logger.warning("warm_model failed (non-fatal): {}", e)


def clear_kv_cache(hot: bool = True) -> None:
    """Wipe oMLX's SSD prefix/KV cache (and, by default, the in-memory hot cache).

    oMLX prefix-caches every prompt into an SSD-backed paged KV store so repeated prefixes
    skip prefill. Within a single agent that's a big win, but across personas the prefixes
    diverge, so a finished agent's blocks are dead weight that just keeps eating SSD (and the
    hot cache keeps eating RAM) until the disk fills or the prefill memory guard trips. Call
    this at each agent boundary to keep peak SSD/RAM use to roughly one agent's context.

    Auth: oMLX's admin endpoints want a session cookie, minted by POSTing the main API key to
    /admin/api/login; httpx.Client keeps that cookie for the follow-up clears. Fully non-fatal
    — a cache-clear hiccup must never abort a run.
    """
    s = get_settings()
    base = f"http://{s.llm_host}:{s.llm_port}"
    try:
        with httpx.Client(base_url=base, timeout=30) as c:
            c.post(
                "/admin/api/login",
                json={"api_key": s.llm_api_key, "remember": False},
            ).raise_for_status()
            ssd = c.post("/admin/api/ssd-cache/clear")
            ssd.raise_for_status()
            msg = f"cleared oMLX SSD KV cache ({ssd.json().get('total_deleted', '?')} blocks)"
            if hot:
                h = c.post("/admin/api/hot-cache/clear")
                h.raise_for_status()
                msg += f", hot cache reclaimed {h.json().get('bytes_reclaimed', 0) / 1e6:.0f}MB"
            logger.info(msg)
    except Exception as e:
        logger.warning("clear_kv_cache failed (non-fatal): {}", e)


def ensure_ready(warm: bool = True) -> bool:
    start_server()
    ok = wait_healthy(120)
    if ok and warm:
        warm_model()
    logger.info("oMLX ready: {}", ok)
    return ok


if __name__ == "__main__":
    ensure_ready()