#!/bin/bash

###############################################################################
# Layer 1A: VPS Basic Setup
# 
# This script configures users, SSH, firewall, and fail2ban on VPS machines.
# Runs independently - can be executed without Nodito setup.
###############################################################################

set -e  # Exit on error

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

# Project root directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
ANSIBLE_DIR="$PROJECT_ROOT/ansible"

###############################################################################
# Helper Functions
###############################################################################

print_header() {
    echo -e "\n${BLUE}========================================${NC}"
    echo -e "${BLUE}$1${NC}"
    echo -e "${BLUE}========================================${NC}\n"
}

print_success() {
    echo -e "${GREEN}✓${NC} $1"
}

print_error() {
    echo -e "${RED}✗${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}⚠${NC} $1"
}

print_info() {
    echo -e "${BLUE}ℹ${NC} $1"
}

confirm_action() {
    local prompt="$1"
    local response
    
    read -p "$(echo -e ${YELLOW}${prompt}${NC} [y/N]: )" response
    [[ "$response" =~ ^[Yy]$ ]]
}

###############################################################################
# Verification Functions
###############################################################################

check_layer_0_complete() {
    print_header "Verifying Layer 0 Prerequisites"
    
    local errors=0
    
    # Check if venv exists
    if [ ! -d "$PROJECT_ROOT/venv" ]; then
        print_error "Python venv not found. Run Layer 0 first."
        ((errors++))
    else
        print_success "Python venv exists"
    fi
    
    # Check if we're in a venv
    if [ -z "$VIRTUAL_ENV" ]; then
        print_error "Virtual environment not activated"
        echo "Run: source venv/bin/activate"
        ((errors++))
    else
        print_success "Virtual environment activated"
    fi
    
    # Check if Ansible is installed
    if ! command -v ansible &> /dev/null; then
        print_error "Ansible not found"
        ((errors++))
    else
        print_success "Ansible found: $(ansible --version | head -n1)"
    fi
    
    # Check if inventory.ini exists
    if [ ! -f "$ANSIBLE_DIR/inventory.ini" ]; then
        print_error "inventory.ini not found"
        ((errors++))
    else
        print_success "inventory.ini exists"
    fi
    
    # Check if infra_vars.yml exists
    if [ ! -f "$ANSIBLE_DIR/infra_vars.yml" ]; then
        print_error "infra_vars.yml not found"
        ((errors++))
    else
        print_success "infra_vars.yml exists"
    fi
    
    if [ $errors -gt 0 ]; then
        print_error "Layer 0 is not complete. Please run ./scripts/setup_layer_0.sh first"
        exit 1
    fi
    
    print_success "Layer 0 prerequisites verified"
}

get_hosts_from_inventory() {
    local group="$1"
    cd "$ANSIBLE_DIR"
    ansible-inventory -i inventory.ini --list | \
        python3 -c "import sys, json; data=json.load(sys.stdin); print(' '.join(data.get('$group', {}).get('hosts', [])))" 2>/dev/null || echo ""
}

check_vps_configured() {
    print_header "Checking VPS Configuration"
    
    local has_vps=false
    for group in vipy watchtower spacey; do
        local hosts=$(get_hosts_from_inventory "$group")
        if [ -n "$hosts" ]; then
            print_success "$group configured: $hosts"
            has_vps=true
        else
            print_info "$group not configured (skipping)"
        fi
    done
    
    if [ "$has_vps" = false ]; then
        print_error "No VPSs configured in inventory.ini"
        print_info "Add at least one VPS (vipy, watchtower, or spacey) to proceed"
        exit 1
    fi
    
    echo ""
}

check_ssh_connectivity() {
    print_header "Testing SSH Connectivity as Root"
    
    local ssh_key=$(grep "ansible_ssh_private_key_file" "$ANSIBLE_DIR/inventory.ini" | head -n1 | sed 's/.*ansible_ssh_private_key_file=\([^ ]*\).*/\1/')
    ssh_key="${ssh_key/#\~/$HOME}"
    
    print_info "Using SSH key: $ssh_key"
    echo ""
    
    local all_good=true
    
    # Test VPSs (vipy, watchtower, spacey)
    for group in vipy watchtower spacey; do
        local hosts=$(get_hosts_from_inventory "$group")
        if [ -n "$hosts" ]; then
            for host in $hosts; do
                print_info "Testing SSH to $host as root..."
                if timeout 10 ssh -i "$ssh_key" -o StrictHostKeyChecking=no -o BatchMode=yes root@$host "echo 'SSH OK'" &>/dev/null; then
                    print_success "SSH to $host as root: OK"
                else
                    print_error "Cannot SSH to $host as root"
                    print_warning "Make sure your SSH key is added to root on $host"
                    all_good=false
                fi
            done
        fi
    done
    
    if [ "$all_good" = false ]; then
        echo ""
        print_error "SSH connectivity test failed"
        print_info "To fix this:"
        echo "  1. Ensure your VPS provider has added your SSH key to root"
        echo "  2. Test manually: ssh -i $ssh_key root@<host>"
        echo ""
        if ! confirm_action "Continue anyway?"; then
            exit 1
        fi
    fi
    
    echo ""
    print_success "SSH connectivity verified"
}

###############################################################################
# VPS Setup Functions
###############################################################################

