interactive-mlx-setup.zsh

#!/bin/zsh
#
# Interactive MLX Server for macOS, version v0.11.
# This script checks python3 and MLX package status, installs or updates mlx/mlx-lm,
# manages MLX-compatible Hugging Face model downloads, starts or stops the MLX server,
# and provides cleanup and uninstall options from one menu-driven interface.
# AI generation note: generated and edited with the Codex App.
# Requested model attribution label: GPT-5.4 Extra High Model.

set -u
set -o pipefail
setopt null_glob

APP_NAME="Interactive MLX Server"
APP_VERSION="v0.11"
SCRIPT_DIR="${0:A:h}"
CONFIG_DIR="${HOME}/.interactive-mlx-setup"
CONFIG_FILE="${CONFIG_DIR}/config.env"
STATE_DIR="${CONFIG_DIR}/state"
PID_FILE="${STATE_DIR}/mlx-server.pid"
LAST_LOG_FILE="${STATE_DIR}/last-server-log.txt"
ORIGINAL_MODEL_DIR="${HOME}/.cache/huggingface"
DEFAULT_MODEL_DIR="${HOME}/MLX-Models"
MODEL_DIR="${ORIGINAL_MODEL_DIR}"
MODEL_REPO_ID_MARKER=".interactive-mlx-repo-id"
SELECTED_MODEL_PATH=""
SERVER_HOST="127.0.0.1"
SERVER_PORT="8080"
SERVER_MODEL_STATUS=""
SERVER_MODEL_TYPE=""
SERVER_MODEL_EFFECTIVE_TYPE=""
SERVER_MODEL_ARCHITECTURES=""
SERVER_MODEL_REASON=""
READ_WITH_QUIT_RESULT=""

use_local_config_dir() {
  CONFIG_DIR="${SCRIPT_DIR}/.interactive-mlx-setup"
  CONFIG_FILE="${CONFIG_DIR}/config.env"
  STATE_DIR="${CONFIG_DIR}/state"
  PID_FILE="${STATE_DIR}/mlx-server.pid"
  LAST_LOG_FILE="${STATE_DIR}/last-server-log.txt"
  mkdir -p "${CONFIG_DIR}" "${STATE_DIR}"
}

if ! mkdir -p "${CONFIG_DIR}" "${STATE_DIR}" 2>/dev/null; then
  use_local_config_dir
elif ! touch "${STATE_DIR}/.write-test" 2>/dev/null; then
  use_local_config_dir
else
  rm -f "${STATE_DIR}/.write-test"
fi

if [[ -f "${CONFIG_FILE}" ]]; then
  source "${CONFIG_FILE}"
fi

USE_COLOR=0
if [[ -t 1 && -z "${NO_COLOR:-}" ]]; then
  USE_COLOR=1
fi

COLOR_RESET=""
COLOR_HEADER_FRAME=""
COLOR_HEADER_TITLE=""
COLOR_INFO=""
COLOR_WARN=""
COLOR_ERROR=""
COLOR_DESC=""
COLOR_MENU_KEY=""
COLOR_MENU_TEXT=""
COLOR_PROMPT=""

if (( USE_COLOR )); then
  COLOR_RESET=$'\033[0m'
  COLOR_HEADER_FRAME=$'\033[38;5;39m'
  COLOR_HEADER_TITLE=$'\033[1;38;5;45m'
  COLOR_INFO=$'\033[38;5;78m'
  COLOR_WARN=$'\033[38;5;220m'
  COLOR_ERROR=$'\033[1;38;5;203m'
  COLOR_DESC=$'\033[38;5;252m'
  COLOR_MENU_KEY=$'\033[1;38;5;81m'
  COLOR_MENU_TEXT=$'\033[38;5;159m'
  COLOR_PROMPT=$'\033[1;38;5;213m'
fi

print_colored_line() {
  local color="${1}"
  local text="${2}"
  if [[ -n "${color}" ]]; then
    printf '%b%s%b\n' "${color}" "${text}" "${COLOR_RESET}"
  else
    print -- "${text}"
  fi
}

expand_path() {
  local raw_path="${1}"
  if [[ -z "${raw_path}" ]]; then
    return 1
  fi
  print -r -- "${~raw_path}"
}

save_config() {
  local config_payload

  config_payload="$(cat <<EOF
MODEL_DIR=$(printf '%q' "${MODEL_DIR}")
SELECTED_MODEL_PATH=$(printf '%q' "${SELECTED_MODEL_PATH}")
SERVER_HOST=$(printf '%q' "${SERVER_HOST}")
SERVER_PORT=$(printf '%q' "${SERVER_PORT}")
EOF
)"

  if ! print -r -- "${config_payload}" > "${CONFIG_FILE}" 2>/dev/null; then
    use_local_config_dir
    print -r -- "${config_payload}" > "${CONFIG_FILE}" 2>/dev/null || {
      print_warn "Could not save config to ${CONFIG_FILE}."
      return 1
    }
  fi
}

exit_script() {
  local answer=""
  local server_pid=""

  if server_is_running; then
    server_pid="$(<"${PID_FILE}")"
    print
    print_warn "The MLX server is still running."
    while true; do
      if [[ -n "${COLOR_PROMPT}" ]]; then
        printf '%bDo you want to stop the running MLX server as well? [y/n/q]: %b' "${COLOR_PROMPT}" "${COLOR_RESET}"
      else
        printf "Do you want to stop the running MLX server as well? [y/n/q]: "
      fi

      if ! IFS= read -r answer; then
        break
      fi

      case "${answer:l}" in
        y|yes)
          if [[ -n "${server_pid}" ]]; then
            kill "${server_pid}" >/dev/null 2>&1 || true
            sleep 2
            if kill -0 "${server_pid}" >/dev/null 2>&1; then
              kill -9 "${server_pid}" >/dev/null 2>&1 || true
            fi
            rm -f "${PID_FILE}"
            print_info "Server stopped."
          fi
          break
          ;;
        n|no|q|quit|exit|stop|"")
          break
          ;;
        *)
          print_warn "Please choose y, n, or q."
          ;;
      esac
    done
  fi

  print
  print_colored_line "${COLOR_WARN}" "Stopping ${APP_NAME} ${APP_VERSION}. Goodbye."
  exit 0
}

