#!/bin/bash
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

set -eu

# Use WORKSPACE_DIR if set (should be set by Dockerfile.local_dev)
# Otherwise, use /workspace as fallback
WORKSPACE_DIR="${WORKSPACE_DIR:-/workspace}"

# Ensure we're not running as root
if [ "$(id -u)" -eq 0 ]; then
    echo "❌ ERROR: This script should not be run as root!"
    echo "The script should run as the 'dynamo' user, not root."
    echo "Current user: $(whoami) (UID: $(id -u))"
    exit 1
fi

# Verify we're running as the expected user
if [ "$(whoami)" != "dynamo" ]; then
    echo "⚠️  WARNING: Expected to run as 'dynamo' user, but running as '$(whoami)'"
    echo "This might cause permission issues."
fi

echo "Running post-create script as user: $(whoami) (UID: $(id -u))"

trap 'echo "❌ ERROR: Command failed at line $LINENO: $BASH_COMMAND"; echo "⚠️ This was unexpected and setup was not completed. Can try to resolve yourself and then manually run the rest of the commands in this file or file a bug."' ERR

retry() {
    # retries for connectivity issues in installs
    local retries=3
    local count=0
    until "$@"; do
        exit_code=$?
        wait_time=$((2 ** count))
        echo "Command failed with exit code $exit_code. Retrying in $wait_time seconds..."
        sleep $wait_time
        count=$((count + 1))
        if [ $count -ge $retries ]; then
            echo "Command failed after $retries attempts."
            return $exit_code
        fi
    done
    return 0
}

show_and_run() {
    # Run commands with debug output shown
    set -x
    "$@"
    local exit_code=$?
    { set +x; } 2>/dev/null
    return $exit_code
}

set -x

# Pre-commit hooks
cd $WORKSPACE_DIR && pre-commit install && retry pre-commit install-hooks
pre-commit run --all-files || true # don't fail the build if pre-commit hooks fail

# Use CARGO_TARGET_DIR from environment (should be set by Dockerfile.local_dev)
# Fallback to $WORKSPACE_DIR/target if not set
export CARGO_TARGET_DIR=${CARGO_TARGET_DIR:-$WORKSPACE_DIR/target}
mkdir -p $CARGO_TARGET_DIR

# Note: Build steps moved to after sanity check - see instructions at the end

{ set +x; } 2>/dev/null

echo -e "\n" >> ~/.bashrc
echo "# === This section is generated by the post-create.sh script ===" >> ~/.bashrc

if ! grep -q "export GPG_TTY=" ~/.bashrc; then
    show_and_run echo 'export GPG_TTY=$(tty)' >> ~/.bashrc
fi

# Unset empty tokens/variables to avoid issues with authentication and SSH
if ! grep -q "# Unset empty tokens" ~/.bashrc; then
    echo -e "\n# Unset empty tokens and environment variables" >> ~/.bashrc
    echo '[ -z "${HF_TOKEN:-}" ] && unset HF_TOKEN' >> ~/.bashrc
    echo '[ -z "${GITHUB_TOKEN:-}" ] && unset GITHUB_TOKEN' >> ~/.bashrc
    echo '[ -z "${SSH_AUTH_SOCK:-}" ] && unset SSH_AUTH_SOCK' >> ~/.bashrc
fi

# Check SSH agent forwarding status
if [ -n "${SSH_AUTH_SOCK:-}" ]; then
    if ssh-add -l > /dev/null 2>&1; then
        echo "SSH agent forwarding is working - found $(ssh-add -l | wc -l) key(s):"
        ssh-add -l
    else
        echo "⚠️ SSH_AUTH_SOCK is set but ssh-add failed - agent may not be accessible"
    fi
else
    echo "⚠️ SSH agent forwarding not configured - SSH_AUTH_SOCK is not set"
fi

cat <<EOF

========================================
✅ Pre-commit hooks configured

Now build the project:
  cargo build --features dynamo-llm/block-manager
  cd lib/bindings/python && maturin develop --uv && cd $WORKSPACE_DIR
  uv pip install --no-deps -e $WORKSPACE_DIR

Optional:
  uv pip install -e lib/gpu_memory_service  # GPU memory manager with C++ extension
  cd lib/bindings/kvbm && maturin develop --uv  # For KVBM support

Verify with:
  $WORKSPACE_DIR/deploy/sanity_check.py
========================================
EOF