#!/usr/bin/env bash
# TokenKeeper installer — Claude Code-style `curl | bash` flow.
#
# Designed to be runnable as either:
#
#   curl -fsSL <hosted-url>/install.sh | bash       # remote (when hosted)
#   ./scripts/install.sh                            # in-repo
#
# Flow:
#
#   1. Detects Python 3.10+ (the CLI is pure-stdlib but uses 3.10 syntax).
#   2. Installs the `tokenkeeper` CLI to $TOKENKEEPER_PREFIX/bin (default
#      ~/.local/bin) by either copying cli/tokenkeeper from a local repo
#      checkout or downloading it from $TOKENKEEPER_RAW_URL.
#   3. Detects which harness integrations are wanted. For each of the
#      three supported harnesses (Claude Code, Cursor, Codex) we check
#      whether the corresponding CLI is on PATH. The user gets a small
#      multiselect TUI (arrow keys + space + enter) with installable
#      detected harnesses pre-checked. Harnesses whose integration isn't
#      implemented yet are shown but locked — even if the CLI is detected
#      we can't wire anything up. Pass `--harness=claude,cursor,codex`
#      (or `--harness=all`) to skip the prompt; non-interactive runs
#      install everything detected and implemented.
#   4. Registers the user with the TokenKeeper server (`tokenkeeper register`).
#   5. For each selected harness, runs the corresponding wiring step.
#      Today only Claude Code is implemented.
#   6. Wires Claude Code's OTel env block at the local forwarder (only when
#      Claude Code was selected).
#   7. Enables the Claude Code status-line (only when Claude Code was
#      selected).
#

# Hosting: when no local checkout is found, the CLI is downloaded from
# $TOKENKEEPER_INSTALL_BASE (defaults to https://install.tokenkeeper.ai,
# served by the `release` GitHub workflow via gh-pages). The site lays
# out artifacts as:
#   /install.sh
#   /plugin.tar.gz
#   /bin/tokenkeeper-<os>-<arch>
#   /bin/SHA256SUMS

set -euo pipefail

PREFIX="${TOKENKEEPER_PREFIX:-$HOME/.local}"
BIN_DIR="$PREFIX/bin"
DEST="$BIN_DIR/tokenkeeper"
SERVER="${TOKENKEEPER_SERVER:-http://127.0.0.1:8000}"
INSTALL_BASE="${TOKENKEEPER_INSTALL_BASE:-https://install.tokenkeeper.ai}"

# All supported harnesses. Keep helpers below in sync. We avoid `declare
# -A` because macOS still ships bash 3.2, which lacks associative
# arrays.
ALL_HARNESSES=(claude cursor codex)

harness_label() {
    case "$1" in
        claude) echo "Claude Code" ;;
        cursor) echo "Cursor" ;;
        codex)  echo "Codex" ;;
        *)      echo "$1" ;;
    esac
}

harness_cli() {
    case "$1" in
        claude) echo "claude" ;;
        cursor) echo "cursor" ;;
        codex)  echo "codex" ;;
        *)      echo "$1" ;;
    esac
}

harness_is_known() {
    case "$1" in
        claude|cursor|codex) return 0 ;;
        *) return 1 ;;
    esac
}

harness_is_implemented() {
    case "$1" in
        claude) return 0 ;;
        cursor|codex) return 1 ;;
        *) return 1 ;;
    esac
}

is_harness_detected() {
    command -v "$(harness_cli "$1")" >/dev/null 2>&1
}

# A harness is installable iff its CLI is on PATH AND its tokenkeeper
# integration is wired up. Two independent gates: missing CLI is the
# user's problem to fix (install the harness, re-run), missing
# integration is ours.
harness_is_installable() {
    is_harness_detected "$1" && harness_is_implemented "$1"
}

