#!/bin/sh
# Set up Pi (https://pi.dev) to chat with the i14 API in one command:
#
#   curl -fsSL https://i14.ai/docs/install.sh | sh
#
# What it does:
#   1. Checks Node.js. If it is missing or too old, offers to install the
#      current Node.js LTS just for you (no password), in the folder Pi's own
#      installer uses. Then installs Pi with Pi's official installer if needed.
#   2. Asks for your i14 API key without showing it (or reads I14_API_KEY).
#   3. Checks the key against the i14 API.
#   4. Saves the key to ~/.config/i14/api-key and adds an "i14" provider to
#      Pi's ~/.pi/agent/models.json, keeping everything else in that file.
#   5. Creates ~/i14-playground (if it is missing) and starts Pi on the i14
#      model inside it, so the files Pi writes land there.
#
# Running it again is safe. Optional settings, mainly for testing:
#   I14_API_KEY   use this key instead of asking for one
#   I14_BASE_URL  API base URL (default https://api.i14.ai/v1)
#   I14_MODEL     model ID (default i14-sugarglider-nemotron3-super-120B-v1)
#   I14_NO_INSTALL_NODE  set to 1 (or true, yes) to never install Node.js;
#                        show the steps instead

set -eu

I14_PROVIDER="i14"
I14_DEFAULT_BASE_URL="https://api.i14.ai/v1"
I14_DEFAULT_MODEL="i14-sugarglider-nemotron3-super-120B-v1"
PI_INSTALLER_URL="https://pi.dev/install.sh"
# Test only: lets the tests use a local stand-in for nodejs.org.
NODE_DIST_URL=${I14_NODE_DIST_URL:-https://nodejs.org/dist}
KEY_TRIES=3

base_url=${I14_BASE_URL:-$I14_DEFAULT_BASE_URL}
base_url=${base_url%/}
model=${I14_MODEL:-$I14_DEFAULT_MODEL}
config_dir="$HOME/.config/i14"
key_file="$config_dir/api-key"
pi_agent_dir=${PI_CODING_AGENT_DIR:-$HOME/.pi/agent}
models_file="$pi_agent_dir/models.json"
# The folder Pi's installer uses for a Node.js it installs. Pi's launcher
# looks here, so Pi finds this Node.js in new terminals too.
node_home=${XDG_DATA_HOME:-$HOME/.local/share}/pi-node
# The folder Pi starts in, so the files it writes do not land loose in HOME.
playground="$HOME/i14-playground"
# The PATH a new terminal has. check_node may put another Node.js folder first.
user_path=$PATH

tmp_dir=
node_stage=
tty_state=

say() {
  printf '%s\n' "$*"
}

fail() {
  printf '%s\n' "$*" >&2
  exit 1
}

have_tty() {
  ( : <>/dev/tty ) 2>/dev/null
}

restore_tty() {
  if [ -n "$tty_state" ]; then
    stty "$tty_state" </dev/tty 2>/dev/null || true
    tty_state=
  fi
}

cleanup() {
  restore_tty
  if [ -n "$node_stage" ]; then
    rm -rf "$node_stage"
  fi
  if [ -n "$tmp_dir" ]; then
    rm -rf "$tmp_dir"
  fi
}

trap cleanup EXIT
trap 'cleanup; printf "\n" >&2; exit 130' INT
trap 'cleanup; exit 143' TERM

tmp_dir=$(mktemp -d "${TMPDIR:-/tmp}/i14-setup.XXXXXX")

# --- Node.js -----------------------------------------------------------------

node_help() {
  case "$(uname -s)" in
    Darwin) printf 'Install the current Node.js LTS from https://nodejs.org (or run: brew install node), then run this command again.' ;;
    *) printf 'Install the current Node.js LTS from https://nodejs.org (or with your package manager), then run this command again.' ;;
  esac
}

# Prints ok, missing, old or no-zstd for the node first on PATH.
node_state() {
  if ! command -v node >/dev/null 2>&1; then
    printf missing
  elif ! node -e 'const [a, b] = process.versions.node.split(".").map(Number); process.exit(a > 22 || (a === 22 && b >= 19) ? 0 : 1)' >/dev/null 2>&1; then
    printf old
  elif ! node -e 'process.exit(typeof require("zlib").createZstdDecompress === "function" ? 0 : 1)' >/dev/null 2>&1; then
    printf no-zstd
  else
    printf ok
  fi
}

