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

# Ollama tmux helper and model utilities
# Source this file or call functions directly: source scripts/ollama_tmux.sh

# Ensure environment variables exist
if [ -f "$HOME/.bashrc" ]; then
  set +u
  # shellcheck disable=SC1091
  source "$HOME/.bashrc" || true
  set -u
fi
: "${OLLAMA_MODELS:=$HOME/models/ollama}"

start_ollama_tmux() {
  session="ollama"
  # Use provided command or sensible default. Ollama uses env vars (OLLAMA_HOST)
  # Example: OLLAMA_HOST=0.0.0.0:11434 ollama serve
  default_cmd="${OLLAMA_CMD:-ollama serve}"
  cmd="${1:-$default_cmd}"
  echo "Using command: $cmd"
  if [ -n "${OLLAMA_HOST:-}" ]; then
    echo "OLLAMA_HOST is set to: $OLLAMA_HOST"
  fi

  if ! command -v tmux >/dev/null 2>&1; then
    echo "⚠️  tmux not installed. Install tmux and re-run."
    return 2
  fi

  # Create models dir if needed
  mkdir -p "$OLLAMA_MODELS"

  if tmux has-session -t "$session" 2>/dev/null; then
    echo "ℹ️  tmux session '$session' already running. Attach with: tmux attach -t $session"
    return 0
  fi

  echo "🔁 Starting ollama in tmux session '$session' with command: $cmd"
  tmux new-session -d -s "$session" "$cmd"
  sleep 0.5
  if tmux has-session -t "$session" 2>/dev/null; then
    echo "✅ Started. Attach: tmux attach -t $session"
  else
    echo "❌ Failed to start tmux session. Check logs or run the command directly: $cmd"
    return 1
  fi
}

pull_ollama_model() {
  if [ "$#" -lt 1 ]; then
    echo "Usage: pull_ollama_model <model-ref>"
    return 2
  fi
  if ! command -v ollama >/dev/null 2>&1; then
    echo "⚠️  ollama CLI not found. Install first."
    return 3
  fi
  model="$1"
  echo "⬇️  Pulling model: $model"
  ollama pull "$model"
}

import_local_model() {
  if [ "$#" -lt 1 ]; then
    echo "Usage: import_local_model <path-to-model-dir-or-file>"
    return 2
  fi
  src="$1"
  if [ ! -e "$src" ]; then
    echo "⚠️  Source not found: $src"
    return 3
  fi
  mkdir -p "$OLLAMA_MODELS"
  dest="$OLLAMA_MODELS/$(basename "$src")"
  ln -sfn "$src" "$dest"
  echo "🔗 Linked $src → $dest"
  echo "Note: depending on your ollama installation you may need to reindex or use the ollama CLI to register the model."
}

stop_ollama_tmux() {
  session="ollama"
  if tmux has-session -t "$session" 2>/dev/null; then
    tmux kill-session -t "$session"
    echo "🛑 Stopped tmux session '$session'"
  else
    echo "ℹ️  No tmux session named '$session' running"
  fi
}

print_usage() {
  cat <<EOF
Usage: $(basename "$0") <command> [args]

Commands:
  start [<cmd>]        Start ollama in tmux (optional command overrides default)
  stop                 Stop the tmux session named 'ollama'
  pull <model-ref>     Pull a model using the ollama CLI
  import <path>        Import/link a local model directory/file into OLLAMA_MODELS
  help                 Show this help

If you want to use these as shell functions, source this file instead:
  source scripts/ollama_tmux.sh
EOF
}

if [ "${BASH_SOURCE[0]}" == "${0}" ]; then
  # Script executed directly; provide a simple CLI wrapper
  cmd="${1:-}"
  case "$cmd" in
    start)
      # shift and pass remaining args as the command
      shift || true
      start_ollama_tmux "$@"
      exit $?
      ;;
    stop)
      stop_ollama_tmux
      exit $?
      ;;
    pull)
      shift || true
      pull_ollama_model "$@"
      exit $?
      ;;
    import)
      shift || true
      import_local_model "$@"
      exit $?
      ;;
    help|--help|-h|"")
      print_usage
      exit 0
      ;;
    *)
      echo "Unknown command: $cmd"
      print_usage
      exit 2
      ;;
  esac
else
  echo "Loaded ollama tmux helper. Functions: start_ollama_tmux, stop_ollama_tmux, pull_ollama_model, import_local_model"
fi