setup_models.py

#!/usr/bin/env python3
"""
setup_models.py — RecordNote local AI model setup
==================================================

Cross-platform orchestrator for preparing local models with:
- Runtime capability detection
- Shell-first step executor with Python fallback
- Automatic path/model discovery and registry normalization
- Idempotent .env and shell export-script generation
- Preflight/postflight validation gates
- Safe reruns (skip-if-valid, lock files, atomic writes)
- Structured JSON step logs and final machine-readable summary
"""

from __future__ import annotations

import argparse
import contextlib
import json
import os
import platform
import shutil
import socket
import subprocess
import sys
import tempfile
import textwrap
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple


# ─────────────────────────────────────────────────────────────────────────────
# Terminal colors
# ─────────────────────────────────────────────────────────────────────────────
_USE_COLOR = sys.stdout.isatty() and os.getenv("TERM") != "dumb" and (
    platform.system() != "Windows"
    or os.getenv("WT_SESSION") is not None
    or os.getenv("TERM_PROGRAM") == "vscode"
    or int(os.getenv("ANSICON", "0")) > 0
)


def _c(code: str, text: str) -> str:
    return f"\033[{code}m{text}\033[0m" if _USE_COLOR else text


def ok(msg: str) -> str:
    return _c("32", f"  ✓  {msg}")


def info(msg: str) -> str:
    return _c("36", f"  ℹ  {msg}")


def warn(msg: str) -> str:
    return _c("33", f"  ⚠  {msg}")


def err(msg: str) -> str:
    return _c("31", f"  ✗  {msg}")


def bold(msg: str) -> str:
    return _c("1", msg)


def dim(msg: str) -> str:
    return _c("2", msg)


# ─────────────────────────────────────────────────────────────────────────────
# Constants/catalogs
# ─────────────────────────────────────────────────────────────────────────────
SCRIPT_DIR = Path(__file__).resolve().parent
MODELS_DIR = SCRIPT_DIR / "models"
GGUF_DIR = MODELS_DIR / "gguf"

WHISPER_MODELS: Dict[str, Dict[str, Any]] = {
    "tiny": {
        "repo": "Systran/faster-whisper-tiny",
        "local_dir": MODELS_DIR / "faster-whisper-tiny",
        "description": "Tiny (~39 MB) — fastest",
        "vram_gb": 0.5,
    },
    "base": {
        "repo": "Systran/faster-whisper-base",
        "local_dir": MODELS_DIR / "faster-whisper-base",
        "description": "Base (~74 MB)",
        "vram_gb": 0.5,
    },
    "small": {
        "repo": "Systran/faster-whisper-small",
        "local_dir": MODELS_DIR / "faster-whisper-small",
        "description": "Small (~244 MB)",
        "vram_gb": 1.0,
    },
    "medium": {
        "repo": "Systran/faster-whisper-medium",
        "local_dir": MODELS_DIR / "faster-whisper-medium",
        "description": "Medium (~769 MB) — recommended",
        "vram_gb": 2.0,
    },
    "large-v3": {
        "repo": "Systran/faster-whisper-large-v3",
        "local_dir": MODELS_DIR / "faster-whisper-large-v3",
        "description": "Large v3 (~3 GB)",
        "vram_gb": 6.0,
    },
}

CT2_MODELS: Dict[str, Dict[str, Any]] = {
    "nllb-200-600M": {
        "hf_repo": "facebook/nllb-200-distilled-600M",
        "ct2_dir": MODELS_DIR / "nllb-200-distilled-600M-ct2",
        "description": "NLLB-200 600M (int8)",
        "vram_gb": 1.0,
    },
    "nllb-200-1.3B": {
        "hf_repo": "facebook/nllb-200-distilled-1.3B",
        "ct2_dir": MODELS_DIR / "nllb-200-distilled-1.3B-ct2",
        "description": "NLLB-200 1.3B (int8)",
        "vram_gb": 3.0,
    },
}

GGUF_MODELS: Dict[str, Dict[str, Any]] = {
    "qwen2.5-7b-instruct": {
        "repo": "bartowski/Qwen2.5-7B-Instruct-GGUF",
        "filename": "Qwen2.5-7B-Instruct-Q4_K_M.gguf",
        "description": "Qwen 2.5 7B Q4_K_M — recommended",
        "vram_gb": 6.0,
    },
    "llama3.2-3b-instruct": {
        "repo": "bartowski/Llama-3.2-3B-Instruct-GGUF",
        "filename": "Llama-3.2-3B-Instruct-Q4_K_M.gguf",
        "description": "Llama 3.2 3B Q4_K_M",
        "vram_gb": 3.0,
    },
    "mistral-7b-instruct": {
        "repo": "TheBloke/Mistral-7B-Instruct-v0.2-GGUF",
        "filename": "mistral-7b-instruct-v0.2.Q4_K_M.gguf",
        "description": "Mistral 7B Q4_K_M",
        "vram_gb": 6.0,
    },
}

OLLAMA_MODELS: Dict[str, Dict[str, Any]] = {
    "llama3.1:8b": {"description": "Llama 3.1 8B", "vram_gb": 6.0},
    "qwen2.5:7b": {"description": "Qwen 2.5 7B", "vram_gb": 6.0},
    "phi3:mini": {"description": "Phi-3 Mini", "vram_gb": 3.0},
}


# ─────────────────────────────────────────────────────────────────────────────
# Structured errors/logging
# ─────────────────────────────────────────────────────────────────────────────
class SetupError(Exception):
    def __init__(self, message: str, category: str = "runtime") -> None:
        super().__init__(message)
        self.category = category