use_installed_node() {
  PATH="$node_home/current/bin:$PATH"
  export PATH
  hash -r 2>/dev/null || true
}

sha256_of() {
  if command -v sha256sum >/dev/null 2>&1; then
    sha256sum "$1" | awk '{ print $1 }'
  else
    shasum -a 256 "$1" | awk '{ print $1 }'
  fi
}

# Installs the current Node.js LTS from nodejs.org into $node_home, for this
# user only. Nothing needs sudo or goes into a system folder. The download is
# checked against the SHA-256 sums nodejs.org publishes.
install_node() {
  case "$(uname -s)" in
    Darwin) node_os=darwin ;;
    Linux) node_os=linux ;;
    *) say "This script cannot install Node.js on $(uname -s)."; return 1 ;;
  esac
  case "$(uname -m)" in
    x86_64 | amd64) node_arch=x64 ;;
    arm64 | aarch64) node_arch=arm64 ;;
    *) say "This script cannot install Node.js for a $(uname -m) processor."; return 1 ;;
  esac
  if [ "$node_os" = linux ] && ls /lib/ld-musl-* >/dev/null 2>&1; then
    say "This script cannot install Node.js on a musl-based Linux such as Alpine."
    return 1
  fi
  if ! command -v sha256sum >/dev/null 2>&1 && ! command -v shasum >/dev/null 2>&1; then
    say "This script needs sha256sum or shasum to check the Node.js download."
    return 1
  fi

  dl="$tmp_dir/node"
  mkdir -p "$dl"
  if ! curl -fsSL "$NODE_DIST_URL/index.tab" -o "$dl/index.tab"; then
    say "Could not download the Node.js release list from $NODE_DIST_URL."
    return 1
  fi
  # The newest release with an LTS name in the "lts" column.
  node_release=$(awk -F '\t' 'NR > 1 && $10 != "-" { print $1; exit }' "$dl/index.tab")
  case "$node_release" in
    v[0-9]*) ;;
    *) say "Could not find the current Node.js LTS release at $NODE_DIST_URL."; return 1 ;;
  esac
  node_name="node-$node_release-$node_os-$node_arch"
  say "Installing Node.js $node_release into $node_home (just for you, no password needed)..."
  if ! curl -fsSL "$NODE_DIST_URL/$node_release/SHASUMS256.txt" -o "$dl/SHASUMS256.txt" ||
    ! curl -fsSL "$NODE_DIST_URL/$node_release/$node_name.tar.gz" -o "$dl/$node_name.tar.gz"; then
    say "Could not download Node.js $node_release from $NODE_DIST_URL."
    return 1
  fi
  expected=$(awk -v f="$node_name.tar.gz" '$2 == f { print $1; exit }' "$dl/SHASUMS256.txt")
  if [ -z "$expected" ] || [ "$(sha256_of "$dl/$node_name.tar.gz")" != "$expected" ]; then
    say "The Node.js download did not match its published checksum, so it was not installed."
    return 1
  fi

  mkdir -p "$node_home"
  node_stage="$node_home/.i14-unpack.$$"
  rm -rf "$node_stage"
  mkdir "$node_stage"
  if ! tar -xzf "$dl/$node_name.tar.gz" -C "$node_stage" || [ ! -x "$node_stage/$node_name/bin/node" ]; then
    say "Could not unpack the Node.js download."
    return 1
  fi
  if [ "$(PATH="$node_stage/$node_name/bin:$PATH"; hash -r 2>/dev/null; node_state)" != ok ]; then
    say "Node.js $node_release does not run on this system, so it was not installed."
    return 1
  fi
  rm -rf "${node_home:?}/$node_name"
  mv "$node_stage/$node_name" "$node_home/$node_name"
  rm -rf "$node_stage"
  node_stage=
  rm -f "$node_home/current"
  ln -s "$node_home/$node_name" "$node_home/current"
  use_installed_node
  say "Installed Node.js $node_release."
  say ""
}

no_install_node() {
  case "${I14_NO_INSTALL_NODE:-}" in
    1 | [Tt][Rr][Uu][Ee] | [Yy][Ee][Ss]) return 0 ;;
  esac
  return 1
}

