#!/bin/bash

# Enterprise Infrastructure Scaling Simulation Script
# Simulates scaling operations for infrastructure nodes

set -euo pipefail

# Configuration
DOCKER_COMPOSE_FILE="docker-compose.yml"
INVENTORY_FILE="inventory/hosts.ini"
LOG_FILE="logs/scaling_simulation.log"

# Default values
DIRECTION="${1:-up}"
COUNT="${2:-1}"
NODE_TYPE="${3:-web}"
SIMULATION_MODE="${SIMULATION_MODE:-false}"

# Logging function
log() {
    echo "$(date '+%Y-%m-%d %H:%M:%S') - $*" | tee -a "$LOG_FILE"
}

# Error handling
error_exit() {
    log "ERROR: $1"
    exit 1
}

# Validate inputs
validate_inputs() {
    if [[ "$DIRECTION" != "up" && "$DIRECTION" != "down" ]]; then
        error_exit "Invalid direction: $DIRECTION. Must be 'up' or 'down'"
    fi

    if ! [[ "$COUNT" =~ ^[0-9]+$ ]] || [ "$COUNT" -lt 1 ]; then
        error_exit "Invalid count: $COUNT. Must be a positive integer"
    fi

    case "$NODE_TYPE" in
        web|db|lb|monitor) ;;
        *) error_exit "Invalid node type: $NODE_TYPE. Must be web, db, lb, or monitor" ;;
    esac
}

# Get current node count
get_current_count() {
    local type="$1"
    case "$type" in
        web) docker-compose ps web | grep -c "Up" ;;
        db) docker-compose ps db | grep -c "Up" ;;
        lb) docker-compose ps lb | grep -c "Up" ;;
        monitor) docker-compose ps monitor | grep -c "Up" ;;
    esac
}

# Scale up infrastructure
scale_up() {
    local type="$1"
    local count="$2"

    log "Scaling up $count $type nodes"

    # Update docker-compose replica count
    sed -i.bak "s/replicas: [0-9]\+/replicas: $(( $(get_current_count "$type") + count ))/" "$DOCKER_COMPOSE_FILE"

    # Deploy new containers
    docker-compose up -d --scale "${type}=${count}"

    # Wait for containers to be ready
    log "Waiting for containers to be ready..."
    sleep 30

    # Update inventory
    update_inventory "$type" "$count" "add"

    # Run provisioning playbook on new nodes
    if [ "$SIMULATION_MODE" = false ]; then
        ansible-playbook -i "$INVENTORY_FILE" playbooks/provision.yml --limit "${type}*"
    fi

    log "Successfully scaled up $count $type nodes"
}

# Scale down infrastructure
scale_down() {
    local type="$1"
    local count="$2"

    local current_count=$(get_current_count "$type")
    if [ "$current_count" -lt "$count" ]; then
        error_exit "Cannot scale down $count nodes. Only $current_count $type nodes currently running"
    fi

    log "Scaling down $count $type nodes"

    # Select nodes to remove (oldest first)
    local nodes_to_remove=$(docker-compose ps "$type" | grep "Up" | head -n "$count" | awk '{print $1}')

    # Decommission nodes
    for node in $nodes_to_remove; do
        if [ "$SIMULATION_MODE" = false ]; then
            ansible-playbook -i "$INVENTORY_FILE" playbooks/decommission.yml --limit "$node"
        fi
        docker stop "$node"
        docker rm "$node"
    done

    # Update docker-compose replica count
    sed -i.bak "s/replicas: [0-9]\+/replicas: $(( current_count - count ))/" "$DOCKER_COMPOSE_FILE"

    # Update inventory
    update_inventory "$type" "$count" "remove"

    log "Successfully scaled down $count $type nodes"
}

# Update Ansible inventory
update_inventory() {
    local type="$1"
    local count="$2"
    local action="$3"

    log "Updating inventory for $action $count $type nodes"

    # This would be more complex in a real implementation
    # For simulation, we'll just log the action
    case "$action" in
        add)
            log "Added $count $type nodes to inventory"
            ;;
        remove)
            log "Removed $count $type nodes from inventory"
            ;;
    esac
}

# Health check after scaling
health_check() {
    log "Running health checks after scaling"

    # Check container status
    if ! docker-compose ps | grep -q "Up"; then
        error_exit "Some containers failed to start"
    fi

    # Ansible ping check
    if [ "$SIMULATION_MODE" = false ]; then
        if ! ansible -i "$INVENTORY_FILE" all -m ping >/dev/null 2>&1; then
            log "WARNING: Some nodes failed Ansible ping check"
        fi
    fi

    log "Health checks completed"
}

# Generate scaling report
generate_report() {
    local report_file="reports/scaling_report_$(date +%Y%m%d_%H%M%S).txt"

    cat > "$report_file" << EOF
Scaling Simulation Report
========================

Timestamp: $(date)
Direction: $DIRECTION
Node Type: $NODE_TYPE
Count: $COUNT
Simulation Mode: $SIMULATION_MODE

Current Status:
$(docker-compose ps)

Inventory Status:
$(ansible -i "$INVENTORY_FILE" --list-hosts all 2>/dev/null || echo "Ansible inventory check failed")

Log File: $LOG_FILE
EOF

    log "Scaling report generated: $report_file"
}

# Main execution
main() {
    log "Starting scaling simulation: $DIRECTION $COUNT $NODE_TYPE nodes"

    validate_inputs

    case "$DIRECTION" in
        up)
            scale_up "$NODE_TYPE" "$COUNT"
            ;;
        down)
            scale_down "$NODE_TYPE" "$COUNT"
            ;;
    esac

    health_check
    generate_report

    log "Scaling simulation completed successfully"
}

# Initialize logging
mkdir -p logs reports

# Run main function
main "$@"