#!/bin/bash

# configure reads TITLE and DESCRIPTION
# shellcheck disable=SC2034
TITLE="AI servers"
# shellcheck disable=SC2034
DESCRIPTION="Use LLM servers on your network, e.g. Ollama"

# The servers you've added, one host:port per line. It exists once you've run
# this, even if you added none, so this isn't "to do" on machines without one.
AI_SERVERS_CONF="$HOME/.config/coredesktop/ai-servers"
# The sandbox config every sandboxed agent reads (see sandbox/agent)
SANDBOX_CONF="$HOME/.config/coredesktop/sandbox/base.json"
CLAUDE_SETTINGS="$HOME/.claude/settings.json"
OPENCODE_CONF="$HOME/.config/opencode/opencode.json"
source "$COREDESKTOP/configure.d/lib/mdns.sh"
source "$COREDESKTOP/configure.d/lib/json.sh"

# Ports LLM servers listen on by default, and who uses each. All of them answer
# OpenAI's API at /v1, which is how every agent here talks to them.
LLM_PORTS=(11434 1234 8080 8000)
declare -A LLM_PORT_NAMES=([11434]=Ollama [1234]="LM Studio" [8080]=llama.cpp [8000]=vLLM)

is_installed() {
  command -v avahi-browse
}

is_configured() {
  [ -f "$AI_SERVERS_CONF" ]
}

# The models a server at host:port offers, one per line; fails if it doesn't
# answer like an LLM server
server_models() {
  curl -sf -m 3 "http://$1/v1/models" | jq -er '.data[].id' 2> /dev/null
}

# What the server is, e.g. Ollama, for its name in the agents' settings
server_kind() {
  local port="${1##*:}"
  if curl -sf -m 3 "http://$1/api/version" > /dev/null; then
    echo Ollama
  else
    echo "${LLM_PORT_NAMES[$port]:-LLM server}"
  fi
}

