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

SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
# shellcheck disable=SC1091
. "$SCRIPT_DIR/common.sh"

load_runtime_env

mkdir -p "$STATE_DIR/run" "$STATE_DIR/logs"

ensure_transformerlab_default_user() {
  local helper_python="${TRANSFORMERLAB_DIR}/envs/transformerlab/bin/python"

  if [ ! -x "$helper_python" ]; then
    return 0
  fi

  if [ -z "${TRANSFORMERLAB_DEFAULT_USER_EMAIL:-}" ] || [ -z "${TRANSFORMERLAB_DEFAULT_USER_PASSWORD:-}" ]; then
    return 0
  fi

  "$helper_python" "$SCRIPT_DIR/ensure_transformerlab_user.py" \
    --transformerlab-dir "$TRANSFORMERLAB_DIR" \
    --email "$TRANSFORMERLAB_DEFAULT_USER_EMAIL" \
    --password "$TRANSFORMERLAB_DEFAULT_USER_PASSWORD" \
    --first-name "${TRANSFORMERLAB_DEFAULT_USER_FIRST_NAME:-Student}" \
    --last-name "${TRANSFORMERLAB_DEFAULT_USER_LAST_NAME:-}" >>"$STATE_DIR/logs/transformerlab_default_user.log" 2>&1 || true
}

resolve_targets() {
  if [ $# -eq 0 ]; then
    echo "No target specified." >&2
    exit 1
  fi

  case "$1" in
    core)
      printf '%s\n' "ollama" "open-webui"
      ;;
    all)
      service_list
      ;;
    *)
      printf '%s\n' "$@"
      ;;
  esac
}

has_live_pid() {
  local service=$1
  local pid_file
  pid_file=$(service_pid_file "$service")

  if [ -f "$pid_file" ]; then
    local pid
    pid=$(cat "$pid_file")
    if kill -0 "$pid" >/dev/null 2>&1; then
      return 0
    fi
  fi

  return 1
}

is_running() {
  local service=$1

  has_live_pid "$service" || service_ready "$service"
}

service_ready() {
  local service=$1

  case "$service" in
    ollama)
      curl -fsS "$(service_url "$service")/api/tags" >/dev/null 2>&1
      ;;
    transformerlab)
      curl -fsS "$(service_url "$service")/healthz" >/dev/null 2>&1
      ;;
    promptfoo)
      curl -fsS "$(service_url "$service")/health" >/dev/null 2>&1
      ;;
    open-webui|chunkviz|embedding-atlas|unsloth|wiki)
      curl -fsS "$(service_url "$service")" >/dev/null 2>&1
      ;;
    *)
      return 1
      ;;
  esac
}

start_one() {
  local service=$1
  local cmd
  local log_file
  local pid_file
  local attempt
  local pid_grace_attempts=5

  if has_live_pid "$service"; then
    if [ "$service" = "transformerlab" ]; then
      ensure_transformerlab_default_user
    fi
    echo "$service already running"
    return 0
  fi

  if service_ready "$service"; then
    if [ "$service" = "transformerlab" ]; then
      ensure_transformerlab_default_user
    fi
    echo "$service already available"
    return 0
  fi

  case "$service" in
    open-webui)
      start_one ollama
      ;;
    transformerlab)
      if command -v python3 >/dev/null 2>&1; then
        python3 "$SCRIPT_DIR/repair_transformerlab_plugin_supports.py" \
          --transformerlab-dir "$TRANSFORMERLAB_DIR" \
          --plugin "fastchat_server" \
          --required-support "chat" \
          --required-support "completion" \
          --required-support "visualize_model" \
          --required-support "model_layers" \
          --required-support "rag" \
          --required-support "tools" \
          --required-support "template" \
          --required-support "embeddings" \
          --required-support "tokenize" \
          --required-support "logprobs" \
          --required-support "batched" >>"$STATE_DIR/logs/transformerlab_plugin_supports.log" 2>&1 || true
      fi
      ;;
    *)
      ;;
  esac

  cmd=$(service_command "$service")
  log_file=$(service_log_file "$service")
  pid_file=$(service_pid_file "$service")

  if [ "$service" = "ollama" ]; then
    env \
      OLLAMA_HOST="${COURSEWARE_BIND_HOST}:${COURSEWARE_OLLAMA_PORT}" \
      OLLAMA_MODELS="$OLLAMA_MODELS_DIR" \
      "$OLLAMA_BIN" serve </dev/null >>"$log_file" 2>&1 &
  elif command -v setsid >/dev/null 2>&1; then
    nohup setsid bash -lc "$cmd" </dev/null >>"$log_file" 2>&1 &
  else
    nohup bash -lc "$cmd" </dev/null >>"$log_file" 2>&1 &
  fi
  echo $! >"$pid_file"

  for attempt in $(seq 1 60); do
    if service_ready "$service"; then
      if [ "$service" = "transformerlab" ]; then
        ensure_transformerlab_default_user
      fi
      echo "started $service"
      return 0
    fi

    if ! has_live_pid "$service"; then
      if [ "$attempt" -ge "$pid_grace_attempts" ]; then
        rm -f "$pid_file"
        echo "failed to start $service; check $log_file" >&2
        exit 1
      fi
    fi

    sleep 1
  done

  echo "$service did not become ready in time; check $log_file" >&2
  exit 1
}

