#!/bin/sh
# GPU CLI Installer
# Usage: curl -fsSL https://gpu-cli.sh | sh
#
# Options (via environment variables):
#   GPU_CLI_VERSION     - Specific version to install (default: latest)
#   GPU_CLI_CHANNEL     - Release channel: stable, beta, nightly (default: stable)
#   GPU_CLI_INSTALL_DIR - Installation directory (default: ~/.gpu-cli)
#   GPU_CLI_NO_MODIFY_PATH - Set to 1 to skip PATH modification

set -e

# Configuration
GCS_BUCKET="gpu-cli-releases"
BASE_URL="https://storage.googleapis.com/${GCS_BUCKET}"
INSTALL_DIR="${GPU_CLI_INSTALL_DIR:-$HOME/.gpu-cli}"
BIN_DIR="$INSTALL_DIR/bin"
POD_BINARIES_DIR="$INSTALL_DIR/pod-binaries"

# Colors (disabled if not a terminal)
if [ -t 1 ]; then
    RED='\033[0;31m'
    GREEN='\033[0;32m'
    YELLOW='\033[0;33m'
    BLUE='\033[0;34m'
    BOLD='\033[1m'
    DIM='\033[2m'
    NC='\033[0m'
    NGREEN='\033[1;32m'
else
    RED=''
    GREEN=''
    YELLOW=''
    BLUE=''
    BOLD=''
    DIM=''
    NC=''
    NGREEN=''
fi

info() {
    printf "${BLUE}==>${NC} ${BOLD}%s${NC}\n" "$1"
}

success() {
    printf "${GREEN}==>${NC} ${BOLD}%s${NC}\n" "$1"
}

warn() {
    printf "${YELLOW}Warning:${NC} %s\n" "$1"
}

error() {
    printf "${RED}Error:${NC} %s\n" "$1" >&2
    exit 1
}

# Animated logo (only if terminal supports it)
show_logo() {
    if [ ! -t 1 ]; then
        printf "${BOLD}GPU CLI${NC}\n"
        return
    fi

    printf "  ${DIM}[${NC}${NGREEN}▓${DIM}░░${NC}${DIM}]${NC} ${DIM}GPU CLI${NC}"
    sleep 0.2
    printf "\r  ${DIM}[${NC}${NGREEN}▓▓${DIM}░${NC}${DIM}]${NC} ${DIM}GPU CLI${NC}"
    sleep 0.2
    printf "\r  ${DIM}[${NC}${NGREEN}▓▓▓${NC}${DIM}]${NC} ${BOLD}GPU CLI${NC}"
    sleep 0.15
    printf "\r  ${NGREEN} ▲▲▲ ${NC} ${BOLD}GPU CLI${NC}\n"
}

validate_channel() {
    case "${GPU_CLI_CHANNEL:-}" in
        stable|beta|nightly|"") ;;
        *) error "Invalid channel: $GPU_CLI_CHANNEL. Use stable, beta, or nightly" ;;
    esac
}

detect_platform() {
    OS=$(uname -s | tr '[:upper:]' '[:lower:]')
    ARCH=$(uname -m)

    case "$OS" in
        darwin) OS="apple-darwin" ;;
        linux) OS="unknown-linux-gnu" ;;
        *) error "Unsupported OS: $OS" ;;
    esac

    case "$ARCH" in
        x86_64|amd64) ARCH="x86_64" ;;
        aarch64|arm64) ARCH="aarch64" ;;
        *) error "Unsupported architecture: $ARCH" ;;
    esac

    TARGET="${ARCH}-${OS}"
    info "Detected platform: $TARGET"

    if [ "$OS" = "apple-darwin" ] && [ "$ARCH" = "x86_64" ]; then
        error "Intel Macs are not supported.
GPU CLI requires Apple Silicon (M1/M2/M3/M4)."
    fi
}

get_version() {
    if [ -n "$GPU_CLI_VERSION" ]; then
        VERSION="$GPU_CLI_VERSION"
        info "Installing specified version: $VERSION"
    else
        CHANNEL="${GPU_CLI_CHANNEL:-stable}"
        info "Fetching latest $CHANNEL version..."
        # Add cache-busting query parameter to avoid stale CDN/proxy caches
        CACHE_BUST=$(date +%s)

        # Try channel-specific path first, fall back to latest/ for stable
        if [ "$CHANNEL" = "stable" ]; then
            VERSION=$(curl -fsSL "${BASE_URL}/channels/stable/version.txt?t=${CACHE_BUST}" 2>/dev/null) || \
            VERSION=$(curl -fsSL "${BASE_URL}/latest/version.txt?t=${CACHE_BUST}") || \
            error "Failed to fetch latest version"
        else
            VERSION=$(curl -fsSL "${BASE_URL}/channels/${CHANNEL}/version.txt?t=${CACHE_BUST}") || \
            error "Failed to fetch ${CHANNEL} version. Channel may not exist."
        fi

        info "Latest $CHANNEL version: $VERSION"
    fi
}