# The LLM servers on the network: "host:port|models" per line. mDNS finds the
# machines, but LLM servers don't announce themselves on it, so each machine
# is asked on the usual ports.
find_servers() {
  local dir host port
  dir="$(mktemp -d)" || return 1
  for host in $(mdns_hosts); do
    for port in "${LLM_PORTS[@]}"; do
      (server_models "$host:$port" | wc -l > "$dir/$host:$port") &
    done
  done
  wait
  for host in "$dir"/*; do
    [ -f "$host" ] && [ "$(cat "$host")" -gt 0 ] && echo "$(basename "$host")|$(cat "$host")"
  done
  rm -rf "$dir"
}

# Let every sandboxed agent (opencode, Hermes, Cline) reach host:port, and
# only that port
sandbox_allow() {
  edit_json "$SANDBOX_CONF" '{}' --arg server "$1" \
    '.network.allowedDomains = ((.network.allowedDomains // []) - [$server] + [$server])'
}

sandbox_disallow() {
  [ -f "$SANDBOX_CONF" ] || return 0
  edit_json "$SANDBOX_CONF" '{}' --arg server "$1" \
    '.network.allowedDomains -= [$server]'
}

# Claude Code's sandbox only covers the commands it runs, e.g. curl, not its
# own connection to Anthropic. Its list replaces the defaults sandbox/setup.sh
# gives it, so if it has none yet, it starts from those.
claude_allow() {
  local defaults
  [ -f "$CLAUDE_SETTINGS" ] || return 0
  defaults="$(jq '.network.allowedDomains' "$COREDESKTOP/sandbox/config/base.json")" || return 1
  edit_json "$CLAUDE_SETTINGS" '{}' --arg server "$1" --argjson defaults "$defaults" \
    '.sandbox.network.allowedDomains = ((.sandbox.network.allowedDomains // $defaults) - [$server] + [$server])'
}

claude_disallow() {
  [ -f "$CLAUDE_SETTINGS" ] || return 0
  edit_json "$CLAUDE_SETTINGS" '{}' --arg server "$1" \
    'if .sandbox.network.allowedDomains then .sandbox.network.allowedDomains -= [$server] else . end'
}

# opencode's name for the server: llm.local:11434 -> llm-11434
opencode_provider() {
  local id="${1%%:*}"
  id="${id%.local}-${1##*:}"
  echo "${id//[^a-zA-Z0-9-]/-}" | tr '[:upper:]' '[:lower:]'
}

# Add the server to opencode as a provider with its models, and make $3 the
# model opencode starts with if you say so
opencode_add() {
  local server="$1" kind="$2" model="$3" id models
  id="$(opencode_provider "$server")"
  models="$(server_models "$server" | jq -R '{(.): {name: .}}' | jq -s 'add')" || return 1
  edit_json "$OPENCODE_CONF" '{"$schema": "https://opencode.ai/config.json"}' \
    --arg id "$id" --arg name "$kind on ${server%%:*}" --arg url "http://$server/v1" \
    --argjson models "$models" \
    '.provider[$id] = {npm: "@ai-sdk/openai-compatible", name: $name, options: {baseURL: $url}, models: $models}' \
    || return 1
  echo "opencode: added $kind on ${server%%:*} (choose its models with /models)."
  if confirm_yes "Start opencode with $model?"; then
    edit_json "$OPENCODE_CONF" '{}' --arg model "$id/$model" '.model = $model' || return 1
  fi
}

opencode_remove() {
  local id
  [ -f "$OPENCODE_CONF" ] || return 0
  id="$(opencode_provider "$1")"
  edit_json "$OPENCODE_CONF" '{}' --arg id "$id" \
    'del(.provider[$id]) | if (.model // "") | startswith($id + "/") then del(.model) else . end'
}

# Hermes uses one model at a time, so ask before replacing the one it has
hermes_use() {
  local server="$1" model="$2"
  confirm_yes "Use $model on ${server%%:*} for Hermes?" || return 0
  unsandboxed hermes config set model.provider custom > /dev/null \
    && unsandboxed hermes config set model.base_url "http://$server/v1" > /dev/null \
    && unsandboxed hermes config set model.default "$model" > /dev/null || return 1
  echo "Hermes: now uses $model."
}

# Hermes' gateway reads the sandbox config when it starts, so it can't reach a
# new server until it restarts. (A sandboxed opencode starts its own server
# each time, so it reads the config again then.)
restart_gateway() {
  if systemctl --user is-active -q hermes-gateway 2> /dev/null \
      && confirm_yes "Restart Hermes' gateway so it can reach the server?"; then
    systemctl --user restart hermes-gateway || return 1
  fi
}

saved_servers() {
  grep -v '^$' "$AI_SERVERS_CONF" 2> /dev/null
}

# Choose a model from the server and set the agents up to use it
add_server() {
  local server="$1" kind models model choice i
  mapfile -t models < <(server_models "$server")
  if [ ${#models[@]} -eq 0 ]; then
    echo "$server doesn't answer like an LLM server (http://$server/v1/models)."
    return 1
  fi
  kind="$(server_kind "$server")"

  echo
  echo "$kind on $server has these models:"
  for i in "${!models[@]}"; do
    printf '  %2d) %s\n' $((i + 1)) "${models[$i]}"
  done
  read -rp "Which should the agents use? [1] " choice
  choice="${choice:-1}"
  if ! [[ "$choice" =~ ^[0-9]+$ ]] || (( choice < 1 || choice > ${#models[@]} )); then
    echo "Choose a number from the list."
    return 1
  fi
  model="${models[$((choice - 1))]}"

  sandbox_allow "$server" || return 1
  claude_allow "$server" || return 1
  grep -qsxF "$server" "$AI_SERVERS_CONF" || echo "$server" >> "$AI_SERVERS_CONF"
  echo "The sandbox lets the agents reach $server."

  if [ -x /usr/bin/opencode ]; then
    opencode_add "$server" "$kind" "$model" || return 1
  fi
  if [ -x "$HOME/.local/bin/hermes" ]; then
    hermes_use "$server" "$model" || return 1
  fi
  restart_gateway
}

remove_server() {
  local server="$1" tmp
  sandbox_disallow "$server" || return 1
  claude_disallow "$server" || return 1
  opencode_remove "$server" || return 1
  tmp="$(mktemp)" || return 1
  grep -vxF "$server" "$AI_SERVERS_CONF" > "$tmp"
  mv -f "$tmp" "$AI_SERVERS_CONF"
  echo "The agents can no longer reach $server."
  if grep -qsF "$server" "$HOME/.hermes/config.yaml"; then
    echo "Hermes still uses it; choose another model provider with: configure hermes"
  fi
}

# Find servers on the network and add the one you choose
find_and_add() {
  local servers server models choice i
  echo "Looking for LLM servers on your network..."
  mapfile -t servers < <(find_servers)
  if [ ${#servers[@]} -eq 0 ]; then
    echo "No LLM servers found. A server has to accept connections from other machines:"
    echo "in Ollama's settings, turn on \"Expose Ollama to the network\" (or set OLLAMA_HOST=0.0.0.0);"
    echo "in LM Studio, turn on \"Serve on Local Network\". Or enter its address instead."
    return 0
  fi
  for i in "${!servers[@]}"; do
    IFS='|' read -r server models <<< "${servers[$i]}"
    printf '  %2d) %-28s %s, %s models\n' $((i + 1)) "$server" "$(server_kind "$server")" "$models"
  done
  read -rp "Which server? " choice
  if ! [[ "$choice" =~ ^[0-9]+$ ]] || (( choice < 1 || choice > ${#servers[@]} )); then
    echo "Choose a number from the list."
    return 1
  fi
  add_server "${servers[$((choice - 1))]%%|*}"
}

# host, host:port or a URL such as http://host:port/v1 -> host:port
enter_and_add() {
  local address
  read -rp "The server's address, e.g. llm.local or llm.local:11434: " address
  address="${address#*://}"
  address="${address%%/*}"
  [ -n "$address" ] || return 0
  [[ "$address" == *:* ]] || address="$address:11434"
  add_server "$address"
}

choose_and_remove() {
  local servers choice i
  mapfile -t servers < <(saved_servers)
  for i in "${!servers[@]}"; do
    printf '  %2d) %s\n' $((i + 1)) "${servers[$i]}"
  done
  read -rp "Which server? " choice
  if ! [[ "$choice" =~ ^[0-9]+$ ]] || (( choice < 1 || choice > ${#servers[@]} )); then
    echo "Choose a number from the list."
    return 1
  fi
  remove_server "${servers[$((choice - 1))]}"
}

run_configure() {
  local choice
  mkdir -p "$(dirname "$AI_SERVERS_CONF")" && touch "$AI_SERVERS_CONF" || return 1
  echo "Agents run in a sandbox that only lets them reach the servers you allow. This finds"
  echo "LLM servers on your network, such as Ollama on a Mac, allows them and sets up the agents."

  while true; do
    echo
    if [ -s "$AI_SERVERS_CONF" ]; then
      echo "The agents can reach: $(saved_servers | paste -sd' ')"
      echo
    fi
    echo "  1) Find servers on your network"
    echo "  2) Enter a server's address"
    [ -s "$AI_SERVERS_CONF" ] && echo "  3) Remove a server"
    echo "  q) Done"
    echo
    read -rp "Choose: " choice || return 0
    case "$choice" in
      1) find_and_add ;;
      2) enter_and_add ;;
      3) [ -s "$AI_SERVERS_CONF" ] && choose_and_remove ;;
      *) return 0 ;;
    esac
  done
}