check_node() {
  # A Node.js installed here on an earlier run (or by Pi's installer) is not
  # on PATH in a new terminal, so look in its folder first.
  if [ "$(node_state)" != ok ] && [ -x "$node_home/current/bin/node" ]; then
    use_installed_node
  fi
  state=$(node_state)
  node_version=$(node --version 2>/dev/null || printf 'unknown')
  case "$state" in
    ok) return 0 ;;
    missing)
      problem="Pi needs Node.js 22.19 or newer, and Node.js is not installed."
      question="Node.js is needed and isn't installed. Install it now? [Y/n] "
      ;;
    old)
      problem="Pi needs Node.js 22.19 or newer. You have Node.js $node_version."
      question="Pi needs Node.js 22.19 or newer, and you have $node_version. Install a newer one just for Pi now? [Y/n] "
      ;;
    *)
      problem="Pi needs a Node.js build with zstd support, and Node.js $node_version does not have it."
      question="Pi needs a Node.js build with zstd support, and $node_version does not have it. Install a new one just for Pi now? [Y/n] "
      ;;
  esac

  if no_install_node || ! have_tty; then
    fail "$problem $(node_help)"
  fi
  printf '%s' "$question" >/dev/tty
  answer=
  IFS= read -r answer </dev/tty || answer=
  case "$answer" in
    [Nn] | [Nn][Oo]) fail "$problem $(node_help)" ;;
  esac
  if ! install_node; then
    fail "Node.js was not installed automatically. $(node_help)"
  fi
}

# --- Pi ----------------------------------------------------------------------

find_pi() {
  hash -r 2>/dev/null || true
  if found=$(command -v pi 2>/dev/null) && [ -n "$found" ]; then
    printf '%s' "$found"
    return 0
  fi
  for candidate in "$pi_agent_dir/bin/pi" "$HOME/.local/bin/pi"; do
    if [ -x "$candidate" ]; then
      printf '%s' "$candidate"
      return 0
    fi
  done
  if npm_prefix=$(npm prefix -g 2>/dev/null) && [ -x "$npm_prefix/bin/pi" ]; then
    printf '%s' "$npm_prefix/bin/pi"
    return 0
  fi
  return 1
}

install_pi() {
  if ! command -v npm >/dev/null 2>&1; then
    fail "Pi needs npm, which usually comes with Node.js. $(node_help)"
  fi
  say "Pi is not installed. Installing it with the official installer from pi.dev..."
  say ""
  if ! curl -fsSL "$PI_INSTALLER_URL" -o "$tmp_dir/pi-install.sh"; then
    fail "Could not download the Pi installer from $PI_INSTALLER_URL. Check your internet connection and try again."
  fi
  # Piping the installer's output keeps it from starting Pi on its own, so
  # this script can finish setting up i14 first. Its prompts still use the
  # terminal.
  { sh "$tmp_dir/pi-install.sh" </dev/null; printf '%s' "$?" >"$tmp_dir/pi-install.status"; } | cat
  if [ "$(cat "$tmp_dir/pi-install.status" 2>/dev/null)" != 0 ]; then
    fail "The Pi installer did not finish. Install Pi with: curl -fsSL $PI_INSTALLER_URL | sh  Then run this command again."
  fi
}

# --- API key -----------------------------------------------------------------

# Reads one line from the terminal without echoing it. Echo goes off before
# the prompt appears, so nothing typed or pasted right away is shown.
read_hidden() {
  tty_state=$(stty -g </dev/tty 2>/dev/null || true)
  stty -echo </dev/tty 2>/dev/null || true
  printf '%s' "$1" >/dev/tty
  hidden_value=
  IFS= read -r hidden_value </dev/tty || hidden_value=
  restore_tty
  hidden_value=$(printf '%s' "$hidden_value" | tr -d ' \t\r')
  printf '\n' >/dev/tty
}

# Prints the HTTP status of GET /models for the key in $api_key and leaves the
# response body in $tmp_dir/models.json. The key goes to curl on stdin, so it
# never shows up in the process list.
check_key() {
  printf 'header = "Authorization: Bearer %s"\n' "$api_key" |
    curl -sS --connect-timeout 15 --max-time 60 -K - \
      -o "$tmp_dir/models.json" -w '%{http_code}' "$base_url/models" 2>/dev/null ||
    true
}

key_is_well_formed() {
  case "$1" in
    "" | *[!A-Za-z0-9._~+/=:-]*) return 1 ;;
  esac
  return 0
}

