#!/bin/bash

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PARENT_DIR="$(dirname "$SCRIPT_DIR")"
RESULTS_DIR="/root/workspace/ppanda/mock-test-jsons"
SSH_KEY="/tmp/test_ansible_key"

mkdir -p "$RESULTS_DIR"

echo "Generating SSH key pair..."
if [ ! -f "$SSH_KEY" ]; then
    ssh-keygen -t rsa -b 4096 -f "$SSH_KEY" -N "" > /dev/null 2>&1
fi

echo "Removing any existing test containers..."
for name in ubuntu-20-04-test ubuntu-22-04-test ubuntu-24-04-test debian-11-test debian-12-test debian-13-test alpine-test; do
    docker rm -f "$name" > /dev/null 2>&1 || true
done

echo "Starting test containers..."

echo "Starting ubuntu-20-04-test..."
docker run -d --name ubuntu-20-04-test -p 2220:22 \
    -v "$SSH_KEY.pub:/root/.ssh/authorized_keys:ro" \
    -t ubuntu:20.04 /bin/bash > /dev/null 2>&1

echo "Starting ubuntu-22-04-test..."
docker run -d --name ubuntu-22-04-test -p 2221:22 \
    -v "$SSH_KEY.pub:/root/.ssh/authorized_keys:ro" \
    -t ubuntu:22.04 /bin/bash > /dev/null 2>&1

echo "Starting ubuntu-24-04-test..."
docker run -d --name ubuntu-24-04-test -p 2222:22 \
    -v "$SSH_KEY.pub:/root/.ssh/authorized_keys:ro" \
    -t ubuntu:24.04 /bin/bash > /dev/null 2>&1

echo "Starting debian-11-test..."
docker run -d --name debian-11-test -p 2223:22 \
    -v "$SSH_KEY.pub:/root/.ssh/authorized_keys:ro" \
    -t debian:11 /bin/bash > /dev/null 2>&1

echo "Starting debian-12-test..."
docker run -d --name debian-12-test -p 2224:22 \
    -v "$SSH_KEY.pub:/root/.ssh/authorized_keys:ro" \
    -t debian:12 /bin/bash > /dev/null 2>&1

echo "Starting debian-13-test..."
docker run -d --name debian-13-test -p 2225:22 \
    -v "$SSH_KEY.pub:/root/.ssh/authorized_keys:ro" \
    -t debian:13 /bin/bash > /dev/null 2>&1

echo "Starting alpine-test..."
docker run -d --name alpine-test -p 2226:22 \
    -v "$SSH_KEY.pub:/root/.ssh/authorized_keys:ro" \
    -t alpine:latest /bin/sh > /dev/null 2>&1

echo "Waiting for containers to initialize..."
sleep 10

echo "Setting up Ubuntu/Debian containers..."
for container in ubuntu-20-04-test ubuntu-22-04-test ubuntu-24-04-test debian-11-test debian-12-test debian-13-test; do
    echo "  Setting up $container..."
    docker exec "$container" /bin/bash -c "apt-get update -qq && apt-get install -y -qq openssh-server python3 sudo" > /dev/null 2>&1
    docker exec "$container" /bin/bash -c "mkdir -p /var/run/sshd" > /dev/null 2>&1
    docker exec "$container" /bin/bash -c "sed -i 's/#PermitRootLogin.*/PermitRootLogin yes/' /etc/ssh/sshd_config" > /dev/null 2>&1
    docker exec "$container" /bin/bash -c "sed -i 's/#PasswordAuthentication.*/PasswordAuthentication no/' /etc/ssh/sshd_config" > /dev/null 2>&1
    docker exec "$container" /bin/bash -c "echo 'root:password' | chpasswd" > /dev/null 2>&1
    docker exec "$container" service ssh start > /dev/null 2>&1 || docker exec "$container" /usr/sbin/sshd > /dev/null 2>&1
    sleep 2
done

