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

# ParticleDB Install Script
# =========================
#
# Downloads pre-built binary for your platform and installs to /usr/local/bin.
#
# Usage:
#   curl -fsSL https://particledb.ai/install.sh | bash
#   curl -fsSL https://particledb.ai/install.sh | bash -s -- --uninstall
#
# Supported platforms:
#   macOS (Apple Silicon / Intel)
#   Linux (x86_64 / ARM64 / Graviton)

S3_BASE="https://pdb-releases-193278869920-us-east-1.s3.us-east-1.amazonaws.com/latest"
INSTALL_DIR="/usr/local/bin"
# HOME may be unset in some non-interactive contexts (e.g. systemd, ssm root). Fall back.
DATA_DIR="${HOME:-/root}/.particledb"
BINARY_NAME="particledb"

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

detect_platform() {
    local os arch
    os="$(uname -s)"
    arch="$(uname -m)"

    case "$os" in
        Darwin)
            case "$arch" in
                arm64|aarch64) echo "darwin-arm64" ;;
                x86_64)        echo "darwin-x86_64" ;;
                *) echo "unsupported" ;;
            esac
            ;;
        Linux)
            case "$arch" in
                x86_64|amd64)  echo "linux-x86_64" ;;
                aarch64|arm64) echo "linux-aarch64" ;;
                *) echo "unsupported" ;;
            esac
            ;;
        *) echo "unsupported" ;;
    esac
}

# Detect a faster CPU-tuned SKU when available. Returns either the perf SKU
# name (e.g. linux-x86_64-v4) or the input baseline platform unchanged.
# We never force a perf SKU: if /proc/cpuinfo isn't readable or the marker
# bits are missing, we just stay on the baseline.
detect_perf_variant() {
    local platform="$1"
    case "$platform" in
        linux-x86_64)
            # x86-64-v4 needs avx512f + bmi2 + vaes (rough proxy: avx512f + vaes
            # implies bmi2 in practice on every Intel/AMD core that ships avx512).
            if [ -r /proc/cpuinfo ] \
               && grep -q "avx512f" /proc/cpuinfo \
               && grep -q "bmi2"   /proc/cpuinfo \
               && grep -q "vaes"   /proc/cpuinfo; then
                echo "linux-x86_64-v4"
                return 0
            fi
            ;;
        linux-aarch64)
            # neoverse-v2 / Graviton4 needs sve2 (Armv9-A baseline).
            if [ -r /proc/cpuinfo ] && grep -q "sve2" /proc/cpuinfo; then
                echo "linux-aarch64-neoverse-v2"
                return 0
            fi
            ;;
    esac
    echo "$platform"
}

# Return 0 if the URL is fetchable (HEAD 200), 1 otherwise. Used so we fall
# back to the baseline if a perf SKU GET would 404.
url_exists() {
    local url="$1"
    if command -v curl &>/dev/null; then
        curl -fsSI -o /dev/null --max-time 10 "$url" 2>/dev/null
    elif command -v wget &>/dev/null; then
        wget -q --spider --timeout=10 "$url" 2>/dev/null
    else
        return 1
    fi
}

