#!/usr/bin/env bash
# Train a Mugi tool-calling LoRA adapter with mlx_lm.lora on the bundled dataset.
#
# Defaults:
#   - Base model: mlx_model from ~/.mugi/app_prefs.json (or MUGI_MODEL)
#   - Dataset:    training/mlx_mugi_agent_tools under repo root (or MUGI_DATA_DIR)
#   - Output:     ~/.mugi/mlx/adapters/mugi-agent-tools-<model-slug>
#   - Python:     ~/.mugi/mlx-venv/bin/python (Mugi's MLX venv)
#
# Usage:
#   ./scripts/train_mugi_tool_adapter.sh
#   ./scripts/train_mugi_tool_adapter.sh --dry-run
#   ./scripts/train_mugi_tool_adapter.sh --model /path/to/model --iters 2000 --num-layers 4
#   MUGI_MODEL=mlx-community/Qwen3.5-9B-6bit ./scripts/train_mugi_tool_adapter.sh
#
# In-app: Settings → LLMs → Adapters → Train tool-calling adapter for current model…
#
# After training:
#   Settings → LLMs → Adapters → Refresh → select the new mugi-agent-tools-* adapter.
#   Keep the same base model selected. Do not overwrite evolution_* adapters.
#
set -euo pipefail

REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
MLX_PYTHON="${MUGI_MLX_PYTHON:-${HOME}/.mugi/mlx-venv/bin/python}"

DRY_RUN=0
CLI_MODEL=""
CLI_ITERS=""
CLI_NUM_LAYERS=""
CLI_OVERWRITE=0
CLI_NO_APPLY=0

usage() {
  cat <<'EOF'
Usage: train_mugi_tool_adapter.sh [options]

Options:
  --model PATH         Override base MLX model (default: mlx_model in app_prefs.json)
  --iters N            Training iterations (default: model-size heuristic)
  --num-layers N       LoRA layers to train (default: model-size heuristic)
  --overwrite          Replace an existing adapter directory
  --no-apply           Do not set mlx_active_adapter when training completes
  --dry-run            Print the resolved plan via tool_adapter_trainer (no training)
  -h, --help           Show this help

Environment:
  MUGI_MODEL, MUGI_ITERS, MUGI_NUM_LAYERS, MUGI_MLX_PYTHON, MUGI_REPO_ROOT
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --model) CLI_MODEL="${2:-}"; shift 2 ;;
    --iters) CLI_ITERS="${2:-}"; shift 2 ;;
    --num-layers) CLI_NUM_LAYERS="${2:-}"; shift 2 ;;
    --overwrite) CLI_OVERWRITE=1; shift ;;
    --no-apply) CLI_NO_APPLY=1; shift ;;
    --dry-run) DRY_RUN=1; shift ;;
    -h|--help) usage; exit 0 ;;
    *) echo "train_mugi_tool_adapter: unknown option: $1" >&2; usage >&2; exit 2 ;;
  esac
done

if [[ ! -x "$MLX_PYTHON" ]]; then
  echo "train_mugi_tool_adapter: MLX venv python not found at ${MLX_PYTHON}" >&2
  echo "Install MLX from Mugi Settings → LLMs first, or set MUGI_MLX_PYTHON." >&2
  exit 1
fi

export MUGI_REPO_ROOT="${MUGI_REPO_ROOT:-$REPO_ROOT}"
export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:$PYTHONPATH}"

ARGS=()
[[ -n "${CLI_MODEL:-${MUGI_MODEL:-}}" ]] && ARGS+=(--model "${CLI_MODEL:-${MUGI_MODEL:-}}")
[[ -n "${CLI_ITERS:-${MUGI_ITERS:-}}" ]] && ARGS+=(--iters "${CLI_ITERS:-${MUGI_ITERS:-}}")
[[ -n "${CLI_NUM_LAYERS:-${MUGI_NUM_LAYERS:-}}" ]] && ARGS+=(--num-layers "${CLI_NUM_LAYERS:-${MUGI_NUM_LAYERS:-}}")
[[ "$CLI_OVERWRITE" -eq 1 ]] && ARGS+=(--overwrite)
[[ "$CLI_NO_APPLY" -eq 1 ]] && ARGS+=(--no-apply)

if [[ "$DRY_RUN" -eq 1 ]]; then
  "$MLX_PYTHON" - <<'PY'
import json
import os
import sys
sys.path.insert(0, os.environ["MUGI_REPO_ROOT"])
from tool_adapter_trainer import resolve_tool_adapter_plan
model = os.environ.get("MUGI_MODEL", "").strip() or None
plan = resolve_tool_adapter_plan(model)
print(json.dumps(plan, indent=2))
if not plan.get("ok"):
    sys.exit(1)
PY
  exit 0
fi

exec "$MLX_PYTHON" -m tool_adapter_trainer "${ARGS[@]}"
