#!/bin/bash

###############################################################################
# Layer 2: General Infrastructure Tools
# 
# This script installs rsync and docker on the machines that need them.
# Must be run after Layer 1A (VPS) or Layer 1B (Nodito) is complete.
###############################################################################

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
    
    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
    
    if ! command -v ansible &> /dev/null; then
        print_error "Ansible not found"
        ((errors++))
    else
        print_success "Ansible found"
    fi
    
    if [ ! -f "$ANSIBLE_DIR/inventory.ini" ]; then
        print_error "inventory.ini not found"
        ((errors++))
    else
        print_success "inventory.ini exists"
    fi
    
    if [ $errors -gt 0 ]; then
        print_error "Layer 0 is not complete"
        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_ssh_connectivity() {
    print_header "Testing SSH Connectivity"
    
    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}"
    
    local all_good=true
    
    for group in vipy watchtower spacey nodito; do
        local hosts=$(get_hosts_from_inventory "$group")
        if [ -n "$hosts" ]; then
            for host in $hosts; do
                print_info "Testing SSH to $host as counterweight..."
                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: OK"
                else
                    print_error "Cannot SSH to $host as counterweight"
                    print_warning "Make sure Layer 1A or 1B is complete for this host"
                    all_good=false
                fi
            done
        fi
    done
    
    if [ "$all_good" = false ]; then
        echo ""
        print_error "SSH connectivity test failed"
        print_info "Ensure Layer 1A (VPS) or Layer 1B (Nodito) is complete"
        echo ""
        if ! confirm_action "Continue anyway?"; then
            exit 1
        fi
    fi
    
    echo ""
    print_success "SSH connectivity verified"
}

###############################################################################
# rsync Installation
###############################################################################

install_rsync() {
    print_header "Installing rsync"
    
    cd "$ANSIBLE_DIR"
    
    print_info "rsync is needed for backup operations"
    print_info "Recommended hosts: vipy, watchtower, lapy"
    echo ""
    
    # Show available hosts
    echo "Available hosts in inventory:"
    for group in vipy watchtower spacey nodito lapy; do
        local hosts=$(get_hosts_from_inventory "$group")
        if [ -n "$hosts" ]; then
            echo "  [$group]: $hosts"
        fi
    done
    echo ""
    
    print_info "Installation options:"
    echo "  1. Install on recommended hosts (vipy, watchtower, lapy)"
    echo "  2. Install on all hosts"
    echo "  3. Custom selection (specify groups)"
    echo "  4. Skip rsync installation"
    echo ""
    
    echo -e -n "${BLUE}Choose option${NC} [1-4]: "
    read option
    
    local limit_hosts=""
    case "$option" in
        1)
            limit_hosts="vipy,watchtower,lapy"
            print_info "Installing rsync on: vipy, watchtower, lapy"
            ;;
        2)
            limit_hosts="all"
            print_info "Installing rsync on: all hosts"
            ;;
        3)
            echo -e -n "${BLUE}Enter groups (comma-separated, e.g., vipy,watchtower,nodito)${NC}: "
            read limit_hosts
            print_info "Installing rsync on: $limit_hosts"
            ;;
        4)
            print_warning "Skipping rsync installation"
            return 1
            ;;
        *)
            print_error "Invalid option"
            return 1
            ;;
    esac
    
    echo ""
    if ! confirm_action "Proceed with rsync installation?"; then
        print_warning "Skipped rsync installation"
        return 1
    fi
    
    print_info "Running: ansible-playbook -i inventory.ini infra/900_install_rsync.yml --limit $limit_hosts"
    echo ""
    
    if ansible-playbook -i inventory.ini infra/900_install_rsync.yml --limit "$limit_hosts"; then
        print_success "rsync installation complete"
        return 0
    else
        print_error "rsync installation failed"
        return 1
    fi
}

###############################################################################
# Docker Installation
###############################################################################