get_key() {
  saved_key=
  if [ -r "$key_file" ]; then
    saved_key=$(cat "$key_file" 2>/dev/null || true)
  fi

  if [ -n "${I14_API_KEY:-}" ]; then
    api_key=$(printf '%s' "$I14_API_KEY" | tr -d ' \t\r\n')
    if ! key_is_well_formed "$api_key"; then
      fail "I14_API_KEY does not look like an i14 API key. Check the value and try again."
    fi
    status=$(check_key)
    case "$status" in
      200) return 0 ;;
      401 | 403) fail "The i14 API did not accept the key in I14_API_KEY. Create a new key at https://platform.i14.ai/ and try again." ;;
      *) unreachable "$status" ;;
    esac
  fi

  if ! have_tty; then
    fail "No terminal is available to ask for your API key. Set I14_API_KEY and run this command again."
  fi

  printf '%s\n' "The key stays hidden: nothing shows on screen while you paste or type it." >/dev/tty
  tries=0
  while [ "$tries" -lt "$KEY_TRIES" ]; do
    tries=$((tries + 1))
    if [ -n "$saved_key" ]; then
      read_hidden "What is your i14 API key? (press Enter to keep your saved key) "
      api_key=${hidden_value:-$saved_key}
    else
      read_hidden "What is your i14 API key? "
      api_key=$hidden_value
    fi
    if [ -z "$api_key" ]; then
      say "No key entered. Get a key at https://platform.i14.ai/."
      continue
    fi
    if ! key_is_well_formed "$api_key"; then
      say "That does not look like an i14 API key. Paste the whole key and try again."
      continue
    fi
    status=$(check_key)
    case "$status" in
      200) return 0 ;;
      401 | 403) say "The i14 API did not accept that key. Check it and try again." ;;
      *) unreachable "$status" ;;
    esac
  done
  fail "No valid key after $KEY_TRIES tries. Create a new key at https://platform.i14.ai/ and run this command again."
}

unreachable() {
  if [ "$1" = 000 ] || [ -z "$1" ]; then
    fail "Could not reach the i14 API at $base_url. Check your internet connection and try again."
  fi
  fail "The i14 API at $base_url returned HTTP $1. Try again in a few minutes."
}

check_model() {
  model_state=$(I14_MODEL_ID="$model" node -e '
const fs = require("fs");
let list;
try { list = JSON.parse(fs.readFileSync(process.argv[1], "utf8")).data; } catch { list = null; }
if (!Array.isArray(list)) { console.log("bad-response"); process.exit(0); }
const entry = list.find((m) => m && m.id === process.env.I14_MODEL_ID);
console.log(!entry ? "missing" : entry.loaded === false ? "not-loaded" : "ok");
' "$tmp_dir/models.json")
  case "$model_state" in
    ok) ;;
    not-loaded) fail "The model $model is not running on the i14 API right now. Try again in a few minutes." ;;
    missing) fail "The model $model is not available to your key. Check the models listed at https://platform.i14.ai/." ;;
    *) fail "The i14 API at $base_url sent an unexpected reply. Try again in a few minutes." ;;
  esac
}

save_key() {
  old_umask=$(umask)
  umask 077
  mkdir -p "$config_dir"
  chmod 700 "$config_dir"
  printf '%s\n' "$api_key" >"$key_file.tmp.$$"
  chmod 600 "$key_file.tmp.$$"
  mv -f "$key_file.tmp.$$" "$key_file"
  umask "$old_umask"
}

# --- Pi config ---------------------------------------------------------------

# The command Pi runs to read the key, quoted only when the path needs it.
key_command() {
  case "$key_file" in
    *[!A-Za-z0-9._/-]*)
      quoted=$(printf '%s' "$key_file" | sed "s/'/'\\\\''/g")
      printf "!cat '%s'" "$quoted"
      ;;
    *) printf '!cat %s' "$key_file" ;;
  esac
}