stop_one() {
  local service=$1
  local pid_file

  pid_file=$(service_pid_file "$service")

  if [ ! -f "$pid_file" ]; then
    echo "$service not running"
    return 0
  fi

  local pid
  pid=$(cat "$pid_file")
  if kill -0 "$pid" >/dev/null 2>&1; then
    kill "$pid" >/dev/null 2>&1 || true
    sleep 2
    if kill -0 "$pid" >/dev/null 2>&1; then
      kill -9 "$pid" >/dev/null 2>&1 || true
    fi
  fi

  rm -f "$pid_file"
  echo "stopped $service"
}

status_one() {
  local service=$1

  if service_ready "$service"; then
    printf 'RUNNING  %-15s %s\n' "$service" "$(service_url "$service")"
  elif has_live_pid "$service"; then
    printf 'STARTING %-15s %s\n' "$service" "$(service_url "$service")"
  else
    printf 'STOPPED  %-15s %s\n' "$service" "$(service_url "$service")"
  fi
}

urls() {
  cat <<EOF
Ollama API:      $(service_url ollama)
Open WebUI:      $(service_url open-webui)
TransformerLab:  $(service_url transformerlab)
ChunkViz:        $(service_url chunkviz)
Embedding Atlas: $(service_url embedding-atlas)
Unsloth Studio:  $(service_url unsloth)
Promptfoo CLI:   $PROMPTFOO_BIN
Promptfoo UI:    $(service_url promptfoo)
Wiki:            $(service_url wiki)
Kiln app:        ${KILN_LAUNCH_PATH:-not installed}
EOF
}

open_kiln() {
  local host_os

  host_os=$(uname -s)
  if [ "$host_os" = "Darwin" ] && [ -d "$KILN_MAC_APP" ]; then
    open "$KILN_MAC_APP"
    return 0
  fi

  if [ -x "$KILN_LINUX_BIN" ]; then
    nohup "$KILN_LINUX_BIN" >/dev/null 2>&1 &
    echo "started Kiln from $KILN_LINUX_BIN"
    return 0
  fi

  echo "Kiln is not installed." >&2
  exit 1
}

show_logs() {
  local service=$1
  local log_file

  log_file=$(service_log_file "$service")
  if [ ! -f "$log_file" ]; then
    echo "No log file for $service" >&2
    exit 1
  fi

  tail -n 80 "$log_file"
}

main() {
  local cmd=${1:-}
  shift || true

  ensure_runtime_env

  case "$cmd" in
    start)
      while IFS= read -r service; do
        start_one "$service"
      done < <(resolve_targets "$@")
      ;;
    stop)
      while IFS= read -r service; do
        stop_one "$service"
      done < <(resolve_targets "$@")
      ;;
    status)
      if [ $# -eq 0 ]; then
        set -- all
      fi
      while IFS= read -r service; do
        status_one "$service"
      done < <(resolve_targets "$@")
      ;;
    urls)
      urls
      ;;
    open)
      if [ "${1:-}" != "kiln" ]; then
        echo "Only 'open kiln' is supported." >&2
        exit 1
      fi
      open_kiln
      ;;
    logs)
      if [ $# -ne 1 ]; then
        echo "Usage: ./labctl logs <service>" >&2
        exit 1
      fi
      show_logs "$1"
      ;;
    *)
      echo "Unknown command: $cmd" >&2
      exit 1
      ;;
  esac
}

main "$@"