# Render the trailing status hint shown to the right of each harness in
# the picker. "Installation not detected" wins over "integration not
# yet implemented" when both apply — fixing the missing CLI is the
# user's first step regardless.
harness_status_hint() {
    local key="$1"
    if ! is_harness_detected "$key"; then
        echo "installation not detected"
        return
    fi
    if ! harness_is_implemented "$key"; then
        echo "integration not yet implemented"
        return
    fi
    echo "detected"
}

# `--local` installs the Claude Code plugin from the local checkout instead
# of the default GitHub marketplace. Used by `bazel run //:install` and
# anyone testing edits before publishing.
LOCAL=0
HARNESS_FLAG=""
for arg in "$@"; do
    case "$arg" in
        --local) LOCAL=1 ;;
        --harness=*) HARNESS_FLAG="${arg#--harness=}" ;;
        # Back-compat: --ide= was the original spelling. Accept it
        # silently so older docs / muscle memory keep working.
        --ide=*) HARNESS_FLAG="${arg#--ide=}" ;;
        -h|--help)
            cat <<'USAGE'
Usage: install.sh [--local] [--harness=LIST]

  --local         Install the Claude Code plugin from this local
                  checkout instead of the default GitHub marketplace.
                  Useful for testing plugin edits before publishing.
  --harness=LIST  Comma-separated subset of {claude,cursor,codex}, or
                  `all`. Skips the interactive selector. Non-interactive
                  runs that omit this default to all detected,
                  implemented harnesses.

Environment:
  TOKENKEEPER_PREFIX        Install root (default: ~/.local)
  TOKENKEEPER_SERVER        Server URL (default: http://127.0.0.1:8000)
  TOKENKEEPER_INSTALL_BASE  Base URL for remote install assets
                            (default: https://install.tokenkeeper.ai).
                            Used only when no local checkout is found.
  TOKENKEEPER_PREBUILT_CLI  Path to a prebuilt tokenkeeper binary; if
                            set, skips both local build and remote
                            download.
USAGE
            exit 0
            ;;
    esac
done

# A controlling terminal lets us drive the multiselect even when stdin is
# the curl pipe. We open it on FD 3 if available.
TTY_FD=""
if [ -r /dev/tty ] && [ -w /dev/tty ]; then
    if exec 3<>/dev/tty 2>/dev/null; then
        TTY_FD=3
    fi
fi
if [ -n "$TTY_FD" ]; then
    INTERACTIVE=1
else
    INTERACTIVE=0
fi

bold()  { printf '\033[1m%s\033[0m\n' "$*"; }
info()  { printf '  %s\n' "$*"; }
warn()  { printf '\033[33m  %s\033[0m\n' "$*" >&2; }
fail()  { printf '\033[31merror:\033[0m %s\n' "$*" >&2; exit 1; }

# Quiet Ctrl+C / SIGTERM handler. We exit 130 without a message — the
# multiselect chains here after restoring the TTY (stty echo + cursor)
# so the terminal stays sane regardless of where the signal lands.
abort() { exit 130; }
trap 'abort' INT TERM

ask_yes_no() {
    local prompt="$1" default="${2:-y}" suffix reply
    if [ "$default" = "y" ]; then suffix="[Y/n]"; else suffix="[y/N]"; fi
    if [ "$INTERACTIVE" -eq 0 ]; then
        info "(non-interactive: defaulting to $default)"
        [ "$default" = "y" ]
        return $?
    fi
    printf '  %s %s ' "$prompt" "$suffix" >&"$TTY_FD"
    IFS= read -r reply <&"$TTY_FD" || reply=""
    reply="${reply:-$default}"
    case "$reply" in
        y|Y|yes|YES) return 0 ;;
        *) return 1 ;;
    esac
}

# ---------------------------------------------------------------------------
# Harness selection
# ---------------------------------------------------------------------------