print_header() {
  print
  print_colored_line "${COLOR_HEADER_FRAME}" "============================================================"
  print_colored_line "${COLOR_HEADER_TITLE}" "${1}"
  print_colored_line "${COLOR_HEADER_FRAME}" "============================================================"
}

print_info() {
  if (( USE_COLOR )); then
    printf '%b[INFO]%b %s\n' "${COLOR_INFO}" "${COLOR_RESET}" "${1}"
  else
    print "[INFO] ${1}"
  fi
}

print_warn() {
  if (( USE_COLOR )); then
    printf '%b[WARN]%b %s\n' "${COLOR_WARN}" "${COLOR_RESET}" "${1}"
  else
    print "[WARN] ${1}"
  fi
}

print_error() {
  if (( USE_COLOR )); then
    printf '%b[ERROR]%b %s\n' "${COLOR_ERROR}" "${COLOR_RESET}" "${1}"
  else
    print "[ERROR] ${1}"
  fi
}

print_desc() {
  print_colored_line "${COLOR_DESC}" "${1}"
}

print_menu_item() {
  local key="${1}"
  local text="${2}"
  if (( USE_COLOR )); then
    printf '%b%s%b %b%s%b\n' "${COLOR_MENU_KEY}" "${key}" "${COLOR_RESET}" "${COLOR_MENU_TEXT}" "${text}" "${COLOR_RESET}"
  else
    print "${key} ${text}"
  fi
}

read_with_quit() {
  local prompt="${1}"
  local answer=""
  while true; do
    if [[ -n "${COLOR_PROMPT}" ]]; then
      printf '%b%s%b' "${COLOR_PROMPT}" "${prompt}" "${COLOR_RESET}" >&2
    else
      printf "%s" "${prompt}" >&2
    fi
    if ! IFS= read -r answer; then
      exit_script
    fi
    if [[ "${answer:l}" == "q" || "${answer:l}" == "quit" || "${answer:l}" == "exit" || "${answer:l}" == "stop" ]]; then
      exit_script
    fi
    READ_WITH_QUIT_RESULT="${answer}"
    return 0
  done
}

prompt_yes_no_quit() {
  local prompt="${1}"
  local answer=""
  while true; do
    if [[ -n "${COLOR_PROMPT}" ]]; then
      printf '%b%s [y/n/q]: %b' "${COLOR_PROMPT}" "${prompt}" "${COLOR_RESET}"
    else
      printf "%s [y/n/q]: " "${prompt}"
    fi
    IFS= read -r answer
    case "${answer:l}" in
      y|yes) return 0 ;;
      n|no) return 1 ;;
      q|quit|exit|stop) exit_script ;;
      *) print_warn "Please choose y, n, or q." ;;
    esac
  done
}

pause_for_menu() {
  read_with_quit "Press Enter to return to the home menu, or q to quit: "
}

require_python3() {
  if ! command -v python3 >/dev/null 2>&1; then
    print_error "python3 is not installed or not in PATH."
    print "Please install Python 3 for macOS first, then run this script again."
    return 1
  fi

  local python_version
  python_version="$(python3 --version 2>/dev/null)"
  print_info "Python status: ${python_version}"
  return 0
}

get_installed_version() {
  local package_name="${1}"
  python3 - "${package_name}" <<'PY'
import importlib.metadata
import sys

package_name = sys.argv[1]

try:
    print(importlib.metadata.version(package_name))
except importlib.metadata.PackageNotFoundError:
    print("")
PY
}

get_latest_version() {
  local package_name="${1}"
  python3 - "${package_name}" <<'PY'
import json
import sys
import urllib.error
import urllib.request

package_name = sys.argv[1]
url = f"https://pypi.org/pypi/{package_name}/json"

try:
    with urllib.request.urlopen(url, timeout=10) as response:
        data = json.load(response)
    print(data["info"]["version"])
except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, KeyError, ValueError):
    print("")
PY
}

show_package_status() {
  local package_name="${1}"
  local package_label="${2}"
  local installed_version

  installed_version="$(get_installed_version "${package_name}")"

  if [[ -n "${installed_version}" ]]; then
    print_info "${package_label}: installed (version ${installed_version})"
  else
    print_warn "${package_label}: not installed"
  fi
}

run_pip_install() {
  local -a packages=("$@")

  print_info "Running: python3 -m pip install --upgrade ${packages[*]}"
  if python3 -m pip install --upgrade "${packages[@]}"; then
    return 0
  fi

  print_warn "Standard install failed. Trying a user-level install."
  python3 -m pip install --user --upgrade "${packages[@]}"
}