echo "Setting up Alpine container..."
docker exec alpine-test /bin/sh -c "apk add --no-cache openssh openssh-server python3 sudo" > /dev/null 2>&1
docker exec alpine-test /bin/sh -c "mkdir -p /var/run/sshd" > /dev/null 2>&1
docker exec alpine-test /bin/sh -c "ssh-keygen -A" > /dev/null 2>&1
docker exec alpine-test /bin/sh -c "sed -i 's/#PermitRootLogin.*/PermitRootLogin yes/' /etc/ssh/sshd_config" > /dev/null 2>&1
docker exec alpine-test /bin/sh -c "sed -i 's/#PasswordAuthentication.*/PasswordAuthentication no/' /etc/ssh/sshd_config" > /dev/null 2>&1
docker exec alpine-test /bin/sh -c "echo 'root:password' | chpasswd" > /dev/null 2>&1
docker exec alpine-test /usr/sbin/sshd > /dev/null 2>&1
sleep 2

echo "Waiting for SSH to be ready..."
for port in 2220 2221 2222 2223 2224 2225 2226; do
    echo "  Waiting for port $port..."
    timeout 30 bash -c "until nc -z localhost $port 2>/dev/null; do sleep 1; done" || echo "  Warning: Port $port not ready"
done

echo "Testing SSH connections..."
SSH_READY=true
SSH_FAILED=()

for port in 2220 2221 2222 2223 2224 2225 2226; do
    if ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o ConnectTimeout=5 -i "$SSH_KEY" -p "$port" root@localhost "echo 'SSH OK'" > /dev/null 2>&1; then
        echo "  ✓ Port $port: SSH connection successful"
    else
        echo "  ✗ Port $port: SSH connection failed"
        SSH_READY=false
        SSH_FAILED+=($port)
    fi
done

if [ "$SSH_READY" = false ]; then
    echo ""
    echo "ERROR: SSH connections failed for ports: ${SSH_FAILED[*]}"
    echo "Checking container logs for failed ports..."
    for port in "${SSH_FAILED[@]}"; do
        case $port in
            2220) CONTAINER="ubuntu-20-04-test" ;;
            2221) CONTAINER="ubuntu-22-04-test" ;;
            2222) CONTAINER="ubuntu-24-04-test" ;;
            2223) CONTAINER="debian-11-test" ;;
            2224) CONTAINER="debian-12-test" ;;
            2225) CONTAINER="debian-13-test" ;;
            2226) CONTAINER="alpine-test" ;;
        esac
        echo ""
        echo "Logs for $CONTAINER:"
        docker logs "$CONTAINER" 2>&1 | tail -20
    done
    echo "Aborting tests due to SSH connection failures."
    exit 1
fi

echo ""
echo "=========================================="
echo "All SSH connections successful!"
echo "=========================================="
echo ""

cat > "$SCRIPT_DIR/test_inventory.ini" << 'EOF'
[test_containers]
ubuntu-20-04-test ansible_host=127.0.0.1 ansible_port=2220 ansible_user=root ansible_ssh_private_key_file=/tmp/test_ansible_key ansible_ssh_common_args='-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'
ubuntu-22-04-test ansible_host=127.0.0.1 ansible_port=2221 ansible_user=root ansible_ssh_private_key_file=/tmp/test_ansible_key ansible_ssh_common_args='-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'
ubuntu-24-04-test ansible_host=127.0.0.1 ansible_port=2222 ansible_user=root ansible_ssh_private_key_file=/tmp/test_ansible_key ansible_ssh_common_args='-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'
debian-11-test ansible_host=127.0.0.1 ansible_port=2223 ansible_user=root ansible_ssh_private_key_file=/tmp/test_ansible_key ansible_ssh_common_args='-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'
debian-12-test ansible_host=127.0.0.1 ansible_port=2224 ansible_user=root ansible_ssh_private_key_file=/tmp/test_ansible_key ansible_ssh_common_args='-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'
debian-13-test ansible_host=127.0.0.1 ansible_port=2225 ansible_user=root ansible_ssh_private_key_file=/tmp/test_ansible_key ansible_ssh_common_args='-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'
alpine-test ansible_host=127.0.0.1 ansible_port=2226 ansible_user=root ansible_ssh_private_key_file=/tmp/test_ansible_key ansible_ssh_common_args='-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'

[all:vars]
ansible_python_interpreter=/usr/bin/python3
EOF

echo "Running maintenance script tests..."
echo "=========================================="
echo ""

cd "$PARENT_DIR"

PLAYBOOK_SUCCESS=0
PLAYBOOK_FAILED=0

echo "Running scan_cves.yml..."
if ansible-playbook playbooks/scan_cves.yml -i "$SCRIPT_DIR/test_inventory.ini" 2>&1 | tee "$RESULTS_DIR/scan_cves_run.log"; then
    echo "  ✓ scan_cves.yml completed successfully"
    ((PLAYBOOK_SUCCESS++))
