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

SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
PROJECT_ROOT=$(cd "$SCRIPT_DIR/.." && pwd)

# Load config/environment defaults if available
if [[ -f "$SCRIPT_DIR/knoe_cfg.sh" ]]; then
  # shellcheck disable=SC1090
  source "$SCRIPT_DIR/knoe_cfg.sh"
fi

action="${1:-scan}"
shift || true

DEFAULT_PORT="11434"
port="${OLLAMA_SERVER_PORT:-$DEFAULT_PORT}"
timeout="${OLLAMA_SCAN_TIMEOUT:-2}"
max_hosts="${OLLAMA_SCAN_MAX_HOSTS:-64}"
hosts_file=""
explicit_host=""
explicit_model=""

usage() {
  cat <<EOF
Usage: $0 [scan|configure] [options]

Options:
  --host <host>         Explicit host (optionally host:port)
  --port <port>         Default port (default: ${DEFAULT_PORT})
  --hosts-file <path>   Hosts file to scan (default: KNOE_CONF/hosts.txt or etc/hosts.txt)
  --timeout <seconds>   Curl timeout (default: ${timeout})
  --max-hosts <n>       Max hosts to probe (default: ${max_hosts})
  --model <name>        Preferred model (configure action)

Env:
  OLLAMA_HOST, OLLAMA_SERVER_HOST, OLLAMA_SERVER_PORT, OLLAMA_CANDIDATES
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --host)
      explicit_host="${2:-}"
      shift 2
      ;;
    --port)
      port="${2:-}"
      shift 2
      ;;
    --hosts-file)
      hosts_file="${2:-}"
      shift 2
      ;;
    --timeout)
      timeout="${2:-}"
      shift 2
      ;;
    --max-hosts)
      max_hosts="${2:-}"
      shift 2
      ;;
    --model)
      explicit_model="${2:-}"
      shift 2
      ;;
    --help|-h)
      usage
      exit 0
      ;;
    *)
      shift
      ;;
  esac
done

if [[ -z "$port" ]]; then
  port="$DEFAULT_PORT"
fi

normalize_host() {
  local raw="$1"
  raw="${raw#http://}"
  raw="${raw#https://}"
  raw="${raw%%/*}"
  printf '%s' "$raw"
}

split_host_port() {
  local raw host port_val
  raw=$(normalize_host "$1")
  host="$raw"
  port_val=""
  if [[ "$raw" == *":"* ]]; then
    host="${raw%%:*}"
    port_val="${raw##*:}"
  fi
  printf '%s\t%s\n' "$host" "$port_val"
}

seen="|"
pairs=()

add_pair() {
  local host="$1"
  local port_val="$2"
  host="$(printf '%s' "$host" | tr -d '[:space:]')"
  if [[ -z "$host" ]]; then
    return
  fi
  if [[ -z "$port_val" ]]; then
    port_val="$port"
  fi
  local key="${host}:${port_val}"
  if [[ "$seen" == *"|$key|"* ]]; then
    return
  fi
  seen="${seen}${key}|"
  pairs+=("$key")
}

add_candidates() {
  local list="$1"
  local entry host port_val
  for entry in $(printf '%s' "$list" | tr ',;' '  '); do
    if [[ -z "$entry" ]]; then
      continue
    fi
    IFS=$'\t' read -r host port_val < <(split_host_port "$entry")
    add_pair "$host" "$port_val"
  done
}

if [[ -n "$explicit_host" ]]; then
  IFS=$'\t' read -r host port_val < <(split_host_port "$explicit_host")
  add_pair "$host" "$port_val"
fi

if [[ -n "${OLLAMA_HOST:-}" ]]; then
  IFS=$'\t' read -r host port_val < <(split_host_port "$OLLAMA_HOST")
  add_pair "$host" "$port_val"
fi

if [[ -n "${OLLAMA_SERVER_HOST:-}" ]]; then
  add_pair "$OLLAMA_SERVER_HOST" "$port"
fi

if [[ -n "${OLLAMA_CANDIDATES:-}" ]]; then
  add_candidates "$OLLAMA_CANDIDATES"
fi

if [[ -z "$hosts_file" ]]; then
  if [[ -n "${KNOE_CONF:-}" && -f "$KNOE_CONF/hosts.txt" ]]; then
    hosts_file="$KNOE_CONF/hosts.txt"
  elif [[ -f "$SCRIPT_DIR/hosts.txt" ]]; then
    hosts_file="$SCRIPT_DIR/hosts.txt"
  elif [[ -f "$PROJECT_ROOT/etc/hosts.txt" ]]; then
    hosts_file="$PROJECT_ROOT/etc/hosts.txt"
  fi
