#!/usr/bin/env bash
set -euo pipefail

cd "$(dirname "$0")/.."

export COPYFILE_DISABLE=1

if [[ -f .env ]]; then
  set -a
  # shellcheck disable=SC1091
  source .env
  set +a
fi

CUSTOM_LOCAL_LLM_ROOT="${CUSTOM_LOCAL_LLM_ROOT:-/Volumes/Zoe/custom-local-llm}"
TRAINING_DATA_DIR="${TRAINING_DATA_DIR:-$CUSTOM_LOCAL_LLM_ROOT/training-data}"
TRAINING_DATA_MANIFEST="${TRAINING_DATA_MANIFEST:-$CUSTOM_LOCAL_LLM_ROOT/manifest.llm.json}"
ENABLE_MIMIR_SDG="${ENABLE_MIMIR_SDG:-0}"
MIMIR_SDG_HOST="${MIMIR_SDG_HOST:-100.80.52.47}"
MIMIR_SDG_USER="${MIMIR_SDG_USER:-aaron-pressey}"
MIMIR_SDG_KEY="${MIMIR_SDG_KEY:-$HOME/.ssh/silma_orson_ed25519}"
MIMIR_SDG_MODEL_ALIAS="${MIMIR_SDG_MODEL_ALIAS:-nvidia-text}"
MIMIR_SDG_NUM_RECORDS="${MIMIR_SDG_NUM_RECORDS:-1000}"

if [[ ! -d "$CUSTOM_LOCAL_LLM_ROOT" ]]; then
  echo "missing custom local LLM root: $CUSTOM_LOCAL_LLM_ROOT" >&2
  exit 1
fi

mkdir -p "$TRAINING_DATA_DIR" "$CUSTOM_LOCAL_LLM_ROOT/logs" "$CUSTOM_LOCAL_LLM_ROOT/models" "$CUSTOM_LOCAL_LLM_ROOT/benchmarks"

if [[ ! -s "$TRAINING_DATA_MANIFEST" ]]; then
  echo "missing training data manifest: $TRAINING_DATA_MANIFEST" >&2
  exit 1
fi

cargo run -- dictionary generate
bridge_args=(
  scripts/generate_training_data.py
  --dictionary-dir dictionary/ \
  --output-dir "$TRAINING_DATA_DIR" \
  --manifest "$TRAINING_DATA_MANIFEST"
)

if [[ "$ENABLE_MIMIR_SDG" = "1" ]]; then
  if [[ ! -r "$MIMIR_SDG_KEY" ]]; then
    echo "missing Mimir SDG SSH key: $MIMIR_SDG_KEY" >&2
    exit 1
  fi
  bridge_args+=(
    --sdg-host "$MIMIR_SDG_HOST"
    --sdg-user "$MIMIR_SDG_USER"
    --sdg-key "$MIMIR_SDG_KEY"
    --sdg-model-alias "$MIMIR_SDG_MODEL_ALIAS"
    --num-records "$MIMIR_SDG_NUM_RECORDS"
  )
fi

python3 "${bridge_args[@]}"

python3 - <<'PY'
import os
from pathlib import Path

root = Path(os.environ.get("CUSTOM_LOCAL_LLM_ROOT", "/Volumes/Zoe/custom-local-llm"))
for path in root.rglob("._*"):
    if path.is_file():
        path.unlink()
PY