#!/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_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() {
    # Skip animation if not a terminal
    if [ ! -t 1 ]; then
        printf "${BOLD}GPU CLI${NC}\n"
        return
    fi

    # GPU Power Up animation
    # Phase 1: Charging bar
    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

    # Phase 2: Launch
    printf "\r  ${NGREEN} ▲▲▲ ${NC} ${BOLD}GPU CLI${NC}\n"
}

# Detect OS and architecture
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"

    # Check for unsupported Intel Macs
    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 to install
get_version() {
    if [ -n "$GPU_CLI_VERSION" ]; then
        VERSION="$GPU_CLI_VERSION"
        info "Installing specified version: $VERSION"
    else
        info "Fetching latest version..."
        # Add cache-busting query parameter to avoid stale CDN/proxy caches
        CACHE_BUST=$(date +%s)
        VERSION=$(curl -fsSL "${BASE_URL}/latest/version.txt?t=${CACHE_BUST}") || error "Failed to fetch latest version"
        info "Latest version: $VERSION"
    fi
}

# Check for required tools
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

    # Check for checksum tool
    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 (sha256sum or shasum). Skipping verification."
        CHECKSUM_CMD=""
    fi
}

# Download and verify archive
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"

    # Create temp directory
    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"

    # Verify checksum if tool available
    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

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

# Install binaries
install_binaries() {
    EXTRACT_DIR="$TMP_DIR/gpu-cli-${TARGET}"

    # Create installation directories
    mkdir -p "$BIN_DIR"
    mkdir -p "$POD_BINARIES_DIR"

    # Install client binaries
    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"

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

    # Write version file
    echo "$VERSION" > "$INSTALL_DIR/version.txt"

    success "Installed GPU CLI v${VERSION}"
}

# Detect shell RC file
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"  # Create it if needed
            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"
            ;;
        *)
            # Default to .profile for unknown shells
            echo "$HOME/.profile"
            ;;
    esac
}

# Prompt user for input (returns 0 for yes, 1 for no)
prompt_user() {
    PROMPT_MSG="$1"
    DEFAULT="${2:-y}"

    # Try to use /dev/tty for interactive input (works even when piped via curl | sh)
    if [ -t 0 ] || [ -c /dev/tty ]; then
        # Interactive - prompt the user
        if [ "$DEFAULT" = "y" ]; then
            printf "%s [Y/n]: " "$PROMPT_MSG" >/dev/tty
        else
            printf "%s [y/N]: " "$PROMPT_MSG" >/dev/tty
        fi

        # Read from /dev/tty if available, otherwise stdin
        if [ -c /dev/tty ]; then
            read -r REPLY </dev/tty
        else
            read -r REPLY
        fi

        REPLY=$(echo "$REPLY" | tr '[:upper:]' '[:lower:]')

        case "$REPLY" in
            y|yes) return 0 ;;
            n|no) return 1 ;;
            "") [ "$DEFAULT" = "y" ] && return 0 || return 1 ;;
            *) prompt_user "$PROMPT_MSG" "$DEFAULT" ;;
        esac
    else
        # Truly non-interactive (e.g., CI/CD), use default
        echo "$PROMPT_MSG (non-interactive, using default: $DEFAULT)"
        [ "$DEFAULT" = "y" ] && return 0 || return 1
    fi
}