@dataclass
class RuntimeCapabilities:
    platform: str
    machine: str
    python_version: str
    has_bash: bool
    has_sh: bool
    has_pwsh: bool
    has_powershell: bool
    has_cmd: bool
    ci: bool
    container: bool
    wsl: bool
    virtualenv: bool
    shell_preference: List[str] = field(default_factory=list)


@dataclass
class StepResult:
    step: str
    action: str
    backend: str
    status: str
    duration_ms: int
    error_code: str = ""
    message: str = ""


def _json_log(payload: Dict[str, Any]) -> None:
    print(json.dumps(payload, ensure_ascii=False))


def emit_step_log(result: StepResult) -> None:
    _json_log(
        {
            "event": "setup_step",
            "step": result.step,
            "action": result.action,
            "backend": result.backend,
            "status": result.status,
            "duration_ms": result.duration_ms,
            "error_code": result.error_code or None,
            "message": result.message or None,
        }
    )


# ─────────────────────────────────────────────────────────────────────────────
# Runtime detection
# ─────────────────────────────────────────────────────────────────────────────

def detect_runtime_capabilities() -> RuntimeCapabilities:
    sysname = platform.system().lower()
    if "windows" in sysname:
        p = "windows"
        shell_order = ["pwsh", "powershell", "cmd", "bash", "sh"]
    elif "darwin" in sysname:
        p = "darwin"
        shell_order = ["bash", "sh", "pwsh", "powershell", "cmd"]
    else:
        p = "linux"
        shell_order = ["bash", "sh", "pwsh", "powershell", "cmd"]

    ci = any(os.getenv(k) for k in ("CI", "GITHUB_ACTIONS", "BUILD_ID", "TF_BUILD"))
    container = Path("/.dockerenv").exists() or "container" in os.getenv("container", "").lower()
    wsl = "microsoft" in platform.release().lower() or bool(os.getenv("WSL_DISTRO_NAME"))
    virtualenv = hasattr(sys, "real_prefix") or (getattr(sys, "base_prefix", sys.prefix) != sys.prefix)

    return RuntimeCapabilities(
        platform=p,
        machine=platform.machine(),
        python_version=sys.version.split()[0],
        has_bash=shutil.which("bash") is not None,
        has_sh=shutil.which("sh") is not None,
        has_pwsh=shutil.which("pwsh") is not None,
        has_powershell=shutil.which("powershell") is not None,
        has_cmd=shutil.which("cmd") is not None or p == "windows",
        ci=ci,
        container=container,
        wsl=wsl,
        virtualenv=virtualenv,
        shell_preference=shell_order,
    )


def _best_shell(caps: RuntimeCapabilities) -> Optional[str]:
    available = {
        "bash": caps.has_bash,
        "sh": caps.has_sh,
        "pwsh": caps.has_pwsh,
        "powershell": caps.has_powershell,
        "cmd": caps.has_cmd,
    }
    for name in caps.shell_preference:
        if available.get(name):
            return name
    return None


# ─────────────────────────────────────────────────────────────────────────────
# Two-layer executor: shell-first, python fallback
# ─────────────────────────────────────────────────────────────────────────────

def _run_shell(shell: str, command: str, timeout: int = 20) -> subprocess.CompletedProcess:
    if shell in {"bash", "sh"}:
        return subprocess.run([shell, "-lc", command], capture_output=True, text=True, timeout=timeout)
    if shell in {"pwsh", "powershell"}:
        return subprocess.run([shell, "-NoProfile", "-Command", command], capture_output=True, text=True, timeout=timeout)
    return subprocess.run(["cmd", "/c", command], capture_output=True, text=True, timeout=timeout)


def execute_step(
    *,
    step: str,
    action: str,
    caps: RuntimeCapabilities,
    shell_command_builder: Optional[Callable[[str], str]],
    python_action: Callable[[], Any],
    acceptable_rc: Tuple[int, ...] = (0,),
) -> Tuple[StepResult, Any]:
    started = time.time()
    shell = _best_shell(caps)

    if shell and shell_command_builder is not None:
        try:
            proc = _run_shell(shell, shell_command_builder(shell))
            if proc.returncode in acceptable_rc:
                res = StepResult(step, action, "shell", "ok", int((time.time() - started) * 1000), message=(proc.stdout or "").strip())
                emit_step_log(res)
                return res, proc
        except Exception as exc:
            shell_err = str(exc)
        else:
            shell_err = (proc.stderr or proc.stdout or "").strip()

        # fallback
        _json_log(
            {
                "event": "setup_step_fallback",
                "step": step,
                "action": action,
                "from": "shell",
                "to": "python",
                "reason": shell_err or "shell_failed",
            }
        )

    try:
        value = python_action()
        res = StepResult(step, action, "python", "ok", int((time.time() - started) * 1000))
        emit_step_log(res)
        return res, value
    except Exception as exc:
        res = StepResult(
            step,
            action,
            "python",
            "error",
            int((time.time() - started) * 1000),
            error_code="python_fallback_failed",
            message=str(exc),
        )
        emit_step_log(res)
        return res, None


