#!/bin/bash
#
# VM Post-Installation Provisioning Script
# Handles: System updates, SSH key deployment, Docker installation
#
# Usage: ./vm-post-install.sh <vm-ip> [ssh-user]
# Example: ./vm-post-install.sh 10.10.0.100 cal
#

set -e

# Configuration
VM_IP="$1"
SSH_USER="${2:-cal}"
HOMELAB_KEY="$HOME/.ssh/homelab_rsa"
EMERGENCY_KEY="$HOME/.ssh/emergency_homelab_rsa"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Logging function
log() {
    echo -e "${GREEN}[$(date '+%H:%M:%S')]${NC} $1"
}

warn() {
    echo -e "${YELLOW}[$(date '+%H:%M:%S')] WARNING:${NC} $1"
}

error() {
    echo -e "${RED}[$(date '+%H:%M:%S')] ERROR:${NC} $1"
    exit 1
}

# Validation
if [ -z "$VM_IP" ]; then
    error "Usage: $0 <vm-ip> [ssh-user]"
fi

if [ ! -f "$HOMELAB_KEY" ]; then
    error "Homelab SSH key not found at $HOMELAB_KEY"
fi

if [ ! -f "$EMERGENCY_KEY" ]; then
    error "Emergency SSH key not found at $EMERGENCY_KEY"
fi

log "Starting VM provisioning for $VM_IP as user $SSH_USER"

# Test initial connectivity (assume password auth is still enabled)
log "Testing initial connectivity to $VM_IP..."
if ! timeout 10 bash -c "nc -z $VM_IP 22" >/dev/null 2>&1; then
    error "Cannot reach $VM_IP on port 22"
fi

# Function to run commands on remote VM
run_remote() {
    local cmd="$1"
    local use_key="${2:-false}"
    
    if [ "$use_key" = "true" ]; then
        ssh -i "$HOMELAB_KEY" -o StrictHostKeyChecking=accept-new "$SSH_USER@$VM_IP" "$cmd"
    else
        # Initial connection might need password
        ssh -o StrictHostKeyChecking=accept-new "$SSH_USER@$VM_IP" "$cmd"
    fi
}

# Function to copy files to remote VM
copy_to_remote() {
    local src="$1"
    local dest="$2"
    local use_key="${3:-false}"
    
    if [ "$use_key" = "true" ]; then
        scp -i "$HOMELAB_KEY" -o StrictHostKeyChecking=accept-new "$src" "$SSH_USER@$VM_IP:$dest"
    else
        scp -o StrictHostKeyChecking=accept-new "$src" "$SSH_USER@$VM_IP:$dest"
    fi
}

# Step 1: System Updates
log "Step 1: Running system updates..."
run_remote "
    echo 'Updating package lists...'
    sudo apt update
    echo 'Upgrading packages...'
    sudo apt upgrade -y
    echo 'Installing essential packages...'
    sudo apt install -y curl wget git vim htop unzip software-properties-common apt-transport-https ca-certificates gnupg lsb-release
    echo 'Cleaning up...'
    sudo apt autoremove -y
    sudo apt autoclean
"
log "✅ System updates completed"

# Step 2: SSH Key Installation
log "Step 2: Installing SSH keys..."

# Create .ssh directory if it doesn't exist
run_remote "mkdir -p ~/.ssh && chmod 700 ~/.ssh"

# Deploy primary homelab key
log "Installing primary homelab SSH key..."
cat "$HOMELAB_KEY.pub" | run_remote "cat >> ~/.ssh/authorized_keys"

# Deploy emergency key
log "Installing emergency SSH key..."
cat "$EMERGENCY_KEY.pub" | run_remote "cat >> ~/.ssh/authorized_keys"

# Set proper permissions
run_remote "chmod 600 ~/.ssh/authorized_keys"

# Test key-based authentication
log "Testing key-based authentication..."
if run_remote "echo 'SSH key authentication working'" true; then
    log "✅ SSH key authentication successful"
else
    error "SSH key authentication failed"
fi

