#!/bin/bash

# Reset and Test Script for Debian Vagrant Cluster
# This script fully destroys and recreates the entire stack, then runs all tests

set -e  # Exit on any 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

# Function to print colored output
print_header() {
    echo -e "${BLUE}=== $1 ===${NC}"
}

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

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

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

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

# Function to check if command exists
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# Function to wait for user confirmation (disabled for automation)
confirm() {
    print_info "Auto-confirming operation (automation mode)"
    return 0
}

# Function to check prerequisites
check_prerequisites() {
    print_header "Checking Prerequisites"
    
    local missing_deps=()
    
    if ! command_exists vagrant; then
        missing_deps+=("vagrant")
    fi
    
    if ! command_exists ansible; then
        missing_deps+=("ansible")
    fi
    
    if ! command_exists make; then
        missing_deps+=("make")
    fi
    
    if [ ${#missing_deps[@]} -ne 0 ]; then
        print_error "Missing required dependencies: ${missing_deps[*]}"
        print_info "Please install the missing dependencies and try again"
        exit 1
    fi
    
    print_success "All prerequisites are installed"
}

# Function to destroy everything
destroy_all() {
    print_header "Destroying All Machines"
    
    print_info "Auto-confirming destruction (automation mode)"
    
    print_info "Stopping all machines..."
    vagrant halt 2>/dev/null || true
    
    print_info "Destroying all machines..."
    vagrant destroy -f
    
    print_info "Cleaning up Vagrant files..."
    rm -rf .vagrant/
    
    print_success "All machines destroyed and cleaned up"
}

# Function to start all machines
start_all() {
    print_header "Starting All Machines"
    
    print_info "Starting Vagrant cluster..."
    vagrant up
    
    print_info "Waiting for machines to be ready..."
    sleep 10
    
    print_info "Checking machine status..."
    vagrant status
    
    print_success "All machines started successfully"
}

# Function to test SSH connectivity
test_ssh() {
    print_header "Testing SSH Connectivity"
    
    local machines=("swarm-manager" "swarm-worker1" "swarm-worker2" "swarm-worker3")
    local failed_machines=()
    
    for machine in "${machines[@]}"; do
        print_info "Testing SSH to $machine..."
        if vagrant ssh "$machine" -c "echo 'SSH test successful'" >/dev/null 2>&1; then
            print_success "SSH to $machine: OK"
        else
            print_error "SSH to $machine: FAILED"
            failed_machines+=("$machine")
        fi
    done
    
    if [ ${#failed_machines[@]} -ne 0 ]; then
        print_error "SSH failed for: ${failed_machines[*]}"
        return 1
    fi
    
    print_success "All SSH connections working"
    return 0
}

# Function to test Ansible connectivity
test_ansible() {
    print_header "Testing Ansible Connectivity"
    
    print_info "Running Ansible ping test..."
    if ansible all -i inventory -m ping; then
        print_success "Ansible connectivity test passed"
        return 0
    else
        print_error "Ansible connectivity test failed"
        return 1
    fi
}

# Function to run setup playbook
run_setup() {
    print_header "Running Setup Playbook"
    
    print_info "Installing Ansible roles..."
    if ansible-galaxy install -r ansible-requirements.yml --force; then
        print_success "Ansible roles installed successfully"
    else
        print_error "Failed to install Ansible roles"
        return 1
    fi
    
    print_info "Installing dependencies and creating swap..."
    if ansible-playbook -i inventory setup-playbook.yml; then
        print_success "Setup playbook completed successfully"
        return 0
    else
        print_error "Setup playbook failed"
        return 1
    fi
}

# Function to run deployment playbook
run_deployment() {
    print_header "Running Deployment Playbook"
    
    print_info "Deploying applications and services..."
    if ansible-playbook -i inventory deploy-playbook.yml; then
        print_success "Deployment playbook completed successfully"
        return 0
    else
        print_error "Deployment playbook failed"
        return 1
    fi
}

# Function to run comprehensive tests
run_tests() {
    print_header "Running Comprehensive Tests"
    
    local test_results=()
    
    # Test 1: SSH Connectivity
    if test_ssh; then
        test_results+=("SSH: ✅ PASS")
    else
        test_results+=("SSH: ❌ FAIL")
    fi
    
    # Test 2: Ansible Connectivity
    if test_ansible; then
        test_results+=("Ansible: ✅ PASS")
    else
        test_results+=("Ansible: ❌ FAIL")
    fi
    
    # Test 3: Setup Playbook
    if run_setup; then
        test_results+=("Setup: ✅ PASS")
    else
        test_results+=("Setup: ❌ FAIL")
    fi
    
    # Test 4: Deployment Playbook
    if run_deployment; then
        test_results+=("Deployment: ✅ PASS")
    else
        test_results+=("Deployment: ❌ FAIL")
    fi
    
    # Test 5: Verify swap is active
    print_info "Verifying swap is active..."
    if ansible all -i inventory -m shell -a "cat /proc/swaps" | grep -q "swapfile"; then
        test_results+=("Swap: ✅ PASS")
    else
        test_results+=("Swap: ❌ FAIL")
    fi
    
    # Test 6: Verify Docker is running
    print_info "Verifying Docker is running..."
    if ansible all -i inventory -m shell -a "docker --version" >/dev/null 2>&1; then
        test_results+=("Docker: ✅ PASS")
    else
        test_results+=("Docker: ❌ FAIL")
    fi
    
    # Test 7: Verify Docker Swarm is initialized
    print_info "Verifying Docker Swarm cluster..."
    if ansible swarm_managers -i inventory -m shell -a "docker node ls" >/dev/null 2>&1; then
        test_results+=("Swarm: ✅ PASS")
    else
        test_results+=("Swarm: ❌ FAIL")
    fi
    
    # Display test results
    print_header "Test Results Summary"
    for result in "${test_results[@]}"; do
        echo "  $result"
    done
    
    # Count failures
    local failures=$(printf '%s\n' "${test_results[@]}" | grep -c "❌ FAIL" || true)
    
    if [ "$failures" -eq 0 ]; then
        print_success "All tests passed! 🎉"
        return 0
    else
        print_error "$failures test(s) failed"
        return 1
    fi
}

# Function to show help
show_help() {
    echo "Reset and Test Script for Debian Vagrant Cluster"
    echo ""
    echo "Usage: $0 [COMMAND]"
    echo ""
    echo "Commands:"
    echo "  full-reset     Destroy everything and run full test cycle"
    echo "  destroy-only   Only destroy all machines"
    echo "  start-only     Only start all machines"
    echo "  test-only      Only run tests (assumes machines are running)"
    echo "  ssh-test       Only test SSH connectivity"
    echo "  ansible-test   Only test Ansible connectivity"
    echo "  setup-only     Only run setup playbook"
    echo "  deploy-only    Only run deployment playbook"
    echo "  help           Show this help message"
    echo ""
    echo "Examples:"
    echo "  $0 full-reset     # Complete destroy/recreate/test cycle"
    echo "  $0 test-only       # Run tests on existing machines"
    echo "  $0 ssh-test        # Quick SSH connectivity check"
    echo ""
    echo "This script will:"
    echo "  1. Check prerequisites (vagrant, ansible, make)"
    echo "  2. Destroy all VMs and clean up"
    echo "  3. Start all VMs fresh"
    echo "  4. Test SSH connectivity"
    echo "  5. Test Ansible connectivity"
    echo "  6. Run setup playbook (dependencies, swap)"
    echo "  7. Run deployment playbook (Docker, services)"
    echo "  8. Verify everything is working"
}

# Main script logic
main() {
    local command=${1:-help}
    
    case "$command" in
        full-reset)
            print_header "Full Reset and Test Cycle"
            check_prerequisites
            destroy_all
            start_all
            run_tests
        ;;
        destroy-only)
            print_header "Destroy Only"
            check_prerequisites
            destroy_all
        ;;
        start-only)
            print_header "Start Only"
            check_prerequisites
            start_all
        ;;
        test-only)
            print_header "Test Only"
            check_prerequisites
            run_tests
        ;;
        ssh-test)
            print_header "SSH Test Only"
            check_prerequisites
            test_ssh
        ;;
        ansible-test)
            print_header "Ansible Test Only"
            check_prerequisites
            test_ansible
        ;;
        setup-only)
            print_header "Setup Only"
            check_prerequisites
            run_setup
        ;;
        deploy-only)
            print_header "Deploy Only"
            check_prerequisites
            run_deployment
        ;;
        help|--help|-h)
            show_help
        ;;
        *)
            print_error "Unknown command: $command"
            show_help
            exit 1
        ;;
    esac
}

