#!/usr/bin/env bash
#
# Try PgCache against a PostgreSQL database, in Docker, in one command.
#
#   curl -fsSL https://pgcache.com/try.sh | bash
#   curl -fsSL https://pgcache.com/try.sh | bash -s -- --upstream postgres://user@host:5432/db
#   curl -fsSL https://pgcache.com/try.sh | bash -s -- status     # live hit rate
#   curl -fsSL https://pgcache.com/try.sh | bash -s -- stop       # remove the container
#
# What it does: finds Docker, works out which database to point at (argument,
# DATABASE_URL, .env, PG* variables, or a running postgres container), runs
# `pgcache check` against it, starts pgcache/pgcache as a container named
# pgcache-try, and prints the connection string to use.
#
# Source: https://pgcache.com/try.sh

set -euo pipefail

IMAGE="${PGCACHE_IMAGE:-pgcache/pgcache:latest}"
CONTAINER="${PGCACHE_CONTAINER:-pgcache-try}"
PROXY_PORT="${PGCACHE_PORT:-6432}"
METRICS_PORT="${PGCACHE_METRICS_PORT:-9090}"

UPSTREAM=""
ASSUME_YES=false
CHECK_ONLY=false
COMMAND="start"

# Origin, resolved by origin_discover.
ORIGIN_SOURCE=""
ORIGIN_USER=""
ORIGIN_PASSWORD=""
ORIGIN_HOST=""
ORIGIN_PORT=""
ORIGIN_DATABASE=""
ORIGIN_SSLMODE=""
# Set when the origin is a container on this Docker host.
ORIGIN_CONTAINER=""
# Extra `docker run` flags that make the origin reachable from the container.
# Expanded as ${DOCKER_NET_ARGS[@]+"${DOCKER_NET_ARGS[@]}"}: macOS ships bash 3.2,
# where an empty array is an unbound variable under set -u.
DOCKER_NET_ARGS=()

# --- output ------------------------------------------------------------------

say()  { printf '%s\n' "$*"; }
note() { printf '  %s\n' "$*"; }
die()  { printf 'pgcache: %s\n' "$*" >&2; exit 1; }