configure_pi() {
  mkdir -p "$pi_agent_dir"
  # The ${...} below are JavaScript template literals, not shell expansions.
  # shellcheck disable=SC2016
  result=$(I14_MODELS_FILE="$models_file" I14_PROVIDER="$I14_PROVIDER" I14_BASE_URL="$base_url" \
    I14_MODEL_ID="$model" I14_KEY_COMMAND="$(key_command)" node -e '
const fs = require("fs");
const env = process.env;
const file = env.I14_MODELS_FILE;
let text = null;
let config = {};
try { text = fs.readFileSync(file, "utf8"); } catch (e) { if (e.code !== "ENOENT") throw e; }
if (text !== null && text.trim() !== "") {
  try { config = JSON.parse(text.replace(/^﻿/, "")); } catch { console.log("not-json"); process.exit(0); }
}
if (!config || typeof config !== "object" || Array.isArray(config)) { console.log("not-object"); process.exit(0); }
if (config.providers === undefined) config.providers = {};
if (!config.providers || typeof config.providers !== "object" || Array.isArray(config.providers)) { console.log("not-object"); process.exit(0); }
const old = config.providers[env.I14_PROVIDER];
const provider = old && typeof old === "object" && !Array.isArray(old) ? { ...old } : {};
provider.baseUrl = env.I14_BASE_URL;
provider.api = "openai-completions";
provider.apiKey = env.I14_KEY_COMMAND;
const models = Array.isArray(provider.models) ? provider.models.slice() : [];
if (!models.some((m) => m && m.id === env.I14_MODEL_ID)) models.push({ id: env.I14_MODEL_ID });
provider.models = models;
config.providers[env.I14_PROVIDER] = provider;
const next = JSON.stringify(config, null, 2) + "\n";
if (text !== null && text === next) { console.log("unchanged"); process.exit(0); }
let mode = 0o600;
if (text !== null) {
  mode = fs.statSync(file).mode & 0o777;
  const stamp = new Date().toISOString().replace(/[-:]/g, "").replace(/\..*/, "");
  const backup = `${file}.bak-${stamp}`;
  fs.copyFileSync(file, backup);
  console.log("updated " + backup);
} else {
  console.log("created");
}
const tmp = `${file}.tmp-${process.pid}`;
fs.writeFileSync(tmp, next, { mode });
fs.renameSync(tmp, file);
')
  case "$result" in
    created) say "Created $models_file with the i14 provider." ;;
    unchanged) say "Pi is already set up for i14 in $models_file." ;;
    "updated "*) say "Added the i14 provider to $models_file (backup: ${result#updated })." ;;
    not-json)
      I14_PROVIDER="$I14_PROVIDER" I14_BASE_URL="$base_url" I14_MODEL_ID="$model" I14_KEY_COMMAND="$(key_command)" node -e '
const env = process.env;
const provider = { baseUrl: env.I14_BASE_URL, api: "openai-completions", apiKey: env.I14_KEY_COMMAND, models: [{ id: env.I14_MODEL_ID }] };
console.log(JSON.stringify({ [env.I14_PROVIDER]: provider }, null, 2).replace(/^\{\n|\n\}$/g, ""));
' >"$tmp_dir/provider.json"
      say "$models_file is not plain JSON (it may have comments), so it was left as is."
      say "Add this entry inside its \"providers\" object, then run: pi --model $I14_PROVIDER/$model"
      say ""
      cat "$tmp_dir/provider.json"
      exit 1
      ;;
    *) fail "$models_file has an unexpected layout, so it was left as is. Fix or move it, then run this command again." ;;
  esac
}

# --- Main --------------------------------------------------------------------

main() {
  say "Setting up Pi to chat with i14 ($model)."
  say ""

  check_node
  if ! pi_bin=$(find_pi); then
    install_pi
    if ! pi_bin=$(find_pi); then
      fail "Pi was installed, but this script could not find it. Open a new terminal and run this command again."
    fi
  fi
  say "Using Pi at $pi_bin"

  get_key
  check_model
  save_key
  say "Saved your API key to $key_file"
  configure_pi

  launch="cd ~/i14-playground && $pi_bin --model $I14_PROVIDER/$model"
  if [ -t 1 ] && have_tty; then
    if ! mkdir -p "$playground" || ! cd "$playground"; then
      fail "Could not create the folder $playground for your files. Check that you can write to your home folder, then run this command again."
    fi
    say ""
    say "Pi is working in ~/i14-playground. Your files will be saved there."
    # Name Pi by its full path when the plain pi command would not find it.
    next_pi=pi
    if [ "$(PATH=$user_path; hash -r 2>/dev/null; command -v pi 2>/dev/null)" != "$pi_bin" ]; then
      next_pi=$pi_bin
    fi
    say "Next time, run: cd ~/i14-playground && $next_pi --model $I14_PROVIDER/$model"
    say ""
    cleanup
    trap - EXIT INT TERM
    exec "$pi_bin" --model "$I14_PROVIDER/$model" </dev/tty
  fi
  mkdir -p "$playground" 2>/dev/null || true
  say ""
  say "Setup complete. Start chatting with: $launch"
}

main "$@"