# ─────────────────────────────────────────────────────────────────────────────
# Rerun safety utilities
# ─────────────────────────────────────────────────────────────────────────────
@contextlib.contextmanager
def file_lock(lock_path: Path):
    lock_path.parent.mkdir(parents=True, exist_ok=True)
    fd: Optional[int] = None
    try:
        while True:
            try:
                fd = os.open(str(lock_path), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
                os.write(fd, str(os.getpid()).encode("utf-8"))
                break
            except FileExistsError:
                time.sleep(0.2)
        yield
    finally:
        if fd is not None:
            os.close(fd)
        with contextlib.suppress(FileNotFoundError):
            lock_path.unlink()


def atomic_write_text(path: Path, content: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=str(path.parent), delete=False) as tmp:
        tmp.write(content)
        tmp_name = tmp.name
    os.replace(tmp_name, path)


# ─────────────────────────────────────────────────────────────────────────────
# Generic checks
# ─────────────────────────────────────────────────────────────────────────────

def _has_network(host: str, port: int = 443, timeout: float = 3.0) -> bool:
    try:
        with socket.create_connection((host, port), timeout=timeout):
            return True
    except OSError:
        return False


def _shell_which(name: str, caps: RuntimeCapabilities) -> bool:
    def shell_cmd(shell: str) -> str:
        if shell in {"bash", "sh"}:
            return f"command -v {name} >/dev/null"
        if shell in {"pwsh", "powershell"}:
            return f"if (Get-Command {name} -ErrorAction SilentlyContinue) {{ exit 0 }} else {{ exit 1 }}"
        return f"where {name} >nul 2>nul"

    def py_fallback() -> bool:
        return shutil.which(name) is not None

    _, value = execute_step(
        step="dependency",
        action=f"which:{name}",
        caps=caps,
        shell_command_builder=shell_cmd,
        python_action=py_fallback,
        acceptable_rc=(0,),
    )
    if isinstance(value, subprocess.CompletedProcess):
        return value.returncode == 0
    return bool(value)


def _ollama_model_exists(tag: str, caps: RuntimeCapabilities) -> bool:
    name = tag.split(":")[0]

    def shell_cmd(shell: str) -> str:
        if shell in {"bash", "sh"}:
            return f"ollama list | grep -E '^{name}[: ]' >/dev/null"
        if shell in {"pwsh", "powershell"}:
            return f"$o=ollama list; if($o -match '^{name}[: ]'){{exit 0}} else {{exit 1}}"
        return f"ollama list | findstr /R /C:\"^{name}[: ]\" >nul"

    def py_fallback() -> bool:
        proc = subprocess.run(["ollama", "list"], capture_output=True, text=True, timeout=10)
        return name in proc.stdout

    _, value = execute_step(
        step="ollama_check",
        action=f"check_model:{tag}",
        caps=caps,
        shell_command_builder=shell_cmd,
        python_action=py_fallback,
        acceptable_rc=(0,),
    )
    if isinstance(value, subprocess.CompletedProcess):
        return value.returncode == 0
    return bool(value)


# ─────────────────────────────────────────────────────────────────────────────
# Dependency check
# ─────────────────────────────────────────────────────────────────────────────

def check_dependencies(caps: RuntimeCapabilities, need_ollama: bool) -> Tuple[bool, List[str]]:
    all_ok = True
    warnings_list: List[str] = []

    if sys.version_info < (3, 10):
        print(err(f"Python 3.10+ required (found {sys.version.split()[0]})."))
        all_ok = False
    else:
        print(ok(f"Python {sys.version.split()[0]}"))

    try:
        import huggingface_hub  # noqa: F401
        print(ok("huggingface_hub is installed"))
    except ImportError:
        print(err("huggingface_hub not found. Run: pip install huggingface_hub>=1.9.0"))
        all_ok = False

    try:
        import ctranslate2  # noqa: F401
        print(ok("ctranslate2 is installed"))
    except ImportError:
        msg = "ctranslate2 not found — CTranslate2 conversion unavailable."
        print(warn(msg))
        warnings_list.append(msg)

    try:
        import transformers  # noqa: F401
        print(ok("transformers is installed"))
    except ImportError:
        msg = "transformers not found — CTranslate2 conversion may fail."
        print(warn(msg))
        warnings_list.append(msg)

    try:
        import llama_cpp  # noqa: F401
        print(ok("llama-cpp-python is installed"))
    except ImportError:
        msg = "llama-cpp-python not found — GGUF can download but runtime may fail."
        print(warn(msg))
        warnings_list.append(msg)

    has_ollama = _shell_which("ollama", caps)
    if has_ollama:
        print(ok("Ollama binary found"))
    else:
        msg = "Ollama binary not found."
        if need_ollama:
            print(err(msg))
            all_ok = False
        else:
            print(warn(msg))
            warnings_list.append(msg)

    return all_ok, warnings_list


# ─────────────────────────────────────────────────────────────────────────────
# GPU detection
# ─────────────────────────────────────────────────────────────────────────────

def _nvidia_vram_gb() -> Optional[float]:
    try:
        out = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"],
            stderr=subprocess.DEVNULL,
            text=True,
        )
        mb = float(out.strip().splitlines()[0])
        return mb / 1024
    except Exception:
        return None


def detect_gpu() -> Dict[str, Any]:
    nvidia = _nvidia_vram_gb()
    if nvidia is not None:
        return {"type": "nvidia", "vram_gb": nvidia, "label": f"NVIDIA GPU ({nvidia:.0f} GB VRAM)"}

    if platform.system() == "Darwin" and platform.machine() == "arm64":
        return {"type": "apple_silicon", "vram_gb": None, "label": "Apple Silicon (unified memory)"}

    if shutil.which("rocminfo") is not None or os.path.exists("/dev/kfd"):
        return {"type": "amd_rocm", "vram_gb": None, "label": "AMD GPU (ROCm)"}

    return {"type": "cpu", "vram_gb": 0, "label": "CPU only"}


# ─────────────────────────────────────────────────────────────────────────────
# Selection
# ─────────────────────────────────────────────────────────────────────────────