# Step 3: Secure SSH Configuration
log "Step 3: Securing SSH configuration..."
run_remote "
    sudo cp /etc/ssh/sshd_config /etc/ssh/sshd_config.backup
    sudo sed -i 's/#PasswordAuthentication yes/PasswordAuthentication no/' /etc/ssh/sshd_config
    sudo sed -i 's/PasswordAuthentication yes/PasswordAuthentication no/' /etc/ssh/sshd_config
    sudo sed -i 's/#PubkeyAuthentication yes/PubkeyAuthentication yes/' /etc/ssh/sshd_config
    sudo sed -i 's/PubkeyAuthentication no/PubkeyAuthentication yes/' /etc/ssh/sshd_config
    sudo systemctl restart sshd
" true
log "✅ SSH hardened (password authentication disabled)"

# Step 4: Docker Installation
log "Step 4: Installing Docker and Docker Compose..."
run_remote "
    # Remove any old Docker installations
    sudo apt remove -y docker docker-engine docker.io containerd runc 2>/dev/null || true
    
    # Add Docker GPG key
    curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg
    
    # Add Docker repository
    echo \"deb [arch=\$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu \$(lsb_release -cs) stable\" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
    
    # Update package list
    sudo apt update
    
    # Install Docker
    sudo apt install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
    
    # Add user to docker group
    sudo usermod -aG docker $USER
    
    # Enable Docker service
    sudo systemctl enable docker
    sudo systemctl start docker
" true

# Verify Docker installation
log "Verifying Docker installation..."
run_remote "
    docker --version
    docker compose version
    sudo docker run --rm hello-world
" true
log "✅ Docker and Docker Compose installed successfully"

# Step 5: System Configuration
log "Step 5: Additional system configuration..."
run_remote "
    # Configure automatic security updates
    sudo apt install -y unattended-upgrades
    sudo dpkg-reconfigure -plow unattended-upgrades
    
    # Set timezone (adjust as needed)
    sudo timedatectl set-timezone America/New_York
    
    # Configure system hostname if needed
    echo 'System configuration completed'
" true

# Step 6: Create useful aliases and environment
log "Step 6: Setting up user environment..."
run_remote "
    # Create useful aliases
    cat >> ~/.bashrc << 'EOF'

# Docker aliases
alias dps='docker ps'
alias dlog='docker logs'
alias dexec='docker exec -it'
alias dstop='docker stop \$(docker ps -q)'
alias dprune='docker system prune -f'

# System aliases  
alias ll='ls -alF'
alias la='ls -A'
alias l='ls -CF'
alias ..='cd ..'
alias ...='cd ../..'

# Docker Compose aliases
alias dc='docker compose'
alias dcup='docker compose up -d'
alias dcdown='docker compose down'
alias dclogs='docker compose logs -f'
EOF

    # Source the updated bashrc
    source ~/.bashrc 2>/dev/null || true
" true

log "✅ User environment configured"

# Step 7: Final verification
log "Step 7: Final system verification..."
HOSTNAME=$(run_remote "hostname" true)
KERNEL=$(run_remote "uname -r" true)
DOCKER_VERSION=$(run_remote "docker --version" true)
UPTIME=$(run_remote "uptime" true)

log "🎉 VM provisioning completed successfully!"
echo ""
echo -e "${BLUE}=== PROVISIONING SUMMARY ===${NC}"
echo -e "${GREEN}VM IP:${NC} $VM_IP"
echo -e "${GREEN}User:${NC} $SSH_USER"
echo -e "${GREEN}Hostname:${NC} $HOSTNAME"
echo -e "${GREEN}Kernel:${NC} $KERNEL"
echo -e "${GREEN}Docker:${NC} $DOCKER_VERSION"
echo -e "${GREEN}Uptime:${NC} $UPTIME"
echo ""
echo -e "${YELLOW}Next Steps:${NC}"
echo "1. SSH access: ssh $SSH_USER@$VM_IP"
echo "2. Test Docker: ssh $SSH_USER@$VM_IP 'docker run --rm hello-world'"
echo "3. Deploy applications using Docker Compose"
echo ""
echo -e "${YELLOW}Security Notes:${NC}"
echo "- Password authentication is DISABLED"
echo "- Only key-based SSH access is allowed"
echo "- Both primary and emergency SSH keys are installed"
echo "- Automatic security updates are enabled"