install_docker() {
    print_header "Installing Docker and Docker Compose"
    
    cd "$ANSIBLE_DIR"
    
    print_info "Docker is needed for containerized services"
    print_info "Recommended hosts: vipy, watchtower"
    echo ""
    
    # Show available hosts (exclude lapy - docker on laptop is optional)
    echo "Available hosts in inventory:"
    for group in vipy watchtower spacey nodito; do
        local hosts=$(get_hosts_from_inventory "$group")
        if [ -n "$hosts" ]; then
            echo "  [$group]: $hosts"
        fi
    done
    echo ""
    
    print_info "Installation options:"
    echo "  1. Install on recommended hosts (vipy, watchtower)"
    echo "  2. Install on all hosts"
    echo "  3. Custom selection (specify groups)"
    echo "  4. Skip docker installation"
    echo ""
    
    echo -e -n "${BLUE}Choose option${NC} [1-4]: "
    read option
    
    local limit_hosts=""
    case "$option" in
        1)
            limit_hosts="vipy,watchtower"
            print_info "Installing Docker on: vipy, watchtower"
            ;;
        2)
            limit_hosts="all"
            print_info "Installing Docker on: all hosts"
            ;;
        3)
            echo -e -n "${BLUE}Enter groups (comma-separated, e.g., vipy,watchtower,nodito)${NC}: "
            read limit_hosts
            print_info "Installing Docker on: $limit_hosts"
            ;;
        4)
            print_warning "Skipping Docker installation"
            return 1
            ;;
        *)
            print_error "Invalid option"
            return 1
            ;;
    esac
    
    echo ""
    if ! confirm_action "Proceed with Docker installation?"; then
        print_warning "Skipped Docker installation"
        return 1
    fi
    
    print_info "Running: ansible-playbook -i inventory.ini infra/910_docker_playbook.yml --limit $limit_hosts"
    echo ""
    
    if ansible-playbook -i inventory.ini infra/910_docker_playbook.yml --limit "$limit_hosts"; then
        print_success "Docker installation complete"
        print_warning "You may need to log out and back in for docker group to take effect"
        return 0
    else
        print_error "Docker installation failed"
        return 1
    fi
}

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

verify_installations() {
    print_header "Verifying Installations"
    
    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}"
    
    echo "Checking installed tools on hosts..."
    echo ""
    
    # Check all remote hosts
    for group in vipy watchtower spacey nodito; do
        local hosts=$(get_hosts_from_inventory "$group")
        if [ -n "$hosts" ]; then
            for host in $hosts; do
                print_info "Checking $host..."
                
                # Check rsync
                if timeout 5 ssh -i "$ssh_key" -o StrictHostKeyChecking=no -o BatchMode=yes counterweight@$host "command -v rsync" &>/dev/null; then
                    print_success "$host: rsync installed"
                else
                    print_warning "$host: rsync not found (may not be needed)"
                fi
                
                # Check docker
                if timeout 5 ssh -i "$ssh_key" -o StrictHostKeyChecking=no -o BatchMode=yes counterweight@$host "command -v docker" &>/dev/null; then
                    print_success "$host: docker installed"
                    
                    # Check docker service
                    if timeout 5 ssh -i "$ssh_key" -o StrictHostKeyChecking=no -o BatchMode=yes counterweight@$host "sudo systemctl is-active docker" &>/dev/null; then
                        print_success "$host: docker service running"
                    else
                        print_warning "$host: docker service not running"
                    fi
                else
                    print_warning "$host: docker not found (may not be needed)"
                fi
                
                echo ""
            done
        fi
    done
}

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

print_summary() {
    print_header "Layer 2 Setup Complete! 🎉"
    
    echo "Summary:"
    echo ""
    print_success "Infrastructure tools installed on specified hosts"
    echo ""
    
    print_info "What was installed:"
    echo "  • rsync - for backup operations"
    echo "  • docker + docker compose - for containerized services"
    echo ""
    
    print_info "Next steps:"
    echo "  1. Proceed to Layer 3: ./scripts/setup_layer_3_caddy.sh"
    echo ""
}

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

main() {
    clear
    
    print_header "🔧 Layer 2: General Infrastructure Tools"
    
    echo "This script will install rsync and docker on your infrastructure."
    echo ""
    
    if ! confirm_action "Continue with Layer 2 setup?"; then
        echo "Setup cancelled."
        exit 0
    fi
    
    check_layer_0_complete
    check_ssh_connectivity
    
    # Install tools
    install_rsync
    echo ""
    install_docker
    
    verify_installations
    print_summary
}

# Run main function
main "$@"