def default_selections(gpu: Dict[str, Any], caps: RuntimeCapabilities) -> Tuple[List[str], List[str], List[str], List[str]]:
    vram = gpu.get("vram_gb") or 0
    gtype = gpu["type"]

    if gtype == "apple_silicon" or vram >= 8:
        whisper = ["medium"]
        gguf = ["qwen2.5-7b-instruct"]
        ollama = ["llama3.1:8b"]
    elif vram >= 4:
        whisper = ["small"]
        gguf = ["llama3.2-3b-instruct"]
        ollama = ["phi3:mini"]
    else:
        whisper = ["base"]
        gguf = ["llama3.2-3b-instruct"]
        ollama = ["phi3:mini"]

    if caps.ci:
        ollama = []
    if not _shell_which("ollama", caps):
        ollama = []

    return whisper, ["nllb-200-600M"], gguf, ollama


def _print_catalog(title: str, catalog: Dict[str, Dict[str, Any]], recommended: Optional[str] = None) -> None:
    print(f"\n{bold(title)}")
    print(dim("  " + "─" * 66))
    for i, (key, cfg) in enumerate(catalog.items(), 1):
        mark = _c("32", " ← recommended") if recommended and key == recommended else ""
        print(f"  {dim(str(i)+'.'):>5}  {key:<30} {dim(cfg.get('description', key))}{mark}")
    print(f"  {dim('0.'):>5}  {'(none — skip this category)':<30}")


def _ask_choices(catalog: Dict[str, Dict[str, Any]]) -> List[str]:
    keys = list(catalog.keys())
    while True:
        raw = input(dim("  Enter numbers (e.g. 1 3), 'all', or 0 to skip: ")).strip()
        if raw in ("0", ""):
            return []
        if raw.lower() == "all":
            return keys
        parts = raw.replace(",", " ").split()
        try:
            selected = []
            for part in parts:
                idx = int(part)
                if idx == 0:
                    return []
                if 1 <= idx <= len(keys):
                    selected.append(keys[idx - 1])
                else:
                    raise ValueError
            return list(dict.fromkeys(selected))
        except ValueError:
            print(warn(f"Invalid input '{raw}'."))


def interactive_menu(caps: RuntimeCapabilities) -> Tuple[List[str], List[str], List[str], List[str]]:
    _print_catalog("1. faster-whisper", WHISPER_MODELS, recommended="medium")
    whisper = _ask_choices(WHISPER_MODELS)

    _print_catalog("2. CTranslate2/NLLB", CT2_MODELS, recommended="nllb-200-600M")
    ct2 = _ask_choices(CT2_MODELS)

    _print_catalog("3. GGUF", GGUF_MODELS, recommended="qwen2.5-7b-instruct")
    gguf = _ask_choices(GGUF_MODELS)

    if _shell_which("ollama", caps):
        _print_catalog("4. Ollama", OLLAMA_MODELS, recommended="llama3.1:8b")
        ollama = _ask_choices(OLLAMA_MODELS)
    else:
        print(warn("Ollama not available; skipping."))
        ollama = []

    return whisper, ct2, gguf, ollama


# ─────────────────────────────────────────────────────────────────────────────
# Model discovery / registry normalization
# ─────────────────────────────────────────────────────────────────────────────

def _expand_abs(path_val: str) -> str:
    return str(Path(path_val).expanduser().resolve())


def _first_existing_file(candidates: List[Path]) -> Optional[str]:
    for c in candidates:
        if c.exists() and c.is_file():
            return str(c.resolve())
    return None


def _first_existing_dir_marker(candidates: List[Path], marker: str) -> Optional[str]:
    for c in candidates:
        if c.exists() and c.is_dir() and (c / marker).exists():
            return str(c.resolve())
    return None