do_install() {
    echo ""
    echo -e "${BLUE}╔══════════════════════════════════════╗${NC}"
    echo -e "${BLUE}║       ParticleDB Installer           ║${NC}"
    echo -e "${BLUE}╚══════════════════════════════════════╝${NC}"
    echo ""

    local platform
    platform="$(detect_platform)"

    if [ "$platform" = "unsupported" ]; then
        echo -e "${RED}Error: Unsupported platform ($(uname -s) $(uname -m))${NC}"
        echo "Supported: macOS (ARM64/x86_64), Linux (x86_64/ARM64)"
        exit 1
    fi

    echo -e "${YELLOW}Platform:${NC} $platform"

    # Pick a faster SKU if the host CPU supports it AND the perf binary is
    # actually published. Either condition failing → fall back to baseline.
    local effective_platform
    effective_platform="$(detect_perf_variant "$platform")"
    if [ "$effective_platform" != "$platform" ]; then
        local perf_url="${S3_BASE}/particledb-${effective_platform}"
        if url_exists "$perf_url"; then
            echo -e "${YELLOW}CPU variant:${NC} ${effective_platform} (perf SKU available)"
            platform="$effective_platform"
        else
            echo -e "${YELLOW}CPU variant:${NC} ${effective_platform} not published, using baseline ${platform}"
        fi
    fi

    local url="${S3_BASE}/particledb-${platform}"
    local tmp_file
    tmp_file="$(mktemp)"

    echo -e "${YELLOW}Downloading ParticleDB...${NC}"
    if command -v curl &>/dev/null; then
        curl -fSL --progress-bar "$url" -o "$tmp_file"
    elif command -v wget &>/dev/null; then
        wget -q --show-progress "$url" -O "$tmp_file"
    else
        echo -e "${RED}Error: curl or wget required${NC}"
        rm -f "$tmp_file"
        exit 1
    fi

    chmod +x "$tmp_file"

    # Verify the binary runs
    if ! "$tmp_file" --version &>/dev/null 2>&1; then
        # Some builds don't have --version, try --help
        if ! "$tmp_file" --help &>/dev/null 2>&1; then
            echo -e "${YELLOW}Warning: Could not verify binary, installing anyway${NC}"
        fi
    fi

    # Install to /usr/local/bin (may need sudo)
    echo -e "${YELLOW}Installing to ${INSTALL_DIR}/${BINARY_NAME}...${NC}"
    if [ -w "$INSTALL_DIR" ]; then
        mv "$tmp_file" "${INSTALL_DIR}/${BINARY_NAME}"
    else
        sudo mv "$tmp_file" "${INSTALL_DIR}/${BINARY_NAME}"
        sudo chmod +x "${INSTALL_DIR}/${BINARY_NAME}"
    fi

    # Create data directory
    mkdir -p "$DATA_DIR"

    echo ""
    echo -e "${GREEN}ParticleDB installed successfully!${NC}"
    echo ""
    echo -e "  Binary:  ${INSTALL_DIR}/${BINARY_NAME}"
    echo -e "  Data:    ${DATA_DIR}/data"
    echo ""
    echo -e "${BLUE}Quick start:${NC}"
    echo ""
    echo "  # Start the server"
    echo "  particledb start --no-auth --data-dir ${DATA_DIR}/data"
    echo ""
    echo "  # Connect with psql"
    echo "  psql -h 127.0.0.1 -p 5432"
    echo ""
    echo "  # Or use Docker"
    echo "  docker run -p 5432:5432 particledb/particledb"
    echo ""
}

do_uninstall() {
    echo -e "${YELLOW}Uninstalling ParticleDB...${NC}"

    if [ -f "${INSTALL_DIR}/${BINARY_NAME}" ]; then
        if [ -w "$INSTALL_DIR" ]; then
            rm -f "${INSTALL_DIR}/${BINARY_NAME}"
        else
            sudo rm -f "${INSTALL_DIR}/${BINARY_NAME}"
        fi
        echo -e "${GREEN}Binary removed.${NC}"
    else
        echo -e "${YELLOW}Binary not found at ${INSTALL_DIR}/${BINARY_NAME}${NC}"
    fi

    echo ""
    read -p "Remove data directory ($DATA_DIR)? [y/N] " -r
    if [[ "$REPLY" =~ ^[Yy]$ ]]; then
        rm -rf "$DATA_DIR"
        echo -e "${GREEN}Data directory removed.${NC}"
    else
        echo "Data directory preserved."
    fi
}

do_status() {
    echo -e "${BLUE}ParticleDB Status${NC}"
    echo ""

    if [ -f "${INSTALL_DIR}/${BINARY_NAME}" ]; then
        echo -e "  Binary: ${GREEN}installed${NC} (${INSTALL_DIR}/${BINARY_NAME})"
        local size
        size=$(ls -lh "${INSTALL_DIR}/${BINARY_NAME}" | awk '{print $5}')
        echo "  Size:   $size"
    else
        echo -e "  Binary: ${RED}not installed${NC}"
    fi

    if [ -d "$DATA_DIR" ]; then
        echo -e "  Data:   ${GREEN}exists${NC} ($DATA_DIR)"
    else
        echo -e "  Data:   ${YELLOW}not created${NC}"
    fi

    if pgrep -f "particledb start" &>/dev/null; then
        echo -e "  Server: ${GREEN}running${NC}"
    else
        echo -e "  Server: ${YELLOW}stopped${NC}"
    fi
}

# Parse arguments
case "${1:-install}" in
    install)   do_install ;;
    uninstall) do_uninstall ;;
    status)    do_status ;;
    *)
        echo "Usage: $0 {install|uninstall|status}"
        exit 1
        ;;
esac