# Parse --harness=LIST. Stores result in SELECTED_HARNESSES. Returns 0 if
# a non-empty subset was specified, 1 if the flag was empty/unset.
parse_harness_flag() {
    local list="$1" token
    SELECTED_HARNESSES=()
    if [ -z "$list" ]; then
        return 1
    fi
    if [ "$list" = "all" ]; then
        SELECTED_HARNESSES=("${ALL_HARNESSES[@]}")
        return 0
    fi
    IFS=',' read -ra tokens <<<"$list"
    for token in "${tokens[@]}"; do
        token="$(echo "$token" | tr -d '[:space:]')"
        [ -z "$token" ] && continue
        if ! harness_is_known "$token"; then
            fail "unknown harness in --harness: '$token' (allowed: ${ALL_HARNESSES[*]}, all)"
        fi
        if ! is_harness_detected "$token"; then
            warn "harness '$token' installation not detected on PATH — skipping."
            continue
        fi
        if ! harness_is_implemented "$token"; then
            warn "harness '$token' integration not yet implemented — skipping."
            continue
        fi
        SELECTED_HARNESSES+=("$token")
    done
    [ "${#SELECTED_HARNESSES[@]}" -gt 0 ]
}

# Pure-bash multiselect TUI. Reads from FD 3 (the controlling tty), draws
# to FD 3 as well so the picker is visible under `curl | bash`.
#
#   $1 = title line
#   args after that = "key:checked:enabled" triples. enabled=0 makes the
#   row read-only (rendered dim, space toggle silently does nothing).
#
# On exit, MULTISELECT_RESULT is populated with the selected enabled keys.
multiselect() {
    local title="$1"; shift
    local keys=() checked=() enabled=() i n=0
    for entry in "$@"; do
        local k="${entry%%:*}" rest="${entry#*:}"
        keys+=("$k")
        checked+=("${rest%%:*}")
        enabled+=("${rest#*:}")
        n=$((n + 1))
    done
    MULTISELECT_RESULT=()
    [ "$n" -eq 0 ] && return 0

    local cursor=0 key rest
    # Hide cursor + raw input. Restore on any exit path.
    stty -echo -icanon min 1 time 0 <&"$TTY_FD" >&"$TTY_FD" 2>/dev/null || true
    printf '\033[?25l' >&"$TTY_FD"
    restore() {
        stty echo icanon <&"$TTY_FD" >&"$TTY_FD" 2>/dev/null || true
        printf '\033[?25h' >&"$TTY_FD"
    }
    # On signal: restore TTY first, then delegate to the script-level
    # abort handler so the user sees the same "Installation aborted"
    # message regardless of where Ctrl+C lands.
    trap 'restore' EXIT
    trap 'restore; abort' INT TERM

    printf '%s\n' "$title" >&"$TTY_FD"
    printf '  (\xe2\x86\x91/\xe2\x86\x93 move, space toggle, enter confirm; locked rows can'\''t be toggled)\n' >&"$TTY_FD"
    # Reserve n lines we will redraw in place.
    for ((i = 0; i < n; i++)); do printf '\n' >&"$TTY_FD"; done
    printf '\033[%dA' "$n" >&"$TTY_FD"

    redraw() {
        local idx mark line key label hint dim
        printf '\r' >&"$TTY_FD"
        for ((idx = 0; idx < n; idx++)); do
            key="${keys[$idx]}"
            label="$(harness_label "$key")"
            hint="$(harness_status_hint "$key")"
            if [ "${enabled[$idx]}" = "0" ]; then
                mark="[-]"
                dim=1
            else
                if [ "${checked[$idx]}" = "1" ]; then mark="[x]"; else mark="[ ]"; fi
                dim=0
            fi
            local prefix
            if [ "$idx" -eq "$cursor" ]; then prefix="> "; else prefix="  "; fi
            if [ "$dim" -eq 1 ]; then
                line=$(printf '\033[2m%s%s %s  %s\033[0m' "$prefix" "$mark" "$label" "$hint")
            elif [ "$idx" -eq "$cursor" ]; then
                line=$(printf '\033[36m%s%s %s\033[0m  \033[2m%s\033[0m' "$prefix" "$mark" "$label" "$hint")
            else
                line=$(printf '%s%s %s  \033[2m%s\033[0m' "$prefix" "$mark" "$label" "$hint")
            fi
            printf '\033[2K%s\n' "$line" >&"$TTY_FD"
        done
        printf '\033[%dA' "$n" >&"$TTY_FD"
    }

    redraw
    while :; do
        IFS= read -rsn1 key <&"$TTY_FD" || break
        case "$key" in
            $'\x1b')
                IFS= read -rsn2 rest <&"$TTY_FD" || rest=""
                case "$rest" in
                    '[A') (( cursor > 0 )) && cursor=$((cursor - 1)) ;;
                    '[B') (( cursor < n - 1 )) && cursor=$((cursor + 1)) ;;
                esac
                ;;
            ' ')
                if [ "${enabled[$cursor]}" = "1" ]; then
                    if [ "${checked[$cursor]}" = "1" ]; then
                        checked[$cursor]=0
                    else
                        checked[$cursor]=1
                    fi
                fi
                ;;
            'k') (( cursor > 0 )) && cursor=$((cursor - 1)) ;;
            'j') (( cursor < n - 1 )) && cursor=$((cursor + 1)) ;;
            $'\n'|$'\r'|'') break ;;
            'q') checked=(); for ((i = 0; i < n; i++)); do checked+=(0); done; break ;;
        esac
        redraw
    done
    # Move past the picker block.
    printf '\033[%dB' "$n" >&"$TTY_FD"
    restore
    trap - EXIT
    # Re-arm the script-level abort handler now that we're done with
    # the picker's TTY-restore variant.
    trap 'abort' INT TERM

    for ((i = 0; i < n; i++)); do
        if [ "${enabled[$i]}" = "1" ] && [ "${checked[$i]}" = "1" ]; then
            MULTISELECT_RESULT+=("${keys[$i]}")
        fi
    done
}

