#!/usr/bin/env python3
"""Layer 3 of the install-confidence ladder: first-run model smoke.

This is the layer that targets the failure you actually hit on a clean machine:
the recommended-setup **chat verification 500** (model loads, but the first
generation fails). It collapses the "build → copy → install on a second machine"
loop into one command on a Mac that has the MLX stack.

It reuses the REAL backend primitives (no parallel reimplementation):

1. ``ensure_mlx_venv`` (only with ``--install``) — make sure the MLX stack is there.
2. ``mlx_download_model`` + ``mlx_download_status`` — fetch a model (``--tiny`` uses
   the smallest tool-capable recommended model).
3. ``start_mlx_server`` — launch ``python -m mlx_vlm.server`` exactly as the app does.
4. ``_verify_local_mlx_chat_once`` — the same non-stream chat probe the wizard runs.
5. On failure: tail ``~/.mugi/mlx/server.log`` and classify (OOM / chat-template
   mismatch) with the same heuristic the verify route uses, so the actual cause
   is printed instead of an opaque "HTTP Error 500".

Requires Apple Silicon + an installed MLX venv (or ``--install``). It downloads a
model and starts a real server, so it is opt-in and NOT part of the normal suite.

Usage:
    python scripts/first_run_smoke.py --tiny          # smallest model, end-to-end
    python scripts/first_run_smoke.py --model PATH     # verify a local model dir
    python scripts/first_run_smoke.py --recommended    # this machine's top pick
    python scripts/first_run_smoke.py                  # verify the model in prefs

Exit 0 + "FIRST_RUN_SMOKE_OK" on success; non-zero with the classified cause +
server.log tail on failure.
"""

from __future__ import annotations

import argparse
import os
import sys
import time

REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if REPO_ROOT not in sys.path:
    sys.path.insert(0, REPO_ROOT)

# Smallest vision + tool-capable recommended model (see MLX_RECOMMENDED in app_mlx.py).
TINY_MODEL_REPO = "mlx-community/gemma-4-E2B-it-qat-8bit"


def _classify_server_log(tail: str) -> str:
    """Mirror mugi_backend_mlx_prefs_routes verify_chat classification."""
    low = (tail or "").lower()
    if (
        "out of memory" in low
        or "insufficient memory" in low
        or ("metal" in low and "memory" in low)
    ):
        return "mlx_out_of_memory"
    if "apply_chat_template" in low or "chat template" in low:
        return "mlx_template_mismatch"
    return ""


def _report_failure(stage: str, msg: str) -> None:
    import app_mlx

    print(f"FAIL: {stage}: {msg}", file=sys.stderr)
    try:
        tail = app_mlx._tail_log(os.path.join(app_mlx.MLX_DATA_DIR, "server.log"))
    except Exception:
        tail = ""
    code = _classify_server_log(tail)
    if code:
        print(f"[first-run-smoke] classified cause: {code}", file=sys.stderr)
    if tail and tail not in ("(log is empty)", "(could not read log)"):
        print("--- ~/.mugi/mlx/server.log (tail) ---", file=sys.stderr)
        print(tail, file=sys.stderr)


def _download_and_resolve(repo: str, download_timeout: float) -> str:
    import app_mlx

    print(f"[first-run-smoke] downloading {repo} (up to {int(download_timeout)}s)")
    res = app_mlx.mlx_download_model(repo)
    if not res.get("ok"):
        raise SystemExit(f"FAIL: download could not start: {res.get('message')}")
    job_id = res["job_id"]
    deadline = time.time() + download_timeout
    last_bucket = -1
    while time.time() < deadline:
        st = app_mlx.mlx_download_status(job_id)
        status = str(st.get("status") or "")
        bucket = int((st.get("progress") or 0.0) * 10)
        if bucket != last_bucket:
            print(f"[first-run-smoke]   {bucket * 10}% {st.get('message', '')}".rstrip())
            last_bucket = bucket
        if status == "completed":
            path = st.get("final_model_path")
            if not path:
                raise SystemExit("FAIL: download completed but no final_model_path")
            print(f"[first-run-smoke] model ready -> {path}")
            return str(path)
        if status == "failed":
            raise SystemExit(f"FAIL: download failed: {st.get('message')}")
        time.sleep(2.0)
    raise SystemExit("FAIL: download timed out")