# Check for command conflicts and offer to resolve
check_and_resolve_conflicts() {
    GPU_CONFLICT=""
    CONFLICT_TYPE=""
    RC_FILE=$(detect_shell_rc)

    # First, check RC file for gpu alias (more reliable than runtime check)
    if [ -f "$RC_FILE" ]; then
        if grep -q "alias gpu=" "$RC_FILE" 2>/dev/null; then
            GPU_CONFLICT=$(grep "alias gpu=" "$RC_FILE" | head -1)
            CONFLICT_TYPE="alias"
        fi
    fi

    # Also check for oh-my-zsh git plugin (common source of gpu alias)
    if [ -z "$GPU_CONFLICT" ] && [ -f "$RC_FILE" ]; then
        if grep -q "plugins=.*git" "$RC_FILE" 2>/dev/null && [ -d "$HOME/.oh-my-zsh" ]; then
            GPU_CONFLICT="oh-my-zsh git plugin (gpu='git push upstream')"
            CONFLICT_TYPE="alias"
        fi
    fi

    # Fall back to runtime checks (might not work in subshell, but try anyway)
    if [ -z "$GPU_CONFLICT" ]; then
        if type gpu 2>/dev/null | grep -q "function"; then
            GPU_CONFLICT="shell function"
            CONFLICT_TYPE="function"
        elif command -v gpu >/dev/null 2>&1; then
            # Only treat as conflict if it's not our own binary
            GPU_PATH=$(command -v gpu)
            if [ "$GPU_PATH" != "$LOCAL_BIN/gpu" ] && [ "$GPU_PATH" != "$BIN_DIR/gpu" ]; then
                GPU_CONFLICT="$GPU_PATH"
                CONFLICT_TYPE="command"
            fi
        fi
    fi

    # If conflict found, handle it
    if [ -n "$GPU_CONFLICT" ]; then
        echo ""
        warn "Conflict detected: 'gpu' $CONFLICT_TYPE already exists"

        if [ "$CONFLICT_TYPE" = "alias" ]; then
            echo "  $GPU_CONFLICT"
            echo ""
            echo "This is commonly from oh-my-zsh's git plugin (gpu = 'git push upstream')."
            echo "GPU CLI needs the 'gpu' command for running code on remote GPUs."
            echo ""

            if prompt_user "Override the existing 'gpu' alias?"; then
                OVERRIDE_GPU_ALIAS=true
                success "Will add 'unalias gpu' to your shell config"
            else
                OVERRIDE_GPU_ALIAS=false
                echo ""
                info "GPU CLI installed. To use it, either:"
                echo "  1. Remove the 'gpu' alias from your shell config, then restart your shell"
                echo "  2. Use the full path: $LOCAL_BIN/gpu"
            fi
        else
            echo "  Found: $GPU_CONFLICT"
            echo ""
            warn "Cannot automatically resolve $CONFLICT_TYPE conflict."
            info "You'll need to use the full path: $LOCAL_BIN/gpu"
            OVERRIDE_GPU_ALIAS=false
        fi
        echo ""
    else
        OVERRIDE_GPU_ALIAS=false
    fi
}

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

    # Create symlinks in ~/.local/bin (standard location)
    LOCAL_BIN="$HOME/.local/bin"
    mkdir -p "$LOCAL_BIN"

    # Remove old symlinks if they exist
    rm -f "$LOCAL_BIN/gpu" "$LOCAL_BIN/gpud"

    # Create new symlinks
    ln -sf "$BIN_DIR/gpu" "$LOCAL_BIN/gpu"
    ln -sf "$BIN_DIR/gpud" "$LOCAL_BIN/gpud"

    # Detect shell RC file once
    RC_FILE=$(detect_shell_rc)
    RC_FILE_MODIFIED=false

    # Check for conflicts first
    check_and_resolve_conflicts

    # Check if ~/.local/bin is in PATH
    PATH_IN_SHELL=false
    case ":$PATH:" in
        *":$LOCAL_BIN:"*) PATH_IN_SHELL=true ;;
    esac

    # Add PATH to RC file if needed (automatic, no prompt)
    if [ "$PATH_IN_SHELL" = "false" ]; then
        if [ -f "$RC_FILE" ] && grep -q "\.local/bin.*PATH" "$RC_FILE"; then
            info "PATH entry already exists in $RC_FILE"
        else
            echo ""
            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"
            RC_FILE_MODIFIED=true
        fi
    fi

    # Add unalias if needed
    if [ "$OVERRIDE_GPU_ALIAS" = "true" ]; then
        add_unalias_to_rc "$RC_FILE"
        RC_FILE_MODIFIED=true
    fi

    # Set global variables for verification step
    if [ "$RC_FILE_MODIFIED" = "true" ]; then
        NEEDS_ACTIVATION=true
        ACTIVATION_RC_FILE="$RC_FILE"
    fi
}

# Add unalias to RC file
add_unalias_to_rc() {
    RC_FILE="$1"

    # Check if unalias already exists
    if [ -f "$RC_FILE" ] && grep -q "unalias gpu" "$RC_FILE"; then
        info "'unalias gpu' already in $RC_FILE"
    else
        echo ""
        info "Adding 'unalias gpu' to $RC_FILE"
        {
            echo ""
            echo "# GPU CLI: Override conflicting alias"
            echo "unalias gpu 2>/dev/null || true"
        } >> "$RC_FILE"
        success "Added 'unalias gpu' to $RC_FILE"
    fi

    # Apply to current shell immediately (best effort)
    # Note: This only works if installer is sourced, not piped via curl
    unalias gpu 2>/dev/null || true
}

# Verify installation
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 ""

        # Show activation instructions if RC file was modified
        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 ""
            echo "Then get started with:"
            echo "  ${BOLD}gpu auth login${NC}               # Authenticate with RunPod"
            echo "  ${BOLD}gpu run python train.py${NC}      # Run on remote GPU"
            echo ""
        else
            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 ""
        fi
    else
        error "Installation verification failed"
    fi
}

# Main
main() {
    echo ""
    show_logo
    echo ""

    check_dependencies
    detect_platform
    get_version
    download_and_verify
    install_binaries
    setup_path
    verify_installation
}

main "$@"