# Resolve SELECTED_HARNESSES from --harness flag, TUI, or detection-only
# fallback.
resolve_selected_harnesses() {
    if parse_harness_flag "$HARNESS_FLAG"; then
        return
    fi
    local installable=()
    local detected_pending=()
    local missing=()
    for h in "${ALL_HARNESSES[@]}"; do
        if ! is_harness_detected "$h"; then
            missing+=("$h")
            continue
        fi
        if harness_is_implemented "$h"; then
            installable+=("$h")
        else
            detected_pending+=("$h")
        fi
    done
    if [ "$INTERACTIVE" -eq 0 ]; then
        SELECTED_HARNESSES=("${installable[@]}")
        if [ "${#missing[@]}" -gt 0 ]; then
            info "installation not detected: ${missing[*]}"
        fi
        if [ "${#detected_pending[@]}" -gt 0 ]; then
            info "integration not yet implemented: ${detected_pending[*]}"
        fi
        if [ "${#SELECTED_HARNESSES[@]}" -eq 0 ]; then
            warn "no installable harness on PATH; skipping harness wiring."
        else
            info "non-interactive: selecting ${SELECTED_HARNESSES[*]}"
        fi
        return
    fi
    local entries=()
    for h in "${ALL_HARNESSES[@]}"; do
        local pre=0 ena=0
        if harness_is_installable "$h"; then
            ena=1
            pre=1
        fi
        entries+=("$h:$pre:$ena")
    done
    multiselect "Which harnesses should TokenKeeper wire up?" "${entries[@]}"
    SELECTED_HARNESSES=("${MULTISELECT_RESULT[@]}")
}

# ---------------------------------------------------------------------------
# Per-harness wiring
# ---------------------------------------------------------------------------