uninstall_mlx_stack() {
  local remove_models_choice=1
  local remove_state_choice=1

  print_header "Uninstall MLX"
  print_desc "This removes the MLX framework, mlx-lm server tools, and optionally your downloaded models and saved script state."

  if server_is_running; then
    print_warn "The MLX server is currently running."
    if prompt_yes_no_quit "Stop the running server before uninstalling?"; then
      stop_running_server || return 1
    else
      print_info "Uninstall cancelled."
      return 1
    fi
  fi

  if ! prompt_yes_no_quit "Do you want to uninstall MLX and the MLX server tools now?"; then
    print_info "Uninstall skipped."
    return 1
  fi

  if python3 -m pip uninstall -y mlx mlx-lm huggingface_hub; then
    print_info "Package uninstall completed."
  else
    print_error "Package uninstall failed."
    return 1
  fi

  if prompt_yes_no_quit "Do you also want to remove all downloaded models in ${MODEL_DIR}?"; then
    remove_models_choice=0
  else
    remove_models_choice=1
  fi

  if (( remove_models_choice == 0 )); then
    local -a detected_model_dirs
    local detected_model_dir

    detected_model_dirs=("${(@f)$(get_model_dirs)}")
    if (( ${#detected_model_dirs[@]} == 1 )) && [[ -z "${detected_model_dirs[1]}" ]]; then
      detected_model_dirs=()
    fi

    if (( ${#detected_model_dirs[@]} == 0 )); then
      print_warn "No detected model folders were found to remove."
    else
      for detected_model_dir in "${detected_model_dirs[@]}"; do
        rm -rf -- "${detected_model_dir}"
      done
      print_info "Removed detected model folders from: ${MODEL_DIR}"
      SELECTED_MODEL_PATH=""
    fi
  fi

  if prompt_yes_no_quit "Do you want to remove saved settings, logs, and PID files too?"; then
    remove_state_choice=0
  else
    remove_state_choice=1
  fi

  if (( remove_state_choice == 0 )); then
    if [[ -n "${CONFIG_DIR}" && -d "${CONFIG_DIR}" && "${CONFIG_DIR}" != "/" && "${CONFIG_DIR}" != "${HOME}" ]]; then
      rm -rf -- "${CONFIG_DIR}"
      print_info "Removed saved script state: ${CONFIG_DIR}"
    else
      print_warn "Skipped state removal because the config directory path was not safe to delete."
    fi
  else
    SELECTED_MODEL_PATH=""
    save_config || true
  fi

  print_info "Current package status after uninstall:"
  show_package_status "mlx" "MLX"
  show_package_status "mlx-lm" "MLX LM"
  show_package_status "huggingface_hub" "Hugging Face Hub"
  return 0
}

ensure_python_package_installed() {
  local package_name="${1}"
  local package_label="${2}"

  if [[ -n "$(get_installed_version "${package_name}")" ]]; then
    return 0
  fi

  print_warn "${package_label} is required but is not installed."
  if ! prompt_yes_no_quit "Install ${package_label} now?"; then
    print_info "Skipped installing ${package_label}."
    return 1
  fi

  if run_pip_install "${package_name}"; then
    show_package_status "${package_name}" "${package_label}"
    return 0
  fi

  print_error "Failed to install ${package_label}."
  return 1
}

install_mlx_stack() {
  local latest_mlx latest_mlx_lm

  latest_mlx="$(get_latest_version "mlx")"
  latest_mlx_lm="$(get_latest_version "mlx-lm")"

  print_header "Install MLX"
  if [[ -n "${latest_mlx}" ]]; then
    print_info "Latest MLX version on PyPI: ${latest_mlx}"
  else
    print_warn "Could not confirm the latest MLX version from PyPI."
  fi

  if [[ -n "${latest_mlx_lm}" ]]; then
    print_info "Latest mlx-lm version on PyPI: ${latest_mlx_lm}"
  else
    print_warn "Could not confirm the latest mlx-lm version from PyPI."
  fi

  if ! prompt_yes_no_quit "Install the MLX framework and MLX LM tools now?"; then
    print_info "Install skipped."
    return 1
  fi

  if ! run_pip_install "mlx" "mlx-lm" "huggingface_hub"; then
    print_error "Installation failed."
    return 1
  fi

  print_info "Installation completed. Verifying package status."
  show_package_status "mlx" "MLX"
  show_package_status "mlx-lm" "MLX LM"
  show_package_status "huggingface_hub" "Hugging Face Hub"
  return 0
}

check_for_updates() {
  local installed_mlx latest_mlx installed_mlx_lm latest_mlx_lm
  local need_update=0

  installed_mlx="$(get_installed_version "mlx")"
  latest_mlx="$(get_latest_version "mlx")"
  installed_mlx_lm="$(get_installed_version "mlx-lm")"
  latest_mlx_lm="$(get_latest_version "mlx-lm")"

  print_header "Update Check"

  if [[ -n "${installed_mlx}" ]]; then
    print_info "Installed MLX version: ${installed_mlx}"
  else
    print_warn "MLX is not installed."
  fi

  if [[ -n "${latest_mlx}" ]]; then
    print_info "Latest MLX version: ${latest_mlx}"
    if [[ -n "${installed_mlx}" && "${installed_mlx}" != "${latest_mlx}" ]]; then
      need_update=1
    fi
  else
    print_warn "Could not confirm the latest MLX version."
  fi

  if [[ -n "${installed_mlx_lm}" ]]; then
    print_info "Installed mlx-lm version: ${installed_mlx_lm}"
  else
    print_warn "mlx-lm is not installed."
  fi

  if [[ -n "${latest_mlx_lm}" ]]; then
    print_info "Latest mlx-lm version: ${latest_mlx_lm}"
    if [[ -n "${installed_mlx_lm}" && "${installed_mlx_lm}" != "${latest_mlx_lm}" ]]; then
      need_update=1
    fi
  else
    print_warn "Could not confirm the latest mlx-lm version."
  fi

  if [[ -z "${installed_mlx}" || -z "${installed_mlx_lm}" ]]; then
    print_warn "One or more required packages are missing."
    if prompt_yes_no_quit "Install missing packages now?"; then
      install_mlx_stack
    else
      print_info "Skipped package installation."
    fi
    return 0
  fi

  if (( need_update == 0 )); then
    print_info "MLX and mlx-lm are already up to date."
    return 0
  fi

  if prompt_yes_no_quit "An update is available. Do you want to update now?"; then
    if run_pip_install "mlx" "mlx-lm" "huggingface_hub"; then
      print_info "Update completed."
      show_package_status "mlx" "MLX"
      show_package_status "mlx-lm" "MLX LM"
    else
      print_error "Update failed."
    fi
  else
    print_info "Update skipped. Ignoring for now."
  fi
}

configure_model_dir() {
  local choice prompt_value expanded_path

  print_header "Model Folder"
  print_info "Current model folder: ${MODEL_DIR}"
  print_desc "Original Hugging Face cache location: ${ORIGINAL_MODEL_DIR}"
  print_desc "Suggested dedicated model folder: ${DEFAULT_MODEL_DIR}"
  print_desc "Choose Y to use the suggested folder, N to keep the original cache folder, C for a custom path, or press Enter to keep the current setting."

  while true; do
    read_with_quit "Model folder choice [y/n/c/Enter]: "
    choice="${READ_WITH_QUIT_RESULT}"

    case "${choice:l}" in
      "")
        break
        ;;
      y|yes)
        MODEL_DIR="${DEFAULT_MODEL_DIR}"
        break
        ;;
      n|no)
        MODEL_DIR="${ORIGINAL_MODEL_DIR}"
        break
        ;;
      c|custom)
        read_with_quit "Enter custom model folder path: "
        prompt_value="${READ_WITH_QUIT_RESULT}"
        if [[ -z "${prompt_value}" ]]; then
          print_warn "Custom path cannot be empty."
          continue
        fi
        expanded_path="$(expand_path "${prompt_value}")"
        MODEL_DIR="${expanded_path}"
        break
        ;;
      *)
        print_warn "Please choose Y, N, C, or press Enter to keep the current setting."
        ;;
    esac
  done

  if ! mkdir -p "${MODEL_DIR}" 2>/dev/null; then
    print_error "Could not create or access model folder: ${MODEL_DIR}"
    return 1
  fi

  save_config || true
  print_info "Model folder set to: ${MODEL_DIR}"
}

get_model_dirs() {
  python3 - "${MODEL_DIR}" <<'PY'
import os
import sys
from pathlib import Path

root = Path(sys.argv[1]).expanduser()
if not root.exists():
    raise SystemExit(0)

max_depth = 6
candidates = set()

for dirpath, dirnames, filenames in os.walk(root):
    current_path = Path(dirpath)
    try:
        relative_parts = current_path.relative_to(root).parts
    except ValueError:
        continue

    depth = len(relative_parts)
    if depth > max_depth:
        dirnames[:] = []
        continue

    if "config.json" in filenames and any(name.endswith(".safetensors") for name in filenames):
        candidates.add(str(current_path))
        dirnames[:] = []
        continue

    dirnames[:] = [
        name
        for name in dirnames
        if name not in {".git", "__pycache__", ".cache"}
    ]

for candidate in sorted(candidates):
    print(candidate)
PY
}

get_model_display_name() {
  local model_path="${1}"
  python3 - "${model_path}" "${MODEL_REPO_ID_MARKER}" <<'PY'
import sys
from pathlib import Path

model_path = Path(sys.argv[1])
marker_name = sys.argv[2]

marker_path = model_path / marker_name
if marker_path.exists():
    try:
        repo_id = marker_path.read_text().strip()
        if repo_id:
            print(repo_id)
            raise SystemExit(0)
    except Exception:
        pass

for parent in [model_path, *model_path.parents]:
    if parent.name.startswith("models--"):
        slug = parent.name[len("models--"):]
        print(slug.replace("--", "/"))
        raise SystemExit(0)

name = model_path.name
if "--" in name:
    print(name.replace("--", "/", 1))
else:
    print(name)
PY
}

normalize_hf_repo_id() {
  local raw_value="${1}"
  python3 - "${raw_value}" <<'PY'
import sys
from urllib.parse import urlparse

raw_value = sys.argv[1].strip()

if not raw_value:
    print("")
    raise SystemExit(0)

if raw_value.startswith(("https://", "http://")):
    parsed = urlparse(raw_value)
    if parsed.netloc not in {"huggingface.co", "www.huggingface.co"}:
        print("")
        raise SystemExit(0)

    parts = [part for part in parsed.path.split("/") if part]
    if len(parts) >= 2:
        print(f"{parts[0]}/{parts[1]}")
    else:
        print("")
else:
    parts = [part for part in raw_value.split("/") if part]
    if len(parts) == 2:
        print(f"{parts[0]}/{parts[1]}")
    else:
        print("")
PY
}

model_validation_summary() {
  local model_path="${1}"
  python3 - "${model_path}" <<'PY'
from pathlib import Path
import sys

model_path = Path(sys.argv[1])
safetensors = list(model_path.glob("*.safetensors"))
config = model_path / "config.json"

print(f"Path: {model_path}")
print(f"Safetensors files found: {len(safetensors)}")
print(f"config.json present: {'yes' if config.exists() else 'no'}")
PY
}

get_model_server_status() {
  local model_path="${1}"
  python3 - "${model_path}" <<'PY'
import ast
import importlib.metadata
import json
import sys
from pathlib import Path

model_path = Path(sys.argv[1])
status = "unsupported"
model_type = ""
effective_model_type = ""
architectures = ""
reason = ""

config_path = model_path / "config.json"
if not config_path.exists():
    status = "invalid"
    reason = "config.json not found."
else:
    try:
        config = json.loads(config_path.read_text())
    except Exception as exc:
        status = "invalid"
        reason = f"Unable to read config.json: {exc}"
    else:
        model_type = str(config.get("model_type", "") or "")
        architectures_value = config.get("architectures", [])
        if isinstance(architectures_value, list):
            architectures = ", ".join(str(item) for item in architectures_value)
        else:
            architectures = str(architectures_value)

        if not model_type:
            status = "invalid"
            reason = "config.json does not contain model_type."
        else:
            try:
                dist = importlib.metadata.distribution("mlx-lm")
            except importlib.metadata.PackageNotFoundError:
                status = "invalid"
                reason = "mlx-lm is not installed."
            else:
                package_root = Path(dist.locate_file(""))
                utils_path = package_root / "mlx_lm" / "utils.py"
                model_remapping = {}

                try:
                    tree = ast.parse(utils_path.read_text())
                    for node in tree.body:
                        if isinstance(node, ast.Assign):
                            for target in node.targets:
                                if isinstance(target, ast.Name) and target.id == "MODEL_REMAPPING":
                                    model_remapping = ast.literal_eval(node.value)
                except Exception:
                    model_remapping = {}

                effective_model_type = str(model_remapping.get(model_type, model_type))
                model_module_path = package_root / "mlx_lm" / "models" / f"{effective_model_type}.py"

                if model_module_path.exists():
                    status = "supported"
                    reason = "Supported by the installed mlx-lm server."
                else:
                    status = "unsupported"
                    reason = f"Model type {model_type} is not supported by the installed mlx-lm server."
                    hint_text = f"{model_type} {architectures}".lower()
                    if any(token in hint_text for token in ("bert", "roberta", "xlm", "embed", "embedding")):
                        reason += " It looks like an embedding model, not a text-generation LLM."

fields = [status, model_type, effective_model_type, architectures, reason]
print("\t".join(value.replace("\t", " ").replace("\n", " ").strip() for value in fields))
PY
}

load_model_server_status() {
  local model_path="${1}"
  local status_line

  SERVER_MODEL_STATUS=""
  SERVER_MODEL_TYPE=""
  SERVER_MODEL_EFFECTIVE_TYPE=""
  SERVER_MODEL_ARCHITECTURES=""
  SERVER_MODEL_REASON=""

  status_line="$(get_model_server_status "${model_path}")"
  IFS=$'\t' read -r SERVER_MODEL_STATUS SERVER_MODEL_TYPE SERVER_MODEL_EFFECTIVE_TYPE SERVER_MODEL_ARCHITECTURES SERVER_MODEL_REASON <<< "${status_line}"
}

print_model_server_status() {
  local model_path="${1}"

  load_model_server_status "${model_path}"
  if [[ "${SERVER_MODEL_STATUS}" == "supported" ]]; then
    print "Server compatibility: supported"
    if [[ -n "${SERVER_MODEL_TYPE}" ]]; then
      print "Model type: ${SERVER_MODEL_TYPE}"
    fi
    if [[ -n "${SERVER_MODEL_EFFECTIVE_TYPE}" && "${SERVER_MODEL_EFFECTIVE_TYPE}" != "${SERVER_MODEL_TYPE}" ]]; then
      print "Resolved model type: ${SERVER_MODEL_EFFECTIVE_TYPE}"
    fi
  else
    print "Server compatibility: not supported"
    if [[ -n "${SERVER_MODEL_TYPE}" ]]; then
      print "Model type: ${SERVER_MODEL_TYPE}"
    fi
    if [[ -n "${SERVER_MODEL_ARCHITECTURES}" ]]; then
      print "Architectures: ${SERVER_MODEL_ARCHITECTURES}"
    fi
    if [[ -n "${SERVER_MODEL_REASON}" ]]; then
      print "Reason: ${SERVER_MODEL_REASON}"
    fi
  fi
}

download_single_model() {
  local repo_input repo_id repo_slug destination_path

  require_python3 || return 1
  ensure_python_package_installed "huggingface_hub" "Hugging Face Hub" || return 1

  print
  print "Enter a Hugging Face URL or repo id for an MLX-compatible model."
  print "Example: https://huggingface.co/mlx-community/Llama-3.2-3B-Instruct-4bit"
  print "Note: the model must be an MLX model to work with MLX LM."

  read_with_quit "Hugging Face URL or repo id (q to quit): "
  repo_input="${READ_WITH_QUIT_RESULT}"
  repo_id="$(normalize_hf_repo_id "${repo_input}")"

  if [[ -z "${repo_id}" ]]; then
    print_error "That is not a valid Hugging Face URL or owner/repo id."
    return 1
  fi

  repo_slug="${repo_id//\//--}"
  destination_path="${MODEL_DIR}/${repo_slug}"

  print_info "Downloading ${repo_id}"
  print_info "Destination: ${destination_path}"
  mkdir -p "${destination_path}"

  if ! python3 - "${repo_id}" "${destination_path}" <<'PY'
import sys

repo_id = sys.argv[1]
destination_path = sys.argv[2]

try:
    from huggingface_hub import snapshot_download
except Exception as exc:
    print(f"Failed to import huggingface_hub: {exc}")
    raise SystemExit(1)

try:
    snapshot_download(repo_id=repo_id, local_dir=destination_path)
except Exception as exc:
    print(f"Download failed: {exc}")
    raise SystemExit(1)

print(f"Download complete: {destination_path}")
PY
  then
    print_error "Model download failed."
    return 1
  fi

  print_info "Download completed successfully."
  print -r -- "${repo_id}" > "${destination_path}/${MODEL_REPO_ID_MARKER}" 2>/dev/null || true
  model_validation_summary "${destination_path}"
  print_model_server_status "${destination_path}"
  if [[ "${SERVER_MODEL_STATUS}" != "supported" ]]; then
    print_warn "This model was downloaded, but it is not usable with the MLX text-generation server."
    print_warn "Choose an instruction/chat MLX model for server use."
  fi
  return 0
}

download_models_flow() {
  local download_again=0

  while true; do
    if ! download_single_model; then
      if ! prompt_yes_no_quit "Do you want to try another model download?"; then
        return 0
      fi
      continue
    fi

    if prompt_yes_no_quit "Do you want to download another model?"; then
      download_again=1
    else
      download_again=0
    fi

    if (( download_again == 0 )); then
      return 0
    fi
  done
}

show_downloaded_models() {
  local -a model_dirs
  local dir index model_label

  print_header "Downloaded Models"
  model_dirs=("${(@f)$(get_model_dirs)}")
  if (( ${#model_dirs[@]} == 1 )) && [[ -z "${model_dirs[1]}" ]]; then
    model_dirs=()
  fi

  if (( ${#model_dirs[@]} == 0 )); then
    print_warn "No models downloaded in ${MODEL_DIR}"
    if prompt_yes_no_quit "Do you want to download a model now?"; then
      download_models_flow
    fi
    return 0
  fi

  index=1
  for dir in "${model_dirs[@]}"; do
    model_label="$(get_model_display_name "${dir}")"
    load_model_server_status "${dir}"
    print "${index}. ${model_label}"
    print "   Path: ${dir}"
    if [[ "${SERVER_MODEL_STATUS}" == "supported" ]]; then
      print "   MLX server: supported (${SERVER_MODEL_EFFECTIVE_TYPE:-${SERVER_MODEL_TYPE}})"
    else
      print "   MLX server: not supported"
      if [[ -n "${SERVER_MODEL_REASON}" ]]; then
        print "   Reason: ${SERVER_MODEL_REASON}"
      fi
    fi
    (( index++ ))
  done
}

choose_model_for_server() {
  local -a model_dirs
  local answer selected_index
  local dir index model_label

  while true; do
    print_header "Choose Model For MLX Server"
    model_dirs=("${(@f)$(get_model_dirs)}")
    if (( ${#model_dirs[@]} == 1 )) && [[ -z "${model_dirs[1]}" ]]; then
      model_dirs=()
    fi

    if (( ${#model_dirs[@]} == 0 )); then
      print_warn "No models are available."
      if prompt_yes_no_quit "Do you want to download a model now?"; then
        download_models_flow
        continue
      fi
      return 1
    fi

    index=1
    for dir in "${model_dirs[@]}"; do
      model_label="$(get_model_display_name "${dir}")"
      load_model_server_status "${dir}"
      print "${index}. ${model_label}"
      print "   Path: ${dir}"
      if [[ "${SERVER_MODEL_STATUS}" == "supported" ]]; then
        print "   MLX server: supported (${SERVER_MODEL_EFFECTIVE_TYPE:-${SERVER_MODEL_TYPE}})"
      else
        print "   MLX server: not supported"
      fi
      (( index++ ))
    done
    print_menu_item "d." "Download another model"
    print_menu_item "b." "Back to home menu"
    print_menu_item "q." "Quit"

    read_with_quit "Choose a model number, d, b, or q: "
    answer="${READ_WITH_QUIT_RESULT}"

    case "${answer:l}" in
      b) return 1 ;;
      d)
        download_models_flow
        continue
        ;;
    esac

    if [[ "${answer}" == <-> ]]; then
      selected_index="${answer}"
      if (( selected_index >= 1 && selected_index <= ${#model_dirs[@]} )); then
        load_model_server_status "${model_dirs[selected_index]}"
        if [[ "${SERVER_MODEL_STATUS}" != "supported" ]]; then
          print_error "That model cannot be used with the MLX server."
          print_model_server_status "${model_dirs[selected_index]}"
          print_warn "Please choose a server-compatible text-generation model."
          continue
        fi

        SELECTED_MODEL_PATH="${model_dirs[selected_index]}"
        save_config
        print_info "Selected model: $(get_model_display_name "${SELECTED_MODEL_PATH}")"
        print_info "Selected model path: ${SELECTED_MODEL_PATH}"
        return 0
      fi
    fi

    print_warn "Invalid selection."
  done
}

server_is_running() {
  if [[ ! -f "${PID_FILE}" ]]; then
    return 1
  fi

  local server_pid
  server_pid="$(<"${PID_FILE}")"

  if [[ -z "${server_pid}" ]]; then
    return 1
  fi

  if kill -0 "${server_pid}" >/dev/null 2>&1; then
    return 0
  fi

  rm -f "${PID_FILE}"
  return 1
}

show_runtime_stats() {
  local server_pid="${1}"

  print_header "Runtime Status"
  print_info "Server PID: ${server_pid}"
  print
  print "Server process usage:"
  ps -p "${server_pid}" -o pid,%cpu,rss,etime,command 2>/dev/null || print_warn "ps output was not available."
  print
  print "System memory snapshot:"
  python3 - <<'PY' 2>/dev/null || print_warn "vm_stat output was not available."
import re
import subprocess


def fmt_bytes(num_bytes: int) -> str:
    units = ["B", "KiB", "MiB", "GiB", "TiB"]
    value = float(num_bytes)
    for unit in units:
        if value < 1024.0 or unit == units[-1]:
            if unit == "B":
                return f"{int(value)} {unit}"
            return f"{value:.2f} {unit}"
        value /= 1024.0


vm_output = subprocess.check_output(["vm_stat"], text=True)
page_size_match = re.search(r"page size of (\d+) bytes", vm_output)
page_size = int(page_size_match.group(1)) if page_size_match else 4096

keys = [
    ("Pages free", "Free"),
    ("Pages active", "Active"),
    ("Pages inactive", "Inactive"),
    ("Pages wired down", "Wired"),
    ("Pages occupied by compressor", "Compressed"),
]

values = {}
for line in vm_output.splitlines():
    if ":" not in line:
        continue
    key, raw_value = line.split(":", 1)
    key = key.strip()
    digits = re.sub(r"[^0-9]", "", raw_value)
    if digits:
        values[key] = int(digits)

try:
    total_ram = int(subprocess.check_output(["sysctl", "-n", "hw.memsize"], text=True).strip())
except Exception:
    total_ram = None

if total_ram is not None:
    print(f"Total RAM: {fmt_bytes(total_ram)}")

for raw_key, label in keys:
    pages = values.get(raw_key)
    if pages is None:
        continue
    byte_value = pages * page_size
    print(f"{label}: {fmt_bytes(byte_value)} ({pages:,} pages)")
PY
  print
  print "GPU memory / utilization snapshot (best effort on macOS):"
  ioreg -r -d 1 -w 0 -c IOAccelerator | grep -i 'Device Utilization %\|Renderer Utilization %\|Alloc system memory\|In use system memory' 2>/dev/null | head -n 1 || print_warn "GPU stats were not available from ioreg."
}

remove_downloaded_model() {
  local -a model_dirs
  local answer selected_index model_path
  local dir index model_label

  if server_is_running; then
    print_warn "Stop the running MLX server before removing downloaded models."
    return 1
  fi

  while true; do
    print_header "Remove Downloaded Model"
    model_dirs=("${(@f)$(get_model_dirs)}")
    if (( ${#model_dirs[@]} == 1 )) && [[ -z "${model_dirs[1]}" ]]; then
      model_dirs=()
    fi

    if (( ${#model_dirs[@]} == 0 )); then
      print_warn "No downloaded models are available to remove."
      return 0
    fi

    index=1
    for dir in "${model_dirs[@]}"; do
      model_label="$(get_model_display_name "${dir}")"
      print "${index}. ${model_label}"
      print "   Path: ${dir}"
      (( index++ ))
    done
    print_menu_item "b." "Back to home menu"
    print_menu_item "q." "Quit"

    read_with_quit "Choose a model number to remove, b, or q: "
    answer="${READ_WITH_QUIT_RESULT}"

    case "${answer:l}" in
      b) return 0 ;;
    esac

    if [[ "${answer}" == <-> ]]; then
      selected_index="${answer}"
      if (( selected_index >= 1 && selected_index <= ${#model_dirs[@]} )); then
        model_path="${model_dirs[selected_index]}"
        print_warn "This will permanently remove the model folder:"
        print_warn "${model_path}"
        if ! prompt_yes_no_quit "Do you want to delete this downloaded model?"; then
          print_info "Model removal skipped."
          continue
        fi

        if rm -rf -- "${model_path}"; then
          print_info "Removed model: $(get_model_display_name "${model_path}")"
          if [[ "${SELECTED_MODEL_PATH}" == "${model_path}" ]]; then
            SELECTED_MODEL_PATH=""
            save_config || true
            print_info "Cleared the selected server model because it was deleted."
          fi
          return 0
        fi

        print_error "Failed to remove model: ${model_path}"
        return 1
      fi
    fi

    print_warn "Invalid selection."
  done
}

show_server_details() {
  local server_pid log_file

  if ! server_is_running; then
    print_warn "No MLX server is currently running."
    return 1
  fi

  server_pid="$(<"${PID_FILE}")"
  log_file=""
  if [[ -f "${LAST_LOG_FILE}" ]]; then
    log_file="$(<"${LAST_LOG_FILE}")"
  fi

  print_header "MLX Server Details"
  print_info "Model: ${SELECTED_MODEL_PATH:-Not selected}"
  print_info "Host: ${SERVER_HOST}"
  print_info "Port: ${SERVER_PORT}"
  print_info "PID: ${server_pid}"
  print_info "API Base URL: http://${SERVER_HOST}:${SERVER_PORT}/v1"
  if [[ -n "${log_file}" ]]; then
    print_info "Log file: ${log_file}"
  fi

  show_runtime_stats "${server_pid}"

  if [[ -n "${log_file}" && -f "${log_file}" ]]; then
    print
    print "Recent server log lines:"
    tail -n 15 "${log_file}" 2>/dev/null || true
  fi
}

stop_running_server() {
  local server_pid

  if ! server_is_running; then
    print_warn "No MLX server is running."
    return 0
  fi

  server_pid="$(<"${PID_FILE}")"
  if ! prompt_yes_no_quit "A server is already running with PID ${server_pid}. Do you want to stop it?"; then
    return 1
  fi

  kill "${server_pid}" >/dev/null 2>&1 || true
  sleep 2
  if kill -0 "${server_pid}" >/dev/null 2>&1; then
    print_warn "The server is still running. Trying a stronger stop."
    kill -9 "${server_pid}" >/dev/null 2>&1 || true
  fi

  rm -f "${PID_FILE}"
  print_info "Server stopped."
  return 0
}

start_mlx_server() {
  local confirm_run log_file server_pid startup_wait max_wait

  require_python3 || return 1
  ensure_python_package_installed "mlx" "MLX" || return 1
  ensure_python_package_installed "mlx-lm" "MLX LM" || return 1

  if [[ -z "${SELECTED_MODEL_PATH}" || ! -d "${SELECTED_MODEL_PATH}" ]]; then
    print_warn "No model is currently selected for the server."
    if ! choose_model_for_server; then
      return 1
    fi
  fi

  load_model_server_status "${SELECTED_MODEL_PATH}"
  if [[ "${SERVER_MODEL_STATUS}" != "supported" ]]; then
    print_error "The selected model is not compatible with the MLX server."
    print_model_server_status "${SELECTED_MODEL_PATH}"
    if prompt_yes_no_quit "Do you want to choose another model now?"; then
      choose_model_for_server || return 1
    else
      return 1
    fi
  fi

  if server_is_running; then
    print_warn "An MLX server is already running."
    show_server_details
    if ! stop_running_server; then
      print_info "Keeping the existing server running."
      return 1
    fi
  fi

  print_header "Start MLX Server"
  print_info "Selected model: $(get_model_display_name "${SELECTED_MODEL_PATH}")"
  print_info "Selected model path: ${SELECTED_MODEL_PATH}"
  print_info "Host: ${SERVER_HOST}"
  print_info "Port: ${SERVER_PORT}"

  confirm_run=1
  if ! prompt_yes_no_quit "Ready to start the MLX server with this model?"; then
    confirm_run=0
  fi

  if (( confirm_run == 0 )); then
    print_info "Returning to the home page menu."
    return 1
  fi

  log_file="${STATE_DIR}/mlx-server-$(date +%Y%m%d-%H%M%S).log"
  print -r -- "${log_file}" > "${LAST_LOG_FILE}"

  print_info "Starting server in the background."
  print_info "Command: python3 -m mlx_lm server --model \"${SELECTED_MODEL_PATH}\" --host ${SERVER_HOST} --port ${SERVER_PORT}"
  nohup python3 -m mlx_lm server --model "${SELECTED_MODEL_PATH}" --host "${SERVER_HOST}" --port "${SERVER_PORT}" > "${log_file}" 2>&1 &
  server_pid=$!
  print -r -- "${server_pid}" > "${PID_FILE}"

  startup_wait=0
  max_wait=30
  while (( startup_wait < max_wait )); do
    if ! kill -0 "${server_pid}" >/dev/null 2>&1; then
      print_error "The server exited during startup. Check the log file below."
      print_error "${log_file}"
      tail -n 20 "${log_file}" 2>/dev/null || true
      rm -f "${PID_FILE}"
      return 1
    fi

    if grep -q "Starting httpd at" "${log_file}" 2>/dev/null; then
      break
    fi

    sleep 1
    (( startup_wait++ ))
  done

  print_header "Server Started"
  print_info "MLX server status: running"
  print_info "Model: $(get_model_display_name "${SELECTED_MODEL_PATH}")"
  print_info "Model path: ${SELECTED_MODEL_PATH}"
  print_info "Host: ${SERVER_HOST}"
  print_info "Port: ${SERVER_PORT}"
  print_info "PID: ${server_pid}"
  print_info "API Base URL: http://${SERVER_HOST}:${SERVER_PORT}/v1"
  print_info "Log file: ${log_file}"
  print_info "Showing best-effort CPU/GPU and memory status after startup."

  show_runtime_stats "${server_pid}"

  print
  print "The MLX server is running in the background."
  print "Use the home menu to view details or stop it."
  return 0
}

startup_checks() {
  local installed_mlx

  print_header "${APP_NAME} ${APP_VERSION}"
  print_desc "This script checks Python, installs or updates MLX when needed,"
  print_desc "manages MLX-compatible Hugging Face model downloads, and starts an MLX server."
  print_desc "You can quit at any prompt by typing q."

  if ! require_python3; then
    return 1
  fi

  show_package_status "mlx" "MLX"
  show_package_status "mlx-lm" "MLX LM"

  installed_mlx="$(get_installed_version "mlx")"
  if [[ -z "${installed_mlx}" ]]; then
    install_mlx_stack || true
  else
    check_for_updates
  fi

  configure_model_dir
  show_downloaded_models

  if [[ -z "${SELECTED_MODEL_PATH}" || ! -d "${SELECTED_MODEL_PATH}" ]]; then
    if prompt_yes_no_quit "Do you want to choose a model now for the MLX server?"; then
      choose_model_for_server || true
    fi
  fi

  if [[ -n "${SELECTED_MODEL_PATH}" && -d "${SELECTED_MODEL_PATH}" ]]; then
    if prompt_yes_no_quit "Do you want to start the MLX server now?"; then
      start_mlx_server || true
    else
      print_info "Returning to the home page menu."
    fi
  fi
}

home_menu() {
  local answer

  while true; do
    print_header "Home Menu"
    print_menu_item "1." "Show Python and MLX status"
    print_menu_item "2." "Install or update MLX"
    print_menu_item "3." "Set model folder"
    print_menu_item "4." "Show downloaded models"
    print_menu_item "5." "Download MLX model from Hugging Face"
    print_menu_item "6." "Choose model for MLX server"
    print_menu_item "7." "Start MLX server"
    print_menu_item "8." "Show running server details"
    print_menu_item "9." "Stop running server"
    print_menu_item "10." "Remove downloaded model"
    print_menu_item "11." "Completely uninstall MLX and server"
    print_menu_item "q." "Quit"

    read_with_quit "Choose an option: "
    answer="${READ_WITH_QUIT_RESULT}"

    case "${answer}" in
      1)
        print_header "Python and MLX Status"
        require_python3 || true
        show_package_status "mlx" "MLX"
        show_package_status "mlx-lm" "MLX LM"
        show_package_status "huggingface_hub" "Hugging Face Hub"
        pause_for_menu
        ;;
      2)
        check_for_updates
        pause_for_menu
        ;;
      3)
        configure_model_dir
        pause_for_menu
        ;;
      4)
        show_downloaded_models
        pause_for_menu
        ;;
      5)
        download_models_flow
        pause_for_menu
        ;;
      6)
        choose_model_for_server || true
        pause_for_menu
        ;;
      7)
        start_mlx_server || true
        pause_for_menu
        ;;
      8)
        show_server_details || true
        pause_for_menu
        ;;
      9)
        stop_running_server || true
        pause_for_menu
        ;;
      10)
        remove_downloaded_model || true
        pause_for_menu
        ;;
      11)
        uninstall_mlx_stack || true
        pause_for_menu
        ;;
      *)
        print_warn "Invalid option."
        ;;
    esac
  done
}

main() {
  startup_checks || exit 1
  save_config
  home_menu
}

main "$@"