def _resolve_model(args) -> str:
    import app_mlx

    if args.model:
        local = os.path.expanduser(args.model)
        if os.path.isdir(local):
            return local
        if args.download:
            return _download_and_resolve(args.model, args.download_timeout)
        # Repo id without --download: let start_mlx_server/mlx-vlm resolve it.
        return args.model
    if args.tiny:
        return _download_and_resolve(TINY_MODEL_REPO, args.download_timeout)
    if args.recommended:
        recs = app_mlx.get_mlx_recommended_models()
        if not recs:
            raise SystemExit("FAIL: no recommended models available")
        return _download_and_resolve(str(recs[0]["repo"]), args.download_timeout)
    model_path, _ctx = app_mlx._selected_mlx_model_and_context_from_prefs()
    if not model_path:
        raise SystemExit(
            "FAIL: no model to verify. Pass --tiny, --recommended, --model PATH, "
            "or select a model in the app first."
        )
    return model_path


def run_first_run_smoke(args) -> int:
    if sys.platform != "darwin":
        print(
            "[first-run-smoke] WARNING: MLX targets Apple Silicon; this run may not "
            "reflect real behavior on a non-darwin host.",
            file=sys.stderr,
        )

    import app_mlx

    if not app_mlx.mlx_venv_exists():
        if not args.install:
            print(
                "FAIL: MLX venv not installed. Re-run with --install (slow: installs "
                "the MLX stack) or install it from the app's Settings → LLMs first.",
                file=sys.stderr,
            )
            return 1
        print("[first-run-smoke] installing MLX stack (ensure_mlx_venv)...")
        ok, msg = app_mlx.ensure_mlx_venv()
        if not ok:
            print(f"FAIL: MLX install failed: {msg}", file=sys.stderr)
            return 1
        print(f"[first-run-smoke] MLX stack ready: {msg}")

    model_path = _resolve_model(args)
    print(f"[first-run-smoke] verifying model: {model_path}")

    server_started = False
    try:
        ok_start, msg_start = app_mlx.start_mlx_server(model_path)
        if not ok_start:
            _report_failure("MLX server start", msg_start)
            return 1
        server_started = True

        # Wait for the process to bind, then give weights time to load.
        for _ in range(30):
            if app_mlx.mlx_server_running():
                break
            time.sleep(1.0)
        time.sleep(5.0)

        last_err = ""
        for attempt in range(4):
            ok_verify, err = app_mlx._verify_local_mlx_chat_once(
                timeout_seconds=args.verify_timeout
            )
            if ok_verify:
                print("FIRST_RUN_SMOKE_OK")
                print(
                    "PASS: model loaded and the chat-verification probe generated a "
                    "response (the exact step that 500'd on the clean machine)."
                )
                return 0
            last_err = err
            print(f"[first-run-smoke] verify attempt {attempt + 1}/4 failed: {err}")
            if attempt < 3:
                time.sleep(5.0)

        _report_failure("chat verification", last_err)
        return 1
    finally:
        if server_started and not args.keep_running:
            print("[first-run-smoke] stopping MLX server")
            try:
                app_mlx.stop_mlx_server()
            except Exception:
                pass


def main() -> int:
    parser = argparse.ArgumentParser(description="First-run model smoke (MLX).")
    src = parser.add_mutually_exclusive_group()
    src.add_argument("--tiny", action="store_true", help=f"use {TINY_MODEL_REPO}")
    src.add_argument("--recommended", action="store_true", help="this machine's top pick")
    src.add_argument("--model", help="local model dir or HF repo id to verify")
    parser.add_argument("--download", action="store_true", help="download --model if it's a repo id")
    parser.add_argument("--install", action="store_true", help="install MLX stack if missing (slow)")
    parser.add_argument("--keep-running", action="store_true", help="leave the MLX server running")
    parser.add_argument("--verify-timeout", type=float, default=90.0, help="chat probe timeout (s)")
    parser.add_argument("--download-timeout", type=float, default=1800.0, help="model download timeout (s)")
    args = parser.parse_args()
    return run_first_run_smoke(args)


if __name__ == "__main__":
    raise SystemExit(main())