# Run main function with all arguments
main "$@"
    else
        print_error "Setup playbook failed"
        return 1
    fi
}

# Function to run deployment playbook
run_deployment() {
    print_header "Running Deployment Playbook"
    
    print_info "Deploying applications and services..."
    if ansible-playbook -i inventory deploy-playbook.yml; then
        print_success "Deployment playbook completed successfully"
        return 0
    else
        print_error "Deployment playbook failed"
        return 1
    fi
}

# Function to run comprehensive tests
run_tests() {
    print_header "Running Comprehensive Tests"
    
    local test_results=()
    
    # Test 1: SSH Connectivity
    if test_ssh; then
        test_results+=("SSH: ✅ PASS")
    else
        test_results+=("SSH: ❌ FAIL")
    fi
    
    # Test 2: Ansible Connectivity
    if test_ansible; then
        test_results+=("Ansible: ✅ PASS")
    else
        test_results+=("Ansible: ❌ FAIL")
    fi
    
    # Test 3: Setup Playbook
    if run_setup; then
        test_results+=("Setup: ✅ PASS")
    else
        test_results+=("Setup: ❌ FAIL")
    fi
    
    # Test 4: Deployment Playbook
    if run_deployment; then
        test_results+=("Deployment: ✅ PASS")
    else
        test_results+=("Deployment: ❌ FAIL")
    fi
    
    # Test 5: Verify swap is active
    print_info "Verifying swap is active..."
    if ansible all -i inventory -m shell -a "cat /proc/swaps" | grep -q "swapfile"; then
        test_results+=("Swap: ✅ PASS")
    else
        test_results+=("Swap: ❌ FAIL")
    fi
    
    # Test 6: Verify Docker is running
    print_info "Verifying Docker is running..."
    if ansible all -i inventory -m shell -a "docker --version" >/dev/null 2>&1; then
        test_results+=("Docker: ✅ PASS")
    else
        test_results+=("Docker: ❌ FAIL")
    fi
    
    # Test 7: Verify Docker Swarm is initialized
    print_info "Verifying Docker Swarm cluster..."
    if ansible swarm_managers -i inventory -m shell -a "docker node ls" >/dev/null 2>&1; then
        test_results+=("Swarm: ✅ PASS")
    else
        test_results+=("Swarm: ❌ FAIL")
    fi
    
    # Display test results
    print_header "Test Results Summary"
    for result in "${test_results[@]}"; do
        echo "  $result"
    done
    
    # Count failures
    local failures=$(printf '%s\n' "${test_results[@]}" | grep -c "❌ FAIL" || true)
    
    if [ "$failures" -eq 0 ]; then
        print_success "All tests passed! 🎉"
        return 0
    else
        print_error "$failures test(s) failed"
        return 1
    fi
}