fi

if [[ -n "$hosts_file" && -f "$hosts_file" ]]; then
  while IFS= read -r line || [[ -n "$line" ]]; do
    line="${line%%#*}"
    line="$(printf '%s' "$line" | xargs)"
    [[ -z "$line" ]] && continue
    read -r first second _rest <<<"$line"
    [[ -z "$first" ]] && continue
    if [[ "$first" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
      add_pair "$second" "$port"
      add_pair "$first" "$port"
    else
      add_pair "$first" "$port"
      if [[ -n "$second" ]]; then
        add_pair "$second" "$port"
      fi
    fi
  done <"$hosts_file"
fi

if [[ -f "/etc/hosts" ]]; then
  while IFS= read -r line || [[ -n "$line" ]]; do
    line="${line%%#*}"
    line="$(printf '%s' "$line" | xargs)"
    [[ -z "$line" ]] && continue
    read -r first second _rest <<<"$line"
    [[ -z "$first" || -z "$second" ]] && continue
    if [[ "$first" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
      add_pair "$second" "$port"
      add_pair "$first" "$port"
    fi
  done <"/etc/hosts"
fi

scan_pairs() {
  local count=0
  local total=${#pairs[@]}
  echo "[scan] Probing $total candidate(s)..." >&2
  local pair host port_val url json models
  for pair in "${pairs[@]}"; do
    if [[ "$count" -ge "$max_hosts" ]]; then
      break
    fi
    host="${pair%%:*}"
    port_val="${pair##*:}"
    echo "[scan] Probing ${host}:${port_val} ..." >&2
    if [[ "$host" == *":"* && "$host" != "["* ]]; then
      url="http://[${host}]:${port_val}/api/tags"
    else
      url="http://${host}:${port_val}/api/tags"
    fi
    if ! json=$(curl -fsS --connect-timeout "$timeout" --max-time "$timeout" "$url" 2>/dev/null); then
      echo "[scan]   no response" >&2
      continue
    fi
    # Use a separate script file for the Python parser to keep stdin free for the JSON data
    local py_script="${SCRIPT_DIR}/.ollama_parser.py"
    if [[ ! -f "$py_script" ]]; then
      cat > "$py_script" <<'PY'
import json
import sys
try:
    data = sys.stdin.read().strip()
    if not data:
        sys.exit(1)
    payload = json.loads(data)
    if isinstance(payload, list):
        models_list = payload
    elif isinstance(payload, dict):
        models_list = payload.get("models", payload.get("tags", payload.get("models_list", [])))
    else:
        sys.exit(1)
    
    if not isinstance(models_list, list):
        sys.exit(1)
        
    names = []
    for entry in models_list:
        if isinstance(entry, dict):
            name = entry.get("name") or entry.get("model")
        elif isinstance(entry, str):
            name = entry
        else:
            continue
        if name:
            names.append(name)
    print(",".join(names))
except Exception:
    sys.exit(1)
PY
    fi

    if ! models=$(printf '%s' "$json" | python3 "$py_script"); then
      continue
    fi
    echo "[scan]   found: ${host}:${port_val} models=${models}" >&2
    printf '%s\t%s\t%s\n' "$host" "$port_val" "$models"
    count=$((count + 1))
  done
  echo "[scan] Done. ${count} server(s) detected." >&2
}

case "$action" in
  scan)
    scan_pairs
    ;;
  configure)
    first_line="$(scan_pairs | head -n 1)"
    if [[ -z "$first_line" ]]; then
      exit 1
    fi
    host="$(printf '%s' "$first_line" | awk -F'\t' '{print $1}')"
    port_val="$(printf '%s' "$first_line" | awk -F'\t' '{print $2}')"
    models="$(printf '%s' "$first_line" | awk -F'\t' '{print $3}')"
    if [[ -z "$explicit_model" ]]; then
      explicit_model="${models%%,*}"
    fi
    echo "OLLAMA_SERVER_HOST=${host}"
    echo "OLLAMA_SERVER_PORT=${port_val}"
    echo "OLLAMA_HOST=http://${host}:${port_val}"
    if [[ -n "$explicit_model" ]]; then
      echo "OLLAMA_MODEL=${explicit_model}"
    fi
    ;;
  *)
    usage
    exit 2
    ;;
esac