setup_vps_users_and_access() {
    print_header "Setting Up Users and SSH Access on VPSs"
    
    cd "$ANSIBLE_DIR"
    
    print_info "This will:"
    echo "  • Create the 'counterweight' user with sudo access"
    echo "  • Configure SSH key authentication"
    echo "  • Disable root login (optional, configured in playbook)"
    echo ""
    print_info "Running: ansible-playbook -i inventory.ini infra/01_user_and_access_setup_playbook.yml"
    echo ""
    
    if ! confirm_action "Proceed with user and access setup?"; then
        print_warning "Skipped user and access setup"
        return 1
    fi
    
    # Run the playbook with -e 'ansible_user="root"' to use root for this first run
    if ansible-playbook -i inventory.ini infra/01_user_and_access_setup_playbook.yml -e 'ansible_user="root"'; then
        print_success "User and access setup complete"
        return 0
    else
        print_error "User and access setup failed"
        return 1
    fi
}

setup_vps_firewall_and_fail2ban() {
    print_header "Setting Up Firewall and Fail2ban on VPSs"
    
    cd "$ANSIBLE_DIR"
    
    print_info "This will:"
    echo "  • Configure UFW firewall with SSH access"
    echo "  • Install and configure fail2ban for brute force protection"
    echo "  • Install and configure auditd for security logging"
    echo ""
    print_info "Running: ansible-playbook -i inventory.ini infra/02_firewall_and_fail2ban_playbook.yml"
    echo ""
    
    if ! confirm_action "Proceed with firewall and fail2ban setup?"; then
        print_warning "Skipped firewall setup"
        return 1
    fi
    
    # Now use the default counterweight user
    if ansible-playbook -i inventory.ini infra/02_firewall_and_fail2ban_playbook.yml; then
        print_success "Firewall and fail2ban setup complete"
        return 0
    else
        print_error "Firewall setup failed"
        return 1
    fi
}

###############################################################################
# Verification Functions
###############################################################################

verify_layer_1a() {
    print_header "Verifying Layer 1A Completion"
    
    cd "$ANSIBLE_DIR"
    
    local ssh_key=$(grep "ansible_ssh_private_key_file" "$ANSIBLE_DIR/inventory.ini" | head -n1 | sed 's/.*ansible_ssh_private_key_file=\([^ ]*\).*/\1/')
    ssh_key="${ssh_key/#\~/$HOME}"
    
    # Test SSH as counterweight user
    print_info "Testing SSH as counterweight user..."
    echo ""
    
    local all_good=true
    
    for group in vipy watchtower spacey; do
        local hosts=$(get_hosts_from_inventory "$group")
        if [ -n "$hosts" ]; then
            for host in $hosts; do
                if timeout 10 ssh -i "$ssh_key" -o StrictHostKeyChecking=no -o BatchMode=yes counterweight@$host "echo 'SSH OK'" &>/dev/null; then
                    print_success "SSH to $host as counterweight: OK"
                else
                    print_error "Cannot SSH to $host as counterweight"
                    all_good=false
                fi
            done
        fi
    done
    
    echo ""
    if [ "$all_good" = true ]; then
        print_success "All SSH connectivity verified"
    else
        print_warning "Some SSH tests failed - manual verification recommended"
        print_info "Test manually: ssh -i $ssh_key counterweight@<host>"
    fi
}

###############################################################################
# Summary Functions
###############################################################################

print_summary() {
    print_header "Layer 1A: VPS Setup Complete! 🎉"
    
    echo "Summary of what was configured:"
    echo ""
    print_success "counterweight user created on all VPSs"
    print_success "SSH key authentication configured"
    print_success "UFW firewall active and configured"
    print_success "fail2ban protecting against brute force attacks"
    print_success "auditd logging security events"
    echo ""
    
    print_warning "Important Security Changes:"
    echo "  • Root SSH login is now disabled (by design)"
    echo "  • Always use 'counterweight' user for SSH access"
    echo "  • Firewall is active - only SSH allowed by default"
    echo ""
    
    print_info "Next steps:"
    echo "  1. Test SSH access: ssh -i ~/.ssh/counterganzua counterweight@<host>"
    echo "  2. (Optional) Set up Nodito: ./scripts/setup_layer_1b_nodito.sh"
    echo "  3. Proceed to Layer 2: ./scripts/setup_layer_2.sh"
    echo ""
}

###############################################################################
# Main Execution
###############################################################################

main() {
    clear
    
    print_header "🔧 Layer 1A: VPS Basic Setup"
    
    echo "This script will configure users, SSH, firewall, and fail2ban on VPS machines."
    echo ""
    print_info "Targets: vipy, watchtower, spacey"
    echo ""
    
    if ! confirm_action "Continue with Layer 1A setup?"; then
        echo "Setup cancelled."
        exit 0
    fi
    
    check_layer_0_complete
    check_vps_configured
    check_ssh_connectivity
    
    # VPS Setup
    local setup_failed=false
    setup_vps_users_and_access || setup_failed=true
    setup_vps_firewall_and_fail2ban || setup_failed=true
    
    verify_layer_1a
    
    if [ "$setup_failed" = true ]; then
        print_warning "Some steps failed - please review errors above"
    fi
    
    print_summary
}

# Run main function
main "$@"