check_dependencies() {
    for cmd in curl tar; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            error "Required command not found: $cmd"
        fi
    done

    if command -v sha256sum >/dev/null 2>&1; then
        CHECKSUM_CMD="sha256sum"
    elif command -v shasum >/dev/null 2>&1; then
        CHECKSUM_CMD="shasum -a 256"
    else
        warn "No checksum tool found. Skipping verification."
        CHECKSUM_CMD=""
    fi
}

download_and_verify() {
    ARCHIVE_NAME="gpu-cli-${TARGET}.tar.gz"
    ARCHIVE_URL="${BASE_URL}/v${VERSION}/${ARCHIVE_NAME}"
    CHECKSUM_URL="${BASE_URL}/v${VERSION}/${ARCHIVE_NAME}.sha256"

    TMP_DIR=$(mktemp -d)
    trap 'rm -rf "$TMP_DIR"' EXIT

    info "Downloading GPU CLI v${VERSION}..."
    curl -fsSL "$ARCHIVE_URL" -o "$TMP_DIR/$ARCHIVE_NAME" || error "Failed to download archive from $ARCHIVE_URL"

    if [ -n "$CHECKSUM_CMD" ]; then
        info "Verifying checksum..."
        EXPECTED_CHECKSUM=$(curl -fsSL "$CHECKSUM_URL") || error "Failed to download checksum from $CHECKSUM_URL"
        cd "$TMP_DIR"
        ACTUAL_CHECKSUM=$($CHECKSUM_CMD "$ARCHIVE_NAME" | cut -d' ' -f1)
        if [ "$EXPECTED_CHECKSUM" != "$ACTUAL_CHECKSUM" ]; then
            error "Checksum verification failed!
Expected: $EXPECTED_CHECKSUM
Actual:   $ACTUAL_CHECKSUM"
        fi
        success "Checksum verified"
    fi

    info "Extracting..."
    cd "$TMP_DIR"
    tar -xzf "$ARCHIVE_NAME"
}

install_binaries() {
    EXTRACT_DIR="$TMP_DIR/gpu-cli-${TARGET}"
    mkdir -p "$BIN_DIR"
    mkdir -p "$POD_BINARIES_DIR"

    info "Installing binaries to $BIN_DIR..."
    cp "$EXTRACT_DIR/gpu" "$BIN_DIR/"
    cp "$EXTRACT_DIR/gpud" "$BIN_DIR/"
    chmod +x "$BIN_DIR/gpu" "$BIN_DIR/gpud"

    if [ -d "$EXTRACT_DIR/pod-binaries" ]; then
        info "Installing pod binaries..."
        cp -r "$EXTRACT_DIR/pod-binaries"/* "$POD_BINARIES_DIR/"
        find "$POD_BINARIES_DIR" -type f -exec chmod +x {} \;
    fi

    echo "$VERSION" > "$INSTALL_DIR/version.txt"
    success "Installed GPU CLI v${VERSION}"
}

detect_shell_rc() {
    case "${SHELL:-}" in
        */zsh)
            if [ -f "$HOME/.zshrc" ]; then echo "$HOME/.zshrc"
            elif [ -f "$HOME/.zprofile" ]; then echo "$HOME/.zprofile"
            else echo "$HOME/.zshrc"
            fi
            ;;
        */bash)
            if [ -f "$HOME/.bash_profile" ]; then echo "$HOME/.bash_profile"
            elif [ -f "$HOME/.bashrc" ]; then echo "$HOME/.bashrc"
            else echo "$HOME/.bash_profile"
            fi
            ;;
        */fish)
            mkdir -p "$HOME/.config/fish"
            echo "$HOME/.config/fish/config.fish"
            ;;
        *)
            echo "$HOME/.profile"
            ;;
    esac
}