harness_wire_claude() {
    if ! command -v claude >/dev/null 2>&1; then
        warn "\`claude\` not on PATH — skipping plugin registration."
        warn "Once Claude Code is installed: \`tokenkeeper claude-register\`."
        return 0
    fi
    local rc=0
    if [ "$LOCAL" -eq 1 ]; then
        local repo_root
        repo_root="$(cd "$SCRIPT_DIR/.." && pwd)"
        info "installing plugin from local checkout ($repo_root)"
        "$DEST" claude-register --marketplace "$repo_root" || rc=$?
    else
        "$DEST" claude-register || rc=$?
    fi
    if [ "$rc" -ne 0 ]; then
        warn "\`tokenkeeper claude-register\` failed; re-run manually after"
        warn "Claude Code is configured."
    fi
}

# Cursor / Codex wiring lives here once those harnesses are implemented.
# Until then the picker locks them out and the --harness flag refuses to
# select them, so we never end up here with an unimplemented key.

harness_is_selected() {
    local needle="$1" h
    for h in "${SELECTED_HARNESSES[@]:-}"; do
        [ "$h" = "$needle" ] && return 0
    done
    return 1
}

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
LOCAL_GO_SRC="$REPO_ROOT/cli/main.go"
PREBUILT="${TOKENKEEPER_PREBUILT_CLI:-}"

# Three modes:
#   prebuilt = caller pointed us at a binary; just install it.
#   local    = repo checkout present; build via bazel.
#   remote   = no checkout; download a prebuilt from $INSTALL_BASE/bin.
if [ -n "$PREBUILT" ] && [ -x "$PREBUILT" ]; then
    INSTALL_MODE="prebuilt"
elif [ -f "$LOCAL_GO_SRC" ]; then
    INSTALL_MODE="local"
else
    INSTALL_MODE="remote"
fi

# ---------------------------------------------------------------------------
# 1. Prerequisites
# ---------------------------------------------------------------------------
bold "==> Checking prerequisites"
case "$INSTALL_MODE" in
    local)
        if ! command -v bazel >/dev/null 2>&1; then
            fail "bazel not found on PATH. Install Bazel (https://bazel.build) and re-run."
        fi
        info "bazel $(bazel --version | head -1) ok"
        ;;
    remote)
        if ! command -v curl >/dev/null 2>&1; then
            fail "curl not found on PATH; required for remote install."
        fi
        info "remote install from $INSTALL_BASE"
        ;;
    prebuilt)
        info "using prebuilt CLI at $PREBUILT"
        ;;
esac

# ---------------------------------------------------------------------------
# 2. Install the CLI
# ---------------------------------------------------------------------------
bold "==> Installing tokenkeeper CLI to $DEST"
mkdir -p "$BIN_DIR"

case "$INSTALL_MODE" in
    prebuilt)
        install -m 0755 "$PREBUILT" "$DEST"
        ;;
    local)
        info "building //cli:tokenkeeper via bazel"
        (cd "$REPO_ROOT" && bazel build //cli:tokenkeeper)
        BUILT="$REPO_ROOT/bazel-bin/cli/tokenkeeper_/tokenkeeper"
        if [ ! -x "$BUILT" ]; then
            fail "bazel build succeeded but expected output $BUILT is missing."
        fi
        install -m 0755 "$BUILT" "$DEST"
        ;;
    remote)
        case "$(uname -s)" in
            Darwin) os="darwin" ;;
            Linux)  os="linux" ;;
            *) fail "unsupported OS for remote install: $(uname -s). Clone the repo and run scripts/install.sh from a checkout." ;;
        esac
        case "$(uname -m)" in
            arm64|aarch64) arch="arm64" ;;
            x86_64|amd64)  arch="amd64" ;;
            *) fail "unsupported arch for remote install: $(uname -m)." ;;
        esac
        url="$INSTALL_BASE/bin/tokenkeeper-${os}-${arch}"
        info "downloading $url"
        if ! curl -fsSL "$url" -o "$DEST.tmp"; then
            fail "failed to download $url. Set TOKENKEEPER_INSTALL_BASE to override or check connectivity."
        fi
        chmod 0755 "$DEST.tmp"
        mv "$DEST.tmp" "$DEST"
        ;;