# Function to show help
show_help() {
    echo "Reset and Test Script for Debian Vagrant Cluster"
    echo ""
    echo "Usage: $0 [COMMAND]"
    echo ""
    echo "Commands:"
    echo "  full-reset     Destroy everything and run full test cycle"
    echo "  destroy-only   Only destroy all machines"
    echo "  start-only     Only start all machines"
    echo "  test-only      Only run tests (assumes machines are running)"
    echo "  ssh-test       Only test SSH connectivity"
    echo "  ansible-test   Only test Ansible connectivity"
    echo "  setup-only     Only run setup playbook"
    echo "  deploy-only    Only run deployment playbook"
    echo "  help           Show this help message"
    echo ""
    echo "Examples:"
    echo "  $0 full-reset     # Complete destroy/recreate/test cycle"
    echo "  $0 test-only       # Run tests on existing machines"
    echo "  $0 ssh-test        # Quick SSH connectivity check"
    echo ""
    echo "This script will:"
    echo "  1. Check prerequisites (vagrant, ansible, make)"
    echo "  2. Destroy all VMs and clean up"
    echo "  3. Start all VMs fresh"
    echo "  4. Test SSH connectivity"
    echo "  5. Test Ansible connectivity"
    echo "  6. Run setup playbook (dependencies, swap)"
    echo "  7. Run deployment playbook (Docker, services)"
    echo "  8. Verify everything is working"
}

# Main script logic
main() {
    local command=${1:-help}
    
    case "$command" in
        full-reset)
            print_header "Full Reset and Test Cycle"
            check_prerequisites
            destroy_all
            start_all
            run_tests
        ;;
        destroy-only)
            print_header "Destroy Only"
            check_prerequisites
            destroy_all
        ;;
        start-only)
            print_header "Start Only"
            check_prerequisites
            start_all
        ;;
        test-only)
            print_header "Test Only"
            check_prerequisites
            run_tests
        ;;
        ssh-test)
            print_header "SSH Test Only"
            check_prerequisites
            test_ssh
        ;;
        ansible-test)
            print_header "Ansible Test Only"
            check_prerequisites
            test_ansible
        ;;
        setup-only)
            print_header "Setup Only"
            check_prerequisites
            run_setup
        ;;
        deploy-only)
            print_header "Deploy Only"
            check_prerequisites
            run_deployment
        ;;
        help|--help|-h)
            show_help
        ;;
        *)
            print_error "Unknown command: $command"
            show_help
            exit 1
        ;;
    esac
}

# Run main function with all arguments
main "$@"