def discover_model_registry(
    *,
    model_root: Optional[str],
    llama_cpp_model_path: Optional[str],
    whisper_cpp_model_path: Optional[str],
    nllb_model_dir: Optional[str],
    caps: RuntimeCapabilities,
) -> Dict[str, Any]:
    env_root = os.getenv("MODEL_REGISTRY_ROOT", "").strip()
    if model_root:
        root = Path(model_root).expanduser().resolve()
        root_conf = "high"
    elif env_root:
        root = Path(env_root).expanduser().resolve()
        root_conf = "high"
    elif MODELS_DIR.exists():
        root = MODELS_DIR.resolve()
        root_conf = "medium"
    else:
        root = (Path.home() / ".recordnote" / "models").resolve()
        root_conf = "low"

    # faster-whisper dir
    fw_env = os.getenv("STT_PROVIDER_FASTER_WHISPER_MODEL_DIR", "").strip()
    if fw_env and Path(fw_env).exists():
        fw_dir, fw_conf = _expand_abs(fw_env), "high"
    else:
        fw_candidates = [root / "faster-whisper"] + [root / f"faster-whisper-{k}" for k in WHISPER_MODELS]
        fw = None
        for c in fw_candidates:
            if c.exists() and c.is_dir() and any(c.iterdir()):
                fw = str(c.resolve())
                break
        fw_dir = fw or str((root / "faster-whisper-medium").resolve())
        fw_conf = "high" if fw else "low"

    # translator dir marker: model.bin
    if nllb_model_dir and Path(nllb_model_dir).exists() and (Path(nllb_model_dir) / "model.bin").exists():
        translator_dir, tr_conf = _expand_abs(nllb_model_dir), "high"
    elif os.getenv("NLLB_MODEL_DIR", "") and Path(os.getenv("NLLB_MODEL_DIR", "")).exists() and (Path(os.getenv("NLLB_MODEL_DIR", "")) / "model.bin").exists():
        translator_dir, tr_conf = _expand_abs(os.getenv("NLLB_MODEL_DIR", "")), "high"
    else:
        tr = _first_existing_dir_marker(
            [root / "nllb-200-distilled-600M", root / "nllb-200-distilled-600M-ct2", root / "nllb-200-distilled-1.3B-ct2"],
            "model.bin",
        )
        translator_dir = tr or str((root / "nllb-200-distilled-600M-ct2").resolve())
        tr_conf = "high" if tr else "low"

    # llama cpp model marker: *.gguf
    if llama_cpp_model_path and Path(llama_cpp_model_path).exists():
        llama_model, llama_conf = _expand_abs(llama_cpp_model_path), "high"
    elif os.getenv("LLAMA_CPP_MODEL_PATH", "") and Path(os.getenv("LLAMA_CPP_MODEL_PATH", "")).exists():
        llama_model, llama_conf = _expand_abs(os.getenv("LLAMA_CPP_MODEL_PATH", "")), "high"
    else:
        gguf_roots = [root / "gguf", SCRIPT_DIR / "models" / "gguf", Path.home() / ".cache" / "huggingface" / "hub"]
        if caps.platform != "windows":
            gguf_roots.append(Path("/models") / "gguf")
        hits: List[Path] = []
        for d in gguf_roots:
            if d.exists() and d.is_dir():
                hits.extend(sorted(d.rglob("*.gguf")))
        llama_model = str(hits[0].resolve()) if hits else ""
        llama_conf = "medium" if hits else "low"

    # whisper.cpp marker: ggml*.bin
    if whisper_cpp_model_path and Path(whisper_cpp_model_path).exists():
        whisper_cpp_model, whisper_cpp_conf = _expand_abs(whisper_cpp_model_path), "high"
    elif os.getenv("STT_PROVIDER_WHISPER_CPP_MODEL", "") and Path(os.getenv("STT_PROVIDER_WHISPER_CPP_MODEL", "")).exists():
        whisper_cpp_model, whisper_cpp_conf = _expand_abs(os.getenv("STT_PROVIDER_WHISPER_CPP_MODEL", "")), "high"
    else:
        wc = _first_existing_file(sorted((root if root.exists() else MODELS_DIR).rglob("ggml*.bin"))) if (root.exists() or MODELS_DIR.exists()) else None
        whisper_cpp_model = wc or ""
        whisper_cpp_conf = "medium" if wc else "low"

    # ollama assets marker: model appears in ollama list
    ollama_found: List[str] = []
    if _shell_which("ollama", caps):
        for tag in OLLAMA_MODELS:
            if _ollama_model_exists(tag, caps):
                ollama_found.append(tag)

    preferred_whisper = "medium"
    fw_name = Path(fw_dir).name
    if fw_name.startswith("faster-whisper-"):
        preferred_whisper = fw_name.replace("faster-whisper-", "")

    return {
        "root": str(root),
        "faster_whisper_dir": fw_dir,
        "faster_whisper_model": preferred_whisper,
        "translator_dir": translator_dir,
        "llama_cpp_model": llama_model,
        "whisper_cpp_model": whisper_cpp_model,
        "ollama_models": ollama_found,
        "confidence": {
            "root": root_conf,
            "faster_whisper_dir": fw_conf,
            "translator_dir": tr_conf,
            "llama_cpp_model": llama_conf,
            "whisper_cpp_model": whisper_cpp_conf,
            "ollama_models": "high" if ollama_found else "low",
        },
    }


# ─────────────────────────────────────────────────────────────────────────────
# Env management
# ─────────────────────────────────────────────────────────────────────────────

def upsert_env_file(env_path: Path, updates: Dict[str, str]) -> None:
    lines = env_path.read_text(encoding="utf-8").splitlines() if env_path.exists() else []
    seen: set[str] = set()
    out: List[str] = []

    for line in lines:
        stripped = line.strip()
        if not stripped or stripped.startswith("#") or "=" not in line:
            out.append(line)
            continue
        key = stripped.split("=", 1)[0].strip()
        if key in updates:
            if key not in seen:
                out.append(f"{key}={updates[key]}")
                seen.add(key)
            # skip duplicate occurrences
        else:
            if key not in seen:
                out.append(line)
                seen.add(key)

    for k, v in updates.items():
        if k not in seen:
            out.append(f"{k}={v}")

    atomic_write_text(env_path, "\n".join(out).rstrip() + "\n")