esac

case ":$PATH:" in
    *":$BIN_DIR:"*) ;;
    *) warn "$BIN_DIR is not on PATH — add it to your shell profile."
       warn "  export PATH=\"$BIN_DIR:\$PATH\""
       ;;
esac

info "tokenkeeper installed: $("$DEST" --help | head -1)"

# Sanity-check that the CLI we just copied has the subcommands the rest of
# this script is about to call. A stale source checkout is the usual
# culprit: it copies cleanly, then `status-bar --enable` later fails with
# "invalid choice" and the user is left wondering why no status bar
# appeared. Fail fast with a clear message instead.
REQUIRED_SUBCOMMANDS=(claude-register telemetry-enable status-bar)
HELP_OUTPUT="$("$DEST" --help 2>&1 || true)"
MISSING_SUBCOMMANDS=()
for sub in "${REQUIRED_SUBCOMMANDS[@]}"; do
    case "$HELP_OUTPUT" in
        *"$sub"*) ;;
        *) MISSING_SUBCOMMANDS+=("$sub") ;;
    esac
done
if [ "${#MISSING_SUBCOMMANDS[@]}" -gt 0 ]; then
    warn "the installed tokenkeeper CLI is missing: ${MISSING_SUBCOMMANDS[*]}"
    case "$INSTALL_MODE" in
        local)    fail "stale source at $LOCAL_GO_SRC — \`git pull\` in that checkout and re-run install.sh." ;;
        remote)   fail "stale CLI at $INSTALL_BASE — re-run the \`release\` workflow or override TOKENKEEPER_INSTALL_BASE." ;;
        prebuilt) fail "stale prebuilt CLI at $PREBUILT — rebuild and re-run install.sh." ;;
    esac
fi

# ---------------------------------------------------------------------------
# 3. Pick harness integrations
# ---------------------------------------------------------------------------
bold "==> Selecting harness integrations"
SELECTED_HARNESSES=()
resolve_selected_harnesses
if [ "${#SELECTED_HARNESSES[@]}" -eq 0 ]; then
    info "no harness integrations selected."
else
    info "selected: ${SELECTED_HARNESSES[*]}"
fi

# ---------------------------------------------------------------------------
# 4. Register the user
# ---------------------------------------------------------------------------
bold "==> Registering with the TokenKeeper server ($SERVER)"
if ! TOKENKEEPER_SERVER="$SERVER" "$DEST" register; then
    warn "register failed — is the TokenKeeper server running at $SERVER?"
    warn "you can re-run \`tokenkeeper register\` later."
fi

# ---------------------------------------------------------------------------
# 5. Per-harness wiring
# ---------------------------------------------------------------------------
for h in "${SELECTED_HARNESSES[@]:-}"; do
    [ -z "$h" ] && continue
    bold "==> Wiring up $(harness_label "$h")"
    "harness_wire_$h"
done

# ---------------------------------------------------------------------------
# 6. Enable telemetry env block (Claude Code only)
# ---------------------------------------------------------------------------
if harness_is_selected claude; then
    bold "==> Enabling Claude Code telemetry forwarding"
    if ! "$DEST" telemetry-enable; then
        warn "\`tokenkeeper telemetry-enable\` failed; re-run manually."
    fi
fi

# ---------------------------------------------------------------------------
# 7. Enable Claude Code status-line (Claude Code only)
# ---------------------------------------------------------------------------
if harness_is_selected claude; then
    bold "==> Enabling Claude Code status-line"
    if ! "$DEST" status-bar --enable; then
        warn "\`tokenkeeper status-bar --enable\` failed; re-run manually."
    fi
fi

bold "==> Done"
if harness_is_selected claude; then
    info "Open \`claude\` to start streaming live telemetry."
fi
info "Visit $SERVER once the server is running to see your KPIs."