else
    echo "  ✗ scan_cves.yml failed"
    ((PLAYBOOK_FAILED++))
fi
echo ""

echo "Running analyze_disk_space.yml..."
if ansible-playbook playbooks/analyze_disk_space.yml -i "$SCRIPT_DIR/test_inventory.ini" 2>&1 | tee "$RESULTS_DIR/analyze_disk_space_run.log"; then
    echo "  ✓ analyze_disk_space.yml completed successfully"
    ((PLAYBOOK_SUCCESS++))
else
    echo "  ✗ analyze_disk_space.yml failed"
    ((PLAYBOOK_FAILED++))
fi
echo ""

echo "Running cleanup_docker.yml..."
if ansible-playbook playbooks/cleanup_docker.yml -i "$SCRIPT_DIR/test_inventory.ini" 2>&1 | tee "$RESULTS_DIR/cleanup_docker_run.log"; then
    echo "  ✓ cleanup_docker.yml completed successfully"
    ((PLAYBOOK_SUCCESS++))
else
    echo "  ✗ cleanup_docker.yml failed"
    ((PLAYBOOK_FAILED++))
fi
echo ""

echo "Running archive_logs.yml..."
if ansible-playbook playbooks/archive_logs.yml -i "$SCRIPT_DIR/test_inventory.ini" 2>&1 | tee "$RESULTS_DIR/archive_logs_run.log"; then
    echo "  ✓ archive_logs.yml completed successfully"
    ((PLAYBOOK_SUCCESS++))
else
    echo "  ✗ archive_logs.yml failed"
    ((PLAYBOOK_FAILED++))
fi
echo ""

echo "=========================================="
echo "Collecting JSON reports..."
echo "=========================================="

TOTAL_JSON_FILES=0

for container in ubuntu-20-04-test ubuntu-22-04-test ubuntu-24-04-test debian-11-test debian-12-test debian-13-test alpine-test; do
    echo "Fetching reports from $container..."
    
    JSON_COUNT=$(docker exec "$container" /bin/sh -c "find /tmp -name '*_report_*.json' -type f 2>/dev/null" | wc -l)
    
    if [ "$JSON_COUNT" -gt 0 ]; then
        docker exec "$container" /bin/sh -c "find /tmp -name '*_report_*.json' -type f" 2>/dev/null | while read -r file; do
            filename=$(basename "$file")
            echo "  Found: $filename"
            if docker cp "$container:$file" "$RESULTS_DIR/${container}_$filename" 2>/dev/null; then
                ((TOTAL_JSON_FILES++))
            fi
        done
    else
        echo "  No JSON reports found"
    fi
done

echo ""
echo "=========================================="
echo "Test Summary"
echo "=========================================="
echo "Results directory: $RESULTS_DIR"
echo ""
echo "Playbook Results:"
echo "  Successful: $PLAYBOOK_SUCCESS"
echo "  Failed: $PLAYBOOK_FAILED"
echo ""
echo "JSON Reports Collected: $TOTAL_JSON_FILES"
echo ""

if [ -d "$RESULTS_DIR" ]; then
    echo "Collected JSON files:"
    find "$RESULTS_DIR" -name "*.json" -type f -exec basename {} \; 2>/dev/null | sort -u
    echo ""
    
    echo "Log files:"
    find "$RESULTS_DIR" -name "*_run.log" -type f -exec basename {} \; 2>/dev/null
    echo ""
fi

echo "Cleaning up test containers..."
for container in ubuntu-20-04-test ubuntu-22-04-test ubuntu-24-04-test debian-11-test debian-12-test debian-13-test alpine-test; do
    docker stop "$container" > /dev/null 2>&1 || true
    docker rm "$container" > /dev/null 2>&1 || true
done

echo ""
echo "=========================================="
echo "Testing complete!"
echo "=========================================="
echo "All JSON reports and logs are available in: $RESULTS_DIR"

if [ "$PLAYBOOK_FAILED" -gt 0 ]; then
    echo ""
    echo "WARNING: $PLAYBOOK_FAILED playbook(s) failed. Please check the log files."
    exit 1
else
    echo ""
    echo "SUCCESS: All playbooks completed successfully!"
    exit 0
fi