usage() {
    sed -n '2,15p' "$0" 2>/dev/null | sed 's/^# \{0,1\}//' || true
    cat <<USAGE
Options:
  --upstream URL   Origin database (postgres://user[:password]@host[:port]/db)
  --port N         Host port for the proxy (default 6432, next free if taken)
  --image IMAGE    Image to run (default pgcache/pgcache:latest)
  --check-only     Run the origin check and stop
  --yes            Do not ask; accept the discovered origin and defaults
  status | stop    Show live cache counters / remove the container
USAGE
}

# Under `curl | bash` stdin is the script, so every prompt reads the terminal.
tty_available() { { : < /dev/tty; } 2>/dev/null; }

# ask VAR "prompt" "default"
ask() {
    local var="$1" prompt="$2" default="${3:-}" answer
    if ! tty_available; then
        [[ -n "$default" ]] && { printf -v "$var" '%s' "$default"; return 0; }
        die "$prompt: no terminal to ask on; pass --upstream or set DATABASE_URL"
    fi
    if [[ -n "$default" ]]; then
        printf '%s [%s]: ' "$prompt" "$default" > /dev/tty
    else
        printf '%s: ' "$prompt" > /dev/tty
    fi
    read -r answer < /dev/tty
    printf -v "$var" '%s' "${answer:-$default}"
}

ask_secret() {
    local var="$1" prompt="$2" answer
    tty_available || die "$prompt: no terminal to ask on; put the password in --upstream"
    printf '%s: ' "$prompt" > /dev/tty
    read -r -s answer < /dev/tty
    printf '\n' > /dev/tty
    printf -v "$var" '%s' "$answer"
}

# confirm "question" -> 0 for yes
confirm() {
    local answer
    $ASSUME_YES && return 0
    tty_available || return 1
    printf '%s [Y/n]: ' "$1" > /dev/tty
    read -r answer < /dev/tty
    [[ -z "$answer" || "$answer" == [Yy]* ]]
}

# --- docker ------------------------------------------------------------------

docker_require() {
    command -v docker > /dev/null 2>&1 \
        || die "Docker is not installed. Get it from https://docs.docker.com/get-docker/ and run this again."
    docker info > /dev/null 2>&1 \
        || die "Docker is installed but the daemon is not running. Start Docker and run this again."
}

container_exists() { docker inspect "$CONTAINER" > /dev/null 2>&1; }

# --- origin discovery --------------------------------------------------------

url_decode() { printf '%b' "${1//%/\\x}"; }

# Fill ORIGIN_* from a postgres:// URL. Returns 1 when it does not look like one.
url_parse() {
    local url="$1" rest userpart hostport dbpart params
    case "$url" in
        postgres://*)   rest="${url#postgres://}" ;;
        postgresql://*) rest="${url#postgresql://}" ;;
        *) return 1 ;;
    esac
    ORIGIN_USER="" ORIGIN_PASSWORD="" ORIGIN_SSLMODE=""
    if [[ "$rest" == *"@"* ]]; then
        userpart="${rest%%@*}"
        rest="${rest#*@}"
        ORIGIN_USER="$(url_decode "${userpart%%:*}")"
        [[ "$userpart" == *":"* ]] && ORIGIN_PASSWORD="$(url_decode "${userpart#*:}")"
    fi
    [[ -n "$ORIGIN_USER" ]] || ORIGIN_USER="${PGUSER:-postgres}"
    hostport="${rest%%/*}"
    dbpart="${rest#*/}"
    [[ "$rest" == */* ]] || dbpart=""
    ORIGIN_DATABASE="${dbpart%%\?*}"
    if [[ "$dbpart" == *"?"* ]]; then
        params="${dbpart#*\?}"
        if [[ "$params" == *"sslmode="* ]]; then
            params="${params#*sslmode=}"
            ORIGIN_SSLMODE="${params%%&*}"
        fi
    fi
    if [[ "$hostport" == *":"* ]]; then
        ORIGIN_HOST="${hostport%%:*}"
        ORIGIN_PORT="${hostport#*:}"
    else
        ORIGIN_HOST="$hostport"
        ORIGIN_PORT="5432"
    fi
    [[ -n "$ORIGIN_HOST" ]]
}

# DATABASE_URL from a .env file in the working directory, quotes stripped.
dotenv_database_url() {
    [[ -f .env ]] || return 1
    local line
    line="$(grep -E '^(export )?DATABASE_URL=' .env 2>/dev/null | head -1)" || return 1
    line="${line#export }"
    line="${line#DATABASE_URL=}"
    line="${line%\"}"; line="${line#\"}"
    line="${line%\'}"; line="${line#\'}"
    [[ -n "$line" ]] && printf '%s' "$line"
}

# Running containers whose image looks like PostgreSQL: "id name image".
postgres_containers() {
    docker ps --format '{{.ID}} {{.Names}} {{.Image}}' 2>/dev/null \
        | grep -Ei ' (postgres|postgis|timescale|pgvector|bitnami/postgresql)' || true
}

container_env() { docker inspect -f '{{range .Config.Env}}{{println .}}{{end}}' "$1" 2>/dev/null | grep "^$2=" | head -1 | cut -d= -f2- || true; }

# Point ORIGIN_* at a container on this Docker host and choose how to reach it.
origin_from_container() {
    local name="$1" network
    ORIGIN_CONTAINER="$name"
    ORIGIN_USER="$(container_env "$name" POSTGRES_USER)"
    ORIGIN_PASSWORD="$(container_env "$name" POSTGRES_PASSWORD)"
    ORIGIN_DATABASE="$(container_env "$name" POSTGRES_DB)"
    [[ -n "$ORIGIN_USER" ]] || ORIGIN_USER="postgres"
    [[ -n "$ORIGIN_DATABASE" ]] || ORIGIN_DATABASE="$ORIGIN_USER"
    ORIGIN_PORT="5432"
    network="$(docker inspect -f '{{range $k, $v := .NetworkSettings.Networks}}{{$k}} {{end}}' "$name" | awk '{print $1}')"
    if [[ -z "$network" || "$network" == "bridge" ]]; then
        # The default bridge has no name resolution, so use the address.
        ORIGIN_HOST="$(docker inspect -f '{{.NetworkSettings.IPAddress}}' "$name")"
        DOCKER_NET_ARGS=()
    else
        ORIGIN_HOST="$name"
        DOCKER_NET_ARGS=(--network "$network")
    fi
    ORIGIN_SOURCE="container $name"
}

origin_discover() {
    local url containers count choice line
    if [[ -n "$UPSTREAM" ]]; then
        url_parse "$UPSTREAM" || die "--upstream must be a postgres://user[:password]@host[:port]/database URL"
        ORIGIN_SOURCE="--upstream"
        return
    fi
    for candidate in UPSTREAM_URL DATABASE_URL; do
        url="${!candidate:-}"
        if [[ -n "$url" ]] && url_parse "$url"; then
            ORIGIN_SOURCE="\$$candidate"
            return
        fi
    done
    if url="$(dotenv_database_url)" && [[ -n "$url" ]] && url_parse "$url"; then
        ORIGIN_SOURCE=".env DATABASE_URL"
        return
    fi
    if [[ -n "${PGHOST:-}" && -n "${PGDATABASE:-}" ]]; then
        ORIGIN_HOST="$PGHOST"
        ORIGIN_PORT="${PGPORT:-5432}"
        ORIGIN_USER="${PGUSER:-${USER:-postgres}}"
        ORIGIN_PASSWORD="${PGPASSWORD:-}"
        ORIGIN_DATABASE="$PGDATABASE"
        ORIGIN_SSLMODE="${PGSSLMODE:-}"
        ORIGIN_SOURCE="PG* environment"
        return
    fi
    containers="$(postgres_containers)"
    if [[ -n "$containers" ]]; then
        count=$(printf '%s\n' "$containers" | wc -l | tr -d ' ')
        say "Running PostgreSQL containers:"
        printf '%s\n' "$containers" | awk '{printf "  %d) %s  (%s)\n", NR, $2, $3}'
        if [[ "$count" -eq 1 ]]; then
            choice=1
        else
            ask choice "Which one is the origin (1-$count, or 0 for another host)" "1"
        fi
        if [[ "$choice" =~ ^[0-9]+$ ]] && [[ "$choice" -ge 1 && "$choice" -le "$count" ]]; then
            line="$(printf '%s\n' "$containers" | sed -n "${choice}p")"
            origin_from_container "$(printf '%s' "$line" | awk '{print $2}')"
            return
        fi
    fi
    ask url "Origin database URL (postgres://user@host:5432/dbname)"
    url_parse "$url" || die "that is not a postgres:// URL"
    ORIGIN_SOURCE="prompt"
}

# A local origin is reachable from a container only as host.docker.internal.
origin_host_rewrite() {
    [[ -n "$ORIGIN_CONTAINER" ]] && return
    case "$ORIGIN_HOST" in
        localhost|127.0.0.1|::1|"[::1]")
            note "Origin host $ORIGIN_HOST becomes host.docker.internal inside the container."
            ORIGIN_HOST="host.docker.internal"
            ;;
    esac
    # Harmless on Docker Desktop, required on Linux for host.docker.internal.
    DOCKER_NET_ARGS=(--add-host=host.docker.internal:host-gateway)
}

origin_url() {
    local url="postgres://${ORIGIN_USER}@${ORIGIN_HOST}:${ORIGIN_PORT}/${ORIGIN_DATABASE}"
    [[ -n "$ORIGIN_SSLMODE" ]] && url="${url}?sslmode=${ORIGIN_SSLMODE}"
    printf '%s' "$url"
}

origin_show() {
    say "Origin (from $ORIGIN_SOURCE):"
    note "$(origin_url)"
    [[ -n "$ORIGIN_PASSWORD" ]] && note "password: (set)"
    if ! confirm "Use this database?"; then
        ask UPSTREAM "Origin database URL"
        url_parse "$UPSTREAM" || die "that is not a postgres:// URL"
        ORIGIN_CONTAINER=""
        ORIGIN_SOURCE="prompt"
        origin_show
    fi
}

# --- ports -------------------------------------------------------------------

port_in_use() { { : > "/dev/tcp/127.0.0.1/$1"; } 2>/dev/null; }

port_free_from() {
    local port="$1"
    while port_in_use "$port"; do port=$((port + 1)); done
    printf '%s' "$port"
}

# --- check -------------------------------------------------------------------

# Images older than this script have no `check` command; their entrypoint
# would treat the word as a proxy argument and start the whole serve path.
image_check_supported() {
    docker run --rm "$IMAGE" --help 2>/dev/null | grep -q '^  check '
}

# The password travels as an inherited environment variable, never in argv.
check_run() {
    ORIGIN_PASSWORD="$ORIGIN_PASSWORD" docker run --rm ${DOCKER_NET_ARGS[@]+"${DOCKER_NET_ARGS[@]}"} \
        -e ORIGIN_PASSWORD "$IMAGE" check --upstream "$(origin_url)"
}

# Offer to switch a Docker-hosted origin to wal_level=logical and restart it.
wal_level_repair_offer() {
    [[ -n "$ORIGIN_CONTAINER" ]] || return 1
    say ""
    say "The origin is a container on this machine, so this script can fix wal_level for you:"
    note "docker exec $ORIGIN_CONTAINER psql -U $ORIGIN_USER -d $ORIGIN_DATABASE -c 'ALTER SYSTEM SET wal_level = logical'"
    note "docker restart $ORIGIN_CONTAINER"
    confirm "Apply this and restart $ORIGIN_CONTAINER?" || return 1
    docker exec "$ORIGIN_CONTAINER" psql -U "$ORIGIN_USER" -d "$ORIGIN_DATABASE" \
        -c "ALTER SYSTEM SET wal_level = logical" > /dev/null
    docker restart "$ORIGIN_CONTAINER" > /dev/null
    say "Waiting for $ORIGIN_CONTAINER to accept connections..."
    until docker exec "$ORIGIN_CONTAINER" pg_isready -U "$ORIGIN_USER" > /dev/null 2>&1; do
        docker inspect -f '{{.State.Running}}' "$ORIGIN_CONTAINER" 2>/dev/null | grep -q true \
            || die "$ORIGIN_CONTAINER stopped; see: docker logs $ORIGIN_CONTAINER"
        sleep 1
    done
    return 0
}

origin_check() {
    local report status
    say ""
    say "Checking the origin..."
    set +e
    report="$(check_run)"
    status=$?
    set -e
    say "$report"
    [[ "$report" == "pgcache check: origin "* ]] \
        || die "unexpected output from '$IMAGE check'; see above"
    if [[ $status -eq 0 ]]; then
        return 0
    fi
    if [[ "$report" == *"asked for a password"* && -z "$ORIGIN_PASSWORD" ]]; then
        ask_secret ORIGIN_PASSWORD "Password for $ORIGIN_USER"
        origin_check
        return
    fi
    if [[ "$report" == *"FAIL  wal_level"* ]] && wal_level_repair_offer; then
        origin_check
        return
    fi
    die "the origin is not ready; fix the failures above and run this again"
}

# --- start / status / stop ---------------------------------------------------

container_start() {
    if container_exists; then
        confirm "A container named $CONTAINER already exists. Replace it?" \
            || die "leaving $CONTAINER as it is (use 'stop' to remove it)"
        docker rm -f "$CONTAINER" > /dev/null
    fi
    PROXY_PORT="$(port_free_from "$PROXY_PORT")"
    METRICS_PORT="$(port_free_from "$METRICS_PORT")"
    say ""
    say "Starting $CONTAINER (proxy on $PROXY_PORT, metrics on $METRICS_PORT)..."
    ORIGIN_PASSWORD="$ORIGIN_PASSWORD" docker run -d --name "$CONTAINER" ${DOCKER_NET_ARGS[@]+"${DOCKER_NET_ARGS[@]}"} \
        -p "127.0.0.1:${PROXY_PORT}:5432" -p "127.0.0.1:${METRICS_PORT}:9090" \
        -e ORIGIN_PASSWORD "$IMAGE" --upstream "$(origin_url)" > /dev/null
    until curl -fs "http://127.0.0.1:${METRICS_PORT}/healthz" > /dev/null 2>&1; do
        if ! docker inspect -f '{{.State.Running}}' "$CONTAINER" 2>/dev/null | grep -q true; then
            say ""
            docker logs --tail 30 "$CONTAINER" 2>&1 | sed 's/^/  /'
            die "$CONTAINER exited during startup; the log tail is above"
        fi
        sleep 1
    done
}

finish_show() {
    local password_part=""
    [[ -n "$ORIGIN_PASSWORD" ]] && password_part=":<password>"
    say ""
    say "PgCache is running. Point your application at it:"
    note "DATABASE_URL=postgres://${ORIGIN_USER}${password_part}@localhost:${PROXY_PORT}/${ORIGIN_DATABASE}"
    say ""
    say "Or try it from psql:"
    note "psql \"postgres://${ORIGIN_USER}@localhost:${PROXY_PORT}/${ORIGIN_DATABASE}\""
    say ""
    say "Watch it work:   curl -fsSL https://pgcache.com/try.sh | bash -s -- status"
    say "Logs:            docker logs -f $CONTAINER"
    say "Stop and remove: curl -fsSL https://pgcache.com/try.sh | bash -s -- stop"
}

status_field() { printf '%s' "$1" | sed -n "s/.*\"$2\":\([0-9]*\).*/\1/p" | head -1; }

status_run() {
    container_exists || die "no container named $CONTAINER; run the script without arguments first"
    local port json registered hits misses total rate
    port="$(docker port "$CONTAINER" 9090/tcp 2>/dev/null | head -1 | sed 's/.*://')"
    [[ -n "$port" ]] || die "$CONTAINER does not publish its metrics port"
    say "Cache status from http://127.0.0.1:${port}/status (Ctrl-C to stop)"
    while :; do
        json="$(curl -fs "http://127.0.0.1:${port}/status" 2>/dev/null || true)"
        if [[ -z "$json" ]]; then
            printf '\r  waiting for %s...                                  ' "$CONTAINER"
        else
            registered="$(status_field "$json" queries_registered)"
            hits="$(status_field "$json" cache_hits)"
            misses="$(status_field "$json" cache_misses)"
            total=$(( ${hits:-0} + ${misses:-0} ))
            rate="-"
            [[ $total -gt 0 ]] && rate="$(( hits * 100 / total ))%"
            printf '\r  queries cached: %-6s hits: %-8s misses: %-8s hit rate: %-5s' \
                "${registered:-0}" "${hits:-0}" "${misses:-0}" "$rate"
        fi
        sleep 1
    done
}

stop_run() {
    container_exists || { say "no container named $CONTAINER"; return 0; }
    docker rm -f "$CONTAINER" > /dev/null
    say "removed $CONTAINER"
}

# --- main --------------------------------------------------------------------

while [[ $# -gt 0 ]]; do
    case "$1" in
        --upstream=*) UPSTREAM="${1#*=}"; shift ;;
        --upstream)   UPSTREAM="${2:-}"; shift 2 ;;
        --port=*)     PROXY_PORT="${1#*=}"; shift ;;
        --port)       PROXY_PORT="${2:-}"; shift 2 ;;
        --image=*)    IMAGE="${1#*=}"; shift ;;
        --image)      IMAGE="${2:-}"; shift 2 ;;
        --check-only) CHECK_ONLY=true; shift ;;
        --yes|-y)     ASSUME_YES=true; shift ;;
        status|stop)  COMMAND="$1"; shift ;;
        -h|--help)    usage; exit 0 ;;
        *) die "unknown argument: $1 (try --help)" ;;
    esac
done

docker_require
case "$COMMAND" in
    status) status_run; exit 0 ;;
    stop)   stop_run; exit 0 ;;
esac

origin_discover
origin_show
origin_host_rewrite

say ""
say "Pulling $IMAGE..."
# A cached image still works when the registry is unreachable (or the tag is local-only).
docker pull -q "$IMAGE" > /dev/null 2>&1 || docker image inspect "$IMAGE" > /dev/null 2>&1 \
    || die "could not pull $IMAGE"
image_check_supported \
    || die "$IMAGE has no 'check' command, so it predates this script. Use a newer image (--image IMAGE)."

origin_check
$CHECK_ONLY && exit 0

container_start
finish_show