def write_export_scripts(base_dir: Path, updates: Dict[str, str]) -> Dict[str, str]:
    posix_path = base_dir / ".recordnote-models.env"
    ps_path = base_dir / ".recordnote-models.ps1"
    cmd_path = base_dir / ".recordnote-models.cmd"

    posix_lines = [f"export {k}='{v.replace("'", "'\\''")}'" for k, v in updates.items()]
    ps_lines = [f"$Env:{k} = \"{v.replace(chr(34), '`"')}\"" for k, v in updates.items()]
    cmd_lines = [f"set {k}={v}" for k, v in updates.items()]

    atomic_write_text(posix_path, "\n".join(posix_lines) + "\n")
    atomic_write_text(ps_path, "\n".join(ps_lines) + "\n")
    atomic_write_text(cmd_path, "\n".join(cmd_lines) + "\n")

    return {
        "posix": str(posix_path.resolve()),
        "powershell": str(ps_path.resolve()),
        "cmd": str(cmd_path.resolve()),
    }


# ─────────────────────────────────────────────────────────────────────────────
# Preflight/postflight gates
# ─────────────────────────────────────────────────────────────────────────────

def preflight_validate(caps: RuntimeCapabilities, registry: Dict[str, Any], require_network: bool) -> Dict[str, Any]:
    issues: List[Dict[str, str]] = []

    root = Path(registry["root"])
    try:
        root.mkdir(parents=True, exist_ok=True)
        probe = root / ".write-probe"
        probe.write_text("ok", encoding="utf-8")
        probe.unlink(missing_ok=True)
    except Exception as exc:
        issues.append({"category": "filesystem", "message": f"Model root not writable: {exc}"})

    if require_network and not _has_network("huggingface.co"):
        issues.append(
            {
                "category": "dependency",
                "message": "Cannot reach huggingface.co:443 (required for selected downloads).",
            }
        )

    if caps.platform not in {"linux", "darwin", "windows"}:
        issues.append({"category": "detection", "message": f"Unknown platform: {caps.platform}"})

    return {"ok": len(issues) == 0, "issues": issues}


def postflight_validate(
    *,
    whisper_paths: List[str],
    ct2_paths: List[str],
    gguf_paths: List[str],
    ollama_tags: List[str],
    registry: Dict[str, Any],
    caps: RuntimeCapabilities,
) -> Dict[str, Any]:
    checks: List[Dict[str, Any]] = []

    for p in whisper_paths:
        path = Path(p)
        checks.append({"name": "faster_whisper", "path": p, "ok": path.exists() and any(path.iterdir())})

    for p in ct2_paths:
        marker = Path(p) / "model.bin"
        checks.append({"name": "ct2", "path": p, "ok": marker.exists()})

    for p in gguf_paths:
        path = Path(p)
        checks.append({"name": "gguf", "path": p, "ok": path.exists() and path.is_file()})

    for tag in ollama_tags:
        checks.append({"name": "ollama", "path": f"tag:{tag}", "ok": _ollama_model_exists(tag, caps)})

    # ModelRegistryService-style checks
    checks.extend(
        [
            {"name": "registry_root", "path": registry["root"], "ok": Path(registry["root"]).exists()},
            {
                "name": "registry_faster_whisper_dir",
                "path": registry["faster_whisper_dir"],
                "ok": Path(registry["faster_whisper_dir"]).exists(),
            },
            {
                "name": "registry_translator_dir",
                "path": registry["translator_dir"],
                "ok": Path(registry["translator_dir"]).exists(),
            },
            {
                "name": "registry_llama_cpp_model",
                "path": registry.get("llama_cpp_model", ""),
                "ok": (not registry.get("llama_cpp_model")) or Path(registry["llama_cpp_model"]).exists(),
            },
            {
                "name": "registry_whisper_cpp_model",
                "path": registry.get("whisper_cpp_model", ""),
                "ok": (not registry.get("whisper_cpp_model")) or Path(registry["whisper_cpp_model"]).exists(),
            },
        ]
    )

    failed = [c for c in checks if not c["ok"]]
    return {"ok": len(failed) == 0, "checks": checks, "failed": failed}


# ─────────────────────────────────────────────────────────────────────────────
# Download/setup steps (independent, partial-success safe)
# ─────────────────────────────────────────────────────────────────────────────

def _whisper_exists(cfg: Dict[str, Any]) -> bool:
    d = Path(cfg["local_dir"])
    return d.exists() and any(d.iterdir())


def _ct2_exists(cfg: Dict[str, Any]) -> bool:
    return (Path(cfg["ct2_dir"]) / "model.bin").exists()


def _gguf_exists(cfg: Dict[str, Any]) -> bool:
    return (GGUF_DIR / cfg["filename"]).exists()


def download_whisper(keys: List[str], force: bool) -> List[str]:
    try:
        from huggingface_hub import snapshot_download
    except ImportError:
        print(err("huggingface_hub missing; skipping Whisper downloads."))
        return []

    done: List[str] = []
    for key in keys:
        cfg = WHISPER_MODELS[key]
        target = Path(cfg["local_dir"])
        with file_lock(target / ".setup.lock"):
            if not force and _whisper_exists(cfg):
                print(ok(f"Whisper {key} already valid; skipping."))
                done.append(str(target.resolve()))
                continue
            try:
                target.mkdir(parents=True, exist_ok=True)
                snapshot_download(
                    repo_id=cfg["repo"],
                    local_dir=str(target),
                    local_dir_use_symlinks=False,
                    ignore_patterns=["*.msgpack", "*.h5", "flax_model*", "tf_model*"],
                )
                done.append(str(target.resolve()))
                print(ok(f"Whisper {key} ready: {target}"))
            except Exception as exc:
                print(err(f"Whisper {key} failed: {exc}"))
    return done


def download_ctranslate2(keys: List[str], force: bool) -> List[str]:
    try:
        from ctranslate2.converters import TransformersConverter
    except ImportError:
        print(err("ctranslate2 missing; skipping CT2 conversions."))
        return []

    done: List[str] = []
    for key in keys:
        cfg = CT2_MODELS[key]
        target = Path(cfg["ct2_dir"])
        with file_lock(target / ".setup.lock"):
            if not force and _ct2_exists(cfg):
                print(ok(f"CT2 {key} already valid; skipping."))
                done.append(str(target.resolve()))
                continue
            try:
                target.mkdir(parents=True, exist_ok=True)
                converter = TransformersConverter(cfg["hf_repo"], low_cpu_mem_usage=True)
                converter.convert(str(target), quantization="int8", force=force)
                done.append(str(target.resolve()))
                print(ok(f"CT2 {key} ready: {target}"))
            except Exception as exc:
                print(err(f"CT2 {key} failed: {exc}"))
    return done


def download_gguf(keys: List[str], force: bool) -> List[str]:
    try:
        from huggingface_hub import hf_hub_download
    except ImportError:
        print(err("huggingface_hub missing; skipping GGUF downloads."))
        return []

    GGUF_DIR.mkdir(parents=True, exist_ok=True)
    done: List[str] = []
    for key in keys:
        cfg = GGUF_MODELS[key]
        target = GGUF_DIR / cfg["filename"]
        with file_lock(target.with_suffix(target.suffix + ".lock")):
            if not force and _gguf_exists(cfg):
                print(ok(f"GGUF {key} already valid; skipping."))
                done.append(str(target.resolve()))
                continue
            try:
                cached = hf_hub_download(repo_id=cfg["repo"], filename=cfg["filename"])
                cached_path = Path(cached)
                if cached_path.resolve() != target.resolve():
                    shutil.copy2(cached, target)
                done.append(str(target.resolve()))
                print(ok(f"GGUF {key} ready: {target}"))
            except Exception as exc:
                print(err(f"GGUF {key} failed: {exc}"))
    return done


def pull_ollama_models(tags: List[str], caps: RuntimeCapabilities, force: bool) -> List[str]:
    if not _shell_which("ollama", caps):
        print(warn("Ollama missing; skipping Ollama pulls."))
        return []

    done: List[str] = []
    for tag in tags:
        if not force and _ollama_model_exists(tag, caps):
            print(ok(f"Ollama {tag} already present; skipping."))
            done.append(tag)
            continue
        try:
            subprocess.run(["ollama", "pull", tag], check=True)
            done.append(tag)
            print(ok(f"Ollama {tag} ready"))
        except Exception as exc:
            print(err(f"Ollama {tag} failed: {exc}"))
    return done


# ─────────────────────────────────────────────────────────────────────────────
# Reporting
# ─────────────────────────────────────────────────────────────────────────────

def print_summary(
    *,
    whisper_paths: List[str],
    ct2_paths: List[str],
    gguf_paths: List[str],
    ollama_tags: List[str],
    registry: Dict[str, Any],
    env_updates: Dict[str, str],
    export_paths: Dict[str, str],
    postflight: Dict[str, Any],
) -> None:
    print(bold("\n\n╔══════════════════════════════════════════════════╗"))
    print(bold("║            Installation Summary                  ║"))
    print(bold("╚══════════════════════════════════════════════════╝\n"))

    def _section(title: str, items: List[str], prefix: str = "") -> None:
        print(f"  {bold(title)}")
        if items:
            for item in items:
                print(f"      {ok(prefix + item)}")
        else:
            print(f"      {dim('(none)')}")
        print()

    _section("faster-whisper:", whisper_paths)
    _section("CTranslate2:", ct2_paths)
    _section("GGUF:", gguf_paths)
    _section("Ollama:", ollama_tags, prefix="tag:")

    print(f"  {bold('Registry')} ")
    for k in ("root", "faster_whisper_dir", "translator_dir", "llama_cpp_model", "whisper_cpp_model"):
        print(f"      {k}={registry.get(k, '')}")
    print()

    print(f"  {bold('Updated env keys')} ")
    for k, v in env_updates.items():
        print(f"      {k}={v}")
    print()

    print(f"  {bold('Activation instructions')} ")
    print(f"      POSIX: source {export_paths['posix']}")
    print(f"      PowerShell: . {export_paths['powershell']}")
    print(f"      CMD: call {export_paths['cmd']}")
    print()

    if postflight["ok"]:
        print(ok("Postflight validation passed."))
    else:
        print(warn("Postflight validation found issues (partial success)."))
        for row in postflight["failed"]:
            print(f"      {err(row['name'])} {row['path']}")


def machine_summary(
    *,
    caps: RuntimeCapabilities,
    registry: Dict[str, Any],
    gpu: Dict[str, Any],
    preflight: Dict[str, Any],
    postflight: Dict[str, Any],
    env_updates: Dict[str, str],
) -> None:
    print("\n=== SETUP_SUMMARY_JSON ===")
    print(
        json.dumps(
            {
                "event": "setup_models_summary",
                "runtime": {
                    "platform": caps.platform,
                    "machine": caps.machine,
                    "python": caps.python_version,
                    "ci": caps.ci,
                    "container": caps.container,
                    "wsl": caps.wsl,
                    "virtualenv": caps.virtualenv,
                },
                "gpu": gpu,
                "registry": registry,
                "preflight": preflight,
                "postflight": postflight,
                "env_updates": env_updates,
            },
            ensure_ascii=False,
            indent=2,
        )
    )


# ─────────────────────────────────────────────────────────────────────────────
# Entrypoint
# ─────────────────────────────────────────────────────────────────────────────

def main() -> None:
    parser = argparse.ArgumentParser(
        prog="setup_models.py",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=textwrap.dedent(
            """\
            RecordNote model setup tool.

            Orchestrates faster-whisper, CTranslate2, GGUF and Ollama setup.
            """
        ),
    )
    parser.add_argument("--all", action="store_true", help="Use non-interactive defaults.")
    parser.add_argument("--non-interactive", action="store_true", help="Force non-interactive mode.")
    parser.add_argument("--update", action="store_true", help="Refresh even if assets already valid.")
    parser.add_argument("--list", action="store_true", help="List model catalogs and exit.")
    parser.add_argument("--skip-deps", action="store_true", help="Skip dependency check.")
    parser.add_argument("--model-root", default="", help="Explicit model root.")
    parser.add_argument("--llama-cpp-model-path", default="", help="Explicit LLAMA_CPP_MODEL_PATH.")
    parser.add_argument("--whisper-cpp-model-path", default="", help="Explicit STT_PROVIDER_WHISPER_CPP_MODEL.")
    parser.add_argument("--nllb-model-dir", default="", help="Explicit NLLB_MODEL_DIR.")
    parser.add_argument("--env-file", default=str((SCRIPT_DIR / ".env").resolve()), help="Path to .env to update.")
    args = parser.parse_args()

    if args.list:
        print(bold("faster-whisper:"))
        print("  " + ", ".join(WHISPER_MODELS.keys()))
        print(bold("CTranslate2:"))
        print("  " + ", ".join(CT2_MODELS.keys()))
        print(bold("GGUF:"))
        print("  " + ", ".join(GGUF_MODELS.keys()))
        print(bold("Ollama:"))
        print("  " + ", ".join(OLLAMA_MODELS.keys()))
        return

    caps = detect_runtime_capabilities()
    print(bold("\n── Runtime capability detection ─────────────────────────────────────────"))
    print(info(f"platform={caps.platform}, machine={caps.machine}, python={caps.python_version}"))
    print(info(f"ci={caps.ci}, container={caps.container}, wsl={caps.wsl}, virtualenv={caps.virtualenv}"))

    registry = discover_model_registry(
        model_root=args.model_root or None,
        llama_cpp_model_path=args.llama_cpp_model_path or None,
        whisper_cpp_model_path=args.whisper_cpp_model_path or None,
        nllb_model_dir=args.nllb_model_dir or None,
        caps=caps,
    )

    low_conf = [k for k, v in registry.get("confidence", {}).items() if v == "low"]
    if low_conf:
        print(warn(f"Auto-detection uncertain for: {', '.join(low_conf)}. Using conservative defaults."))

    if not args.skip_deps:
        print(bold("\n── Dependency check ─────────────────────────────────────────────────"))
        deps_ok, dep_warn = check_dependencies(caps, need_ollama=False)
        if not deps_ok:
            raise SetupError("Dependency check failed.", category="dependency")
        for w in dep_warn:
            print(warn(w))

    gpu = detect_gpu()
    print(bold("\n── Hardware detection ───────────────────────────────────────────────"))
    print(info(gpu["label"]))

    non_interactive = bool(args.all or args.non_interactive or caps.ci)
    if non_interactive:
        whisper_sel, ct2_sel, gguf_sel, ollama_sel = default_selections(gpu, caps)
    else:
        whisper_sel, ct2_sel, gguf_sel, ollama_sel = interactive_menu(caps)

    if not any([whisper_sel, ct2_sel, gguf_sel, ollama_sel]):
        print(warn("No models selected; nothing to do."))
        return

    preflight = preflight_validate(caps, registry, require_network=bool(whisper_sel or ct2_sel or gguf_sel))
    if not preflight["ok"]:
        print(err("Preflight failed:"))
        for issue in preflight["issues"]:
            print(f"  - [{issue['category']}] {issue['message']}")
        machine_summary(caps=caps, registry=registry, gpu=gpu, preflight=preflight, postflight={"ok": False, "checks": [], "failed": []}, env_updates={})
        raise SetupError("Preflight validation failed.", category="validation")

    print(bold("\n── Downloading / preparing models ───────────────────────────────────"))
    whisper_paths = download_whisper(whisper_sel, force=args.update)
    ct2_paths = download_ctranslate2(ct2_sel, force=args.update)
    gguf_paths = download_gguf(gguf_sel, force=args.update)
    ollama_done = pull_ollama_models(ollama_sel, caps=caps, force=args.update)

    if whisper_paths:
        registry["faster_whisper_dir"] = str(Path(whisper_paths[0]).resolve())
        n = Path(whisper_paths[0]).name
        if n.startswith("faster-whisper-"):
            registry["faster_whisper_model"] = n.replace("faster-whisper-", "")
    if ct2_paths:
        registry["translator_dir"] = str(Path(ct2_paths[0]).resolve())
    if gguf_paths:
        registry["llama_cpp_model"] = str(Path(gguf_paths[0]).resolve())

    env_updates = {
        "MODEL_REGISTRY_ROOT": registry["root"],
        "STT_PROVIDER_FASTER_WHISPER_MODEL_DIR": registry["faster_whisper_dir"],
        "STT_PROVIDER_FASTER_WHISPER_MODEL": registry["faster_whisper_model"],
        "NLLB_MODEL_DIR": registry["translator_dir"],
        "LLAMA_CPP_MODEL_PATH": registry.get("llama_cpp_model", ""),
        "STT_PROVIDER_WHISPER_CPP_MODEL": registry.get("whisper_cpp_model", ""),
    }

    env_file = Path(args.env_file).expanduser().resolve()
    upsert_env_file(env_file, env_updates)
    export_paths = write_export_scripts(SCRIPT_DIR, env_updates)

    postflight = postflight_validate(
        whisper_paths=whisper_paths,
        ct2_paths=ct2_paths,
        gguf_paths=gguf_paths,
        ollama_tags=ollama_done,
        registry=registry,
        caps=caps,
    )

    print_summary(
        whisper_paths=whisper_paths,
        ct2_paths=ct2_paths,
        gguf_paths=gguf_paths,
        ollama_tags=ollama_done,
        registry=registry,
        env_updates=env_updates,
        export_paths=export_paths,
        postflight=postflight,
    )

    machine_summary(
        caps=caps,
        registry=registry,
        gpu=gpu,
        preflight=preflight,
        postflight=postflight,
        env_updates=env_updates,
    )

    if not postflight["ok"]:
        raise SetupError("Completed with partial success; check postflight failures.", category="validation")


if __name__ == "__main__":
    try:
        main()
    except SetupError as exc:
        _json_log({"event": "setup_error", "category": exc.category, "message": str(exc)})
        print(err(f"Setup failed [{exc.category}]: {exc}"))
        sys.exit(1)
    except KeyboardInterrupt:
        _json_log({"event": "setup_error", "category": "runtime", "message": "Interrupted by user"})
        print(err("Interrupted by user."))
        sys.exit(130)