setup_path() {
    if [ "${GPU_CLI_NO_MODIFY_PATH:-0}" = "1" ]; then
        return
    fi

    LOCAL_BIN="$HOME/.local/bin"
    mkdir -p "$LOCAL_BIN"
    rm -f "$LOCAL_BIN/gpu" "$LOCAL_BIN/gpud"
    ln -sf "$BIN_DIR/gpu" "$LOCAL_BIN/gpu"
    ln -sf "$BIN_DIR/gpud" "$LOCAL_BIN/gpud"

    RC_FILE=$(detect_shell_rc)

    # Handle oh-my-zsh gpu plugin alias conflict
    # oh-my-zsh's gpu plugin aliases 'gpu' to 'switcherooctl', shadowing our binary
    if [ -f "$RC_FILE" ]; then
        if ! grep -q 'unalias gpu' "$RC_FILE"; then
            # Check if oh-my-zsh is likely active (ZSH variable or .oh-my-zsh dir)
            if [ -d "$HOME/.oh-my-zsh" ] || [ -n "${ZSH:-}" ]; then
                info "Detected oh-my-zsh - adding unalias for gpu"
                {
                    echo ""
                    echo "# GPU CLI: remove oh-my-zsh gpu alias (switcherooctl) so gpu-cli works"
                    echo 'unalias gpu 2>/dev/null || true'
                } >> "$RC_FILE"
                NEEDS_ACTIVATION=true
                ACTIVATION_RC_FILE="$RC_FILE"
            fi
        fi
    fi

    case ":$PATH:" in
        *":$LOCAL_BIN:"*) ;;
        *)
            if [ -f "$RC_FILE" ] && grep -q "\.local/bin.*PATH\|PATH.*\.local/bin" "$RC_FILE"; then
                info "PATH entry already exists in $RC_FILE"
            else
                info "Adding ~/.local/bin to PATH in $RC_FILE"
                {
                    echo ""
                    echo "# Added by GPU CLI installer"
                    echo 'export PATH="$HOME/.local/bin:$PATH"'
                } >> "$RC_FILE"
                success "Added PATH to $RC_FILE"
                NEEDS_ACTIVATION=true
                ACTIVATION_RC_FILE="$RC_FILE"
            fi
            ;;
    esac
}

verify_installation() {
    if [ -x "$BIN_DIR/gpu" ]; then
        INSTALLED_VERSION=$("$BIN_DIR/gpu" --version 2>/dev/null | head -1 || echo "unknown")
        success "GPU CLI installed successfully!"
        echo ""
        echo "  Version: $INSTALLED_VERSION"
        echo "  Location: $BIN_DIR/gpu"
        echo ""
        echo "GPU CLI is free forever, no account required. Bring your own GPU or a"
        echo "provider API key. You pay the provider directly at their rates, zero markup."
        echo ""

        if [ "${NEEDS_ACTIVATION:-false}" = "true" ]; then
            echo "${YELLOW}+----------------------------------------------------------------+${NC}"
            echo "${YELLOW}|${NC} ${BOLD}IMPORTANT: Activate GPU CLI to use it${NC}                        ${YELLOW}|${NC}"
            echo "${YELLOW}+----------------------------------------------------------------+${NC}"
            echo ""
            echo "Run this command now:"
            echo ""
            echo "    ${GREEN}${BOLD}source $ACTIVATION_RC_FILE${NC}"
            echo ""
            echo "Or restart your terminal."
            echo ""
        fi

        echo "Get started:"
        echo "  ${BOLD}gpu auth login${NC}               # Authenticate with RunPod"
        echo "  ${BOLD}gpu run python train.py${NC}      # Run on remote GPU"
        echo ""
    else
        error "Installation verification failed"
    fi
}

# Set update channel in global config
set_update_channel() {
    CHANNEL="${GPU_CLI_CHANNEL:-stable}"

    # Only set config for non-stable channels
    if [ "$CHANNEL" = "stable" ]; then
        return
    fi

    # Config path (same for all platforms - matches Rust PathConfig)
    CONFIG_DIR="$HOME/.gpu-cli"
    CONFIG_FILE="$CONFIG_DIR/config.toml"

    mkdir -p "$CONFIG_DIR"

    if [ -f "$CONFIG_FILE" ]; then
        # Config exists - check if it already has [updates] section
        if grep -q '^\[updates\]' "$CONFIG_FILE"; then
            # Update existing channel setting
            if grep -q '^channel = ' "$CONFIG_FILE"; then
                # Replace existing channel line
                sed -i.bak "s/^channel = .*/channel = \"$CHANNEL\"/" "$CONFIG_FILE" && rm -f "$CONFIG_FILE.bak"
            else
                # Add channel to existing [updates] section
                sed -i.bak "/^\[updates\]/a\\
channel = \"$CHANNEL\"" "$CONFIG_FILE" && rm -f "$CONFIG_FILE.bak"
            fi
        else
            # Append [updates] section
            echo "" >> "$CONFIG_FILE"
            echo "[updates]" >> "$CONFIG_FILE"
            echo "channel = \"$CHANNEL\"" >> "$CONFIG_FILE"
        fi
    else
        # Create new config file
        cat > "$CONFIG_FILE" << EOF
[updates]
channel = "$CHANNEL"
EOF
    fi

    info "Set update channel to $CHANNEL in $CONFIG_FILE"
}

main() {
    echo ""
    show_logo
    echo ""

    check_dependencies
    validate_channel
    detect_platform
    get_version
    download_and_verify
    install_binaries
    set_update_channel
    setup_path
    verify_installation
}

main "$@"