#!/usr/bin/env bash

#####################################################################################
#### docker-smoke-tests.sh — Verify Docker daemon and containers are healthy     ####
#### Checks daemon, API, lifecycle, networking, volumes, DNS, compose, images.    ####
####                                                                             ####
#### Author:  Phil Connor                                                        ####
#### Contact: contact@mylinux.work                                               ####
#### License: MIT                                                                ####
#### Version: 1.0                                                                ####
####                                                                             ####
#### Usage:   ./docker-smoke-tests.sh                                            ####
####          COMPOSE_FILE=/path/to/docker-compose.yml ./docker-smoke-tests.sh   ####
####                                                                             ####
#### See --help for all options.                                                 ####
#####################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
TEST_IMAGE="${TEST_IMAGE:-alpine:latest}"
COMPOSE_FILE="${COMPOSE_FILE:-}"
SKIP_LIFECYCLE="${SKIP_LIFECYCLE:-false}"
SKIP_NETWORK="${SKIP_NETWORK:-false}"
SKIP_VOLUME="${SKIP_VOLUME:-false}"
SKIP_BUILD="${SKIP_BUILD:-false}"
DNS_TEST_DOMAIN="${DNS_TEST_DOMAIN:-google.com}"
OUTPUT_FORMAT="${OUTPUT_FORMAT:-text}"
COLOR="${COLOR:-auto}"
VERBOSE="${VERBOSE:-false}"

# ── State ─────────────────────────────────────────────────────────────
PASS=0
FAIL=0
SKIP=0
TOTAL=0
RESULTS=()
START_TIME=""

# ── Test artifact names ──────────────────────────────────────────────
SMOKE_PREFIX="smoke-test-$$"
SMOKE_CONTAINER="${SMOKE_PREFIX}-ctr"
SMOKE_PORT_CONTAINER="${SMOKE_PREFIX}-port"
SMOKE_DNS_CONTAINER="${SMOKE_PREFIX}-dns"
SMOKE_VOL_CONTAINER="${SMOKE_PREFIX}-vol"
SMOKE_NET_CONTAINER="${SMOKE_PREFIX}-net"
SMOKE_MEM_CONTAINER="${SMOKE_PREFIX}-mem"
SMOKE_VOLUME="${SMOKE_PREFIX}-vol"
SMOKE_NETWORK="${SMOKE_PREFIX}-net"
SMOKE_BUILD_TAG="${SMOKE_PREFIX}-img"
SMOKE_PORT=$((47000 + (RANDOM % 1000)))

# ── Colors ────────────────────────────────────────────────────────────
setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        RED="" GREEN="" YELLOW="" BLUE="" BOLD="" RESET=""
        return
    fi
    if [[ "$COLOR" == "always" ]] || [[ -t 1 ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[0;33m'
        BLUE='\033[0;34m'
        BOLD='\033[1m'
        RESET='\033[0m'
    else
        RED="" GREEN="" YELLOW="" BLUE="" BOLD="" RESET=""
    fi
}

# ── Logging ───────────────────────────────────────────────────────────
log()     { echo -e "${BLUE}[INFO]${RESET} $*"; }
warn()    { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }
err()     { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
verbose() { if [[ "$VERBOSE" == "true" ]]; then echo -e "${BLUE}[DEBUG]${RESET} $*"; fi; }

# ── Test Result Recording ─────────────────────────────────────────────
record_pass() {
    local name="$1" detail="${2:-}"
    ((PASS++)) || true; ((TOTAL++)) || true
    RESULTS+=("PASS|${name}|${detail}")
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then echo "ok ${TOTAL} - ${name}${detail:+ (${detail})}"
    else echo -e "  ${GREEN}✓${RESET} ${name}${detail:+ — ${detail}}"; fi
}

record_fail() {
    local name="$1" detail="${2:-}"
    ((FAIL++)) || true; ((TOTAL++)) || true
    RESULTS+=("FAIL|${name}|${detail}")
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        echo "not ok ${TOTAL} - ${name}"
        [[ -n "$detail" ]] && echo "  # ${detail}"
    else echo -e "  ${RED}✗${RESET} ${name}${detail:+ — ${detail}}"; fi
}

record_skip() {
    local name="$1" reason="${2:-}"
    ((SKIP++)) || true; ((TOTAL++)) || true
    RESULTS+=("SKIP|${name}|${reason}")
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then echo "ok ${TOTAL} - ${name} # SKIP ${reason}"
    else echo -e "  ${YELLOW}⊘${RESET} ${name}${reason:+ — ${reason}}"; fi
}

# ── Helpers ───────────────────────────────────────────────────────────
has_cmd() { command -v "$1" >/dev/null 2>&1; }

remove_container() {
    local name="$1"
    docker rm -f "$name" >/dev/null 2>&1 || true
}

section() {
    if [[ "$OUTPUT_FORMAT" != "tap" ]]; then echo ""; echo -e "${BOLD}$1${RESET}"; fi
}

# ── Cleanup ───────────────────────────────────────────────────────────
# shellcheck disable=SC2317
cleanup() {
    verbose "Cleaning up test artifacts..."
    remove_container "$SMOKE_CONTAINER"
    remove_container "$SMOKE_PORT_CONTAINER"
    remove_container "$SMOKE_DNS_CONTAINER"
    remove_container "$SMOKE_VOL_CONTAINER"
    remove_container "$SMOKE_NET_CONTAINER"
    remove_container "$SMOKE_MEM_CONTAINER"
    docker volume rm -f "$SMOKE_VOLUME" >/dev/null 2>&1 || true
    docker network rm "$SMOKE_NETWORK" >/dev/null 2>&1 || true
    docker rmi -f "$SMOKE_BUILD_TAG" >/dev/null 2>&1 || true
}
trap cleanup EXIT

# ══════════════════════════════════════════════════════════════════════
#                          TEST FUNCTIONS
# ══════════════════════════════════════════════════════════════════════

# ── 1. Docker daemon running ─────────────────────────────────────────
test_daemon_running() {
    if has_cmd systemctl; then
        if systemctl is-active --quiet docker 2>/dev/null; then record_pass "Docker daemon running" "systemctl active"
        else record_fail "Docker daemon running" "systemctl inactive"; fi
    elif has_cmd service; then
        if service docker status >/dev/null 2>&1; then record_pass "Docker daemon running" "service running"
        else record_fail "Docker daemon running" "service stopped"; fi
    elif docker info >/dev/null 2>&1; then record_pass "Docker daemon running" "docker info ok"
    else record_fail "Docker daemon running" "cannot determine status"; fi
}

# ── 2. Docker API responsive ─────────────────────────────────────────
test_api_responsive() {
    local output
    if output=$(timeout 10 docker info 2>&1); then
        local ver; ver=$(echo "$output" | grep -i "Server Version" | head -1 | awk '{print $NF}') || true
        record_pass "Docker API responsive" "server ${ver:-unknown}"
    else record_fail "Docker API responsive" "docker info timed out or failed"; fi
}

# ── 3. Docker socket accessible ──────────────────────────────────────
test_socket_accessible() {
    local socket="/var/run/docker.sock"
    if [[ -S "$socket" ]]; then
        if [[ -r "$socket" && -w "$socket" ]]; then record_pass "Docker socket accessible" "$socket"
        else record_fail "Docker socket accessible" "$socket not readable/writable"; fi
    elif docker info >/dev/null 2>&1; then record_pass "Docker socket accessible" "non-default socket"
    else record_fail "Docker socket accessible" "$socket not found"; fi
}

# ── 4. Container lifecycle ───────────────────────────────────────────
test_container_lifecycle() {
    if [[ "$SKIP_LIFECYCLE" == "true" ]]; then record_skip "Container lifecycle" "SKIP_LIFECYCLE=true"; return; fi
    remove_container "$SMOKE_CONTAINER"

    if ! docker create --name "$SMOKE_CONTAINER" "$TEST_IMAGE" sleep 30 >/dev/null 2>&1; then
        record_fail "Container lifecycle" "docker create failed"
        return
    fi

    if ! docker start "$SMOKE_CONTAINER" >/dev/null 2>&1; then
        record_fail "Container lifecycle" "docker start failed"
        return
    fi

    local exec_output
    exec_output=$(docker exec "$SMOKE_CONTAINER" echo "smoke-ok" 2>&1) || true
    if [[ "$exec_output" != "smoke-ok" ]]; then
        record_fail "Container lifecycle" "docker exec failed"
        return
    fi

    if ! docker stop -t 5 "$SMOKE_CONTAINER" >/dev/null 2>&1; then
        record_fail "Container lifecycle" "docker stop failed"
        return
    fi

    if ! docker rm "$SMOKE_CONTAINER" >/dev/null 2>&1; then
        record_fail "Container lifecycle" "docker rm failed"
        return
    fi

    record_pass "Container lifecycle" "create/start/exec/stop/rm"
}

# ── 5. Port binding ──────────────────────────────────────────────────
test_port_binding() {
    if [[ "$SKIP_LIFECYCLE" == "true" ]]; then record_skip "Port binding" "SKIP_LIFECYCLE=true"; return; fi
    if ! has_cmd curl; then record_skip "Port binding" "curl not installed"; return; fi
    remove_container "$SMOKE_PORT_CONTAINER"

    if ! docker run -d --name "$SMOKE_PORT_CONTAINER" \
        -p "${SMOKE_PORT}:80" \
        "$TEST_IMAGE" sh -c 'mkdir -p /var/www && echo "smoke-ok" > /var/www/index.html && httpd -f -p 80 -h /var/www 2>/dev/null || { while true; do echo -e "HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\nsmoke-ok\n" | nc -l -p 80 2>/dev/null || break; done; }' >/dev/null 2>&1; then
        record_fail "Port binding" "failed to start container with port mapping"
        return
    fi

    sleep 2
    local response
    response=$(curl -sf --max-time 5 "http://localhost:${SMOKE_PORT}/" 2>/dev/null) || true
    remove_container "$SMOKE_PORT_CONTAINER"

    if [[ "$response" == *"smoke-ok"* ]]; then
        record_pass "Port binding" "curl localhost:${SMOKE_PORT}"
    else
        record_fail "Port binding" "no response on localhost:${SMOKE_PORT}"
    fi
}

# ── 6. Container DNS ─────────────────────────────────────────────────
test_container_dns() {
    if [[ "$SKIP_LIFECYCLE" == "true" ]]; then record_skip "Container DNS" "SKIP_LIFECYCLE=true"; return; fi
    remove_container "$SMOKE_DNS_CONTAINER"

    local dns_output
    dns_output=$(docker run --rm --name "$SMOKE_DNS_CONTAINER" "$TEST_IMAGE" \
        sh -c "nslookup ${DNS_TEST_DOMAIN} 2>/dev/null || getent hosts ${DNS_TEST_DOMAIN} 2>/dev/null || ping -c1 -W3 ${DNS_TEST_DOMAIN} 2>/dev/null" 2>&1) || true

    if [[ -n "$dns_output" ]] && ! echo "$dns_output" | grep -qi "can't resolve\|not found\|failure\|NXDOMAIN"; then
        record_pass "Container DNS" "${DNS_TEST_DOMAIN}"
    else
        record_fail "Container DNS" "failed to resolve ${DNS_TEST_DOMAIN}"
    fi
}

# ── 7. Volume mount ──────────────────────────────────────────────────
test_volume_mount() {
    if [[ "$SKIP_VOLUME" == "true" ]]; then record_skip "Volume mount" "SKIP_VOLUME=true"; return; fi
    docker volume rm -f "$SMOKE_VOLUME" >/dev/null 2>&1 || true
    remove_container "$SMOKE_VOL_CONTAINER"

    if ! docker volume create "$SMOKE_VOLUME" >/dev/null 2>&1; then
        record_fail "Volume mount" "docker volume create failed"
        return
    fi

    local write_result
    write_result=$(docker run --rm --name "$SMOKE_VOL_CONTAINER" \
        -v "${SMOKE_VOLUME}:/data" "$TEST_IMAGE" \
        sh -c 'echo "smoke-vol-ok" > /data/test.txt && cat /data/test.txt' 2>&1) || true

    docker volume rm -f "$SMOKE_VOLUME" >/dev/null 2>&1 || true

    if [[ "$write_result" == "smoke-vol-ok" ]]; then
        record_pass "Volume mount" "write/read verified"
    else
        record_fail "Volume mount" "write/read mismatch"
    fi
}

# ── 8. Network create/connect ────────────────────────────────────────
test_network_create() {
    if [[ "$SKIP_NETWORK" == "true" ]]; then record_skip "Network create/connect" "SKIP_NETWORK=true"; return; fi
    docker network rm "$SMOKE_NETWORK" >/dev/null 2>&1 || true
    remove_container "$SMOKE_NET_CONTAINER"

    if ! docker network create --driver bridge "$SMOKE_NETWORK" >/dev/null 2>&1; then
        record_fail "Network create/connect" "docker network create failed"
        return
    fi

    local net_output
    net_output=$(docker run --rm --name "$SMOKE_NET_CONTAINER" \
        --network "$SMOKE_NETWORK" "$TEST_IMAGE" \
        sh -c 'ip addr show 2>/dev/null || ifconfig 2>/dev/null' 2>&1) || true

    docker network rm "$SMOKE_NETWORK" >/dev/null 2>&1 || true

    if [[ -n "$net_output" ]]; then
        record_pass "Network create/connect" "bridge network"
    else
        record_fail "Network create/connect" "container failed to attach to network"
    fi
}

# ── 9. Image pull ────────────────────────────────────────────────────
test_image_pull() {
    if docker pull "$TEST_IMAGE" >/dev/null 2>&1; then
        record_pass "Image pull" "$TEST_IMAGE"
    else
        record_fail "Image pull" "failed to pull $TEST_IMAGE"
    fi
}

# ── 10. Image build ──────────────────────────────────────────────────
test_image_build() {
    if [[ "$SKIP_BUILD" == "true" ]]; then record_skip "Image build" "SKIP_BUILD=true"; return; fi
    docker rmi -f "$SMOKE_BUILD_TAG" >/dev/null 2>&1 || true

    if echo "FROM alpine:latest" | docker build -t "$SMOKE_BUILD_TAG" - >/dev/null 2>&1; then
        docker rmi -f "$SMOKE_BUILD_TAG" >/dev/null 2>&1 || true
        record_pass "Image build" "inline Dockerfile"
    else
        record_fail "Image build" "docker build failed"
    fi
}

# ── 11. Docker Compose ───────────────────────────────────────────────
test_compose_stack() {
    if [[ -z "$COMPOSE_FILE" ]]; then
        record_skip "Compose stack" "COMPOSE_FILE not set"
        return
    fi
    if [[ ! -f "$COMPOSE_FILE" ]]; then
        record_fail "Compose stack" "${COMPOSE_FILE} not found"
        return
    fi
    local compose_cmd=""
    if docker compose version >/dev/null 2>&1; then
        compose_cmd="docker compose"
    elif has_cmd docker-compose; then
        compose_cmd="docker-compose"
    else
        record_skip "Compose stack" "neither 'docker compose' nor 'docker-compose' available"
        return
    fi

    local ps_output expected_count running_count
    ps_output=$($compose_cmd -f "$COMPOSE_FILE" ps --format json 2>/dev/null) || true

    if [[ -z "$ps_output" ]]; then
        ps_output=$($compose_cmd -f "$COMPOSE_FILE" ps 2>/dev/null) || true
        if [[ -z "$ps_output" ]]; then
            record_fail "Compose stack" "could not read compose project status"
            return
        fi
        expected_count=$(echo "$ps_output" | tail -n +2 | wc -l)
        running_count=$(echo "$ps_output" | tail -n +2 | grep -ciE "up|running" || true)
    else
        expected_count=$(echo "$ps_output" | grep -c '"Service"' 2>/dev/null || echo "$ps_output" | wc -l)
        running_count=$(echo "$ps_output" | grep -ciE '"running"' 2>/dev/null || true)
    fi

    if [[ "$expected_count" -eq 0 ]]; then
        record_fail "Compose stack" "no services found"
    elif [[ "$running_count" -ge "$expected_count" ]]; then
        record_pass "Compose stack" "${running_count}/${expected_count} services running"
    else
        record_fail "Compose stack" "${running_count}/${expected_count} services running"
    fi
}

# ── 12. Resource limits ──────────────────────────────────────────────
test_resource_limits() {
    if [[ "$SKIP_LIFECYCLE" == "true" ]]; then record_skip "Resource limits" "SKIP_LIFECYCLE=true"; return; fi
    remove_container "$SMOKE_MEM_CONTAINER"

    local mem_limit
    mem_limit=$(docker run --rm --name "$SMOKE_MEM_CONTAINER" \
        --memory=64m "$TEST_IMAGE" \
        sh -c 'cat /sys/fs/cgroup/memory.max 2>/dev/null || cat /sys/fs/cgroup/memory/memory.limit_in_bytes 2>/dev/null' 2>&1) || true

    if [[ -z "$mem_limit" ]]; then
        record_skip "Resource limits" "cgroup memory info not available"
        return
    fi

    local limit_bytes=67108864  # 64 MiB
    if [[ "$mem_limit" =~ ^[0-9]+$ ]]; then
        if [[ "$mem_limit" -le $((limit_bytes + 1048576)) ]]; then
            local limit_mb=$((mem_limit / 1048576))
            record_pass "Resource limits" "memory cgroup enforced (${limit_mb}M)"
        else
            record_fail "Resource limits" "memory limit not enforced (got ${mem_limit})"
        fi
    else
        record_skip "Resource limits" "unexpected cgroup value: ${mem_limit}"
    fi
}

# ── 13. Disk space ───────────────────────────────────────────────────
test_disk_space() {
    local df_output
    df_output=$(docker system df 2>/dev/null) || true

    if [[ -z "$df_output" ]]; then
        record_fail "Disk space" "docker system df failed"
        return
    fi

    local docker_root used_pct
    docker_root=$(docker info --format '{{.DockerRootDir}}' 2>/dev/null) || docker_root="/var/lib/docker"
    used_pct=$(df "$docker_root" 2>/dev/null | tail -1 | awk '{print $5}' | tr -d '%') || used_pct=0

    if [[ "$used_pct" -gt 80 ]]; then
        record_fail "Disk space" "${used_pct}% used (threshold 80%)"
    else
        record_pass "Disk space" "${used_pct}% used"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                          OUTPUT
# ══════════════════════════════════════════════════════════════════════

print_tap_header() {
    echo "TAP version 13"
}

print_tap_footer() {
    echo "1..${TOTAL}"
    echo "# pass ${PASS}"
    echo "# fail ${FAIL}"
    echo "# skip ${SKIP}"
}

print_summary() {
    local end_time; end_time=$(date +%s)
    local duration=$(( end_time - START_TIME ))
    echo ""
    echo -e "${BOLD}────────────────────────────────────────${RESET}"
    echo -e "${BOLD}Summary${RESET}  Docker Smoke Tests"
    echo -e "  ${GREEN}${PASS} passed${RESET}  ${RED}${FAIL} failed${RESET}  ${YELLOW}${SKIP} skipped${RESET}  (${duration}s)"
    echo -e "${BOLD}────────────────────────────────────────${RESET}"
    if [[ $FAIL -eq 0 ]]; then echo -e "${GREEN}${BOLD}All tests passed.${RESET}"
    else echo -e "${RED}${BOLD}${FAIL} test(s) failed.${RESET}"; fi
}

# ══════════════════════════════════════════════════════════════════════
#                          MAIN
# ══════════════════════════════════════════════════════════════════════

usage() {
    cat <<EOF
Usage: $(basename "$0") [OPTIONS]

Smoke-test Docker daemon, containers, networking, volumes, images, and compose.

Environment variables (all optional):
  TEST_IMAGE          Image for tests          (default: alpine:latest)
  COMPOSE_FILE        docker-compose.yml path  (default: none)
  SKIP_LIFECYCLE      Skip lifecycle tests     (default: false)
  SKIP_NETWORK        Skip network tests       (default: false)
  SKIP_VOLUME         Skip volume tests        (default: false)
  SKIP_BUILD          Skip build test          (default: false)
  DNS_TEST_DOMAIN     DNS domain to resolve    (default: google.com)
  OUTPUT_FORMAT       text or tap              (default: text)
  COLOR               auto, always, never      (default: auto)
  VERBOSE             Show debug output        (default: false)

Options:
  --skip-lifecycle    Skip container lifecycle tests
  --skip-network      Skip network tests
  --skip-volume       Skip volume tests
  --skip-build        Skip image build test
  --format FORMAT     Output format: text (default), tap
  --verbose           Show debug output
  --no-color          Disable colored output
  --help              Show this help

Examples:
  ./$(basename "$0")
  COMPOSE_FILE=/opt/stacks/web/docker-compose.yml ./$(basename "$0")
  OUTPUT_FORMAT=tap ./$(basename "$0") --skip-build
EOF
}

main() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --skip-lifecycle)  SKIP_LIFECYCLE=true ;;
            --skip-network)    SKIP_NETWORK=true ;;
            --skip-volume)     SKIP_VOLUME=true ;;
            --skip-build)      SKIP_BUILD=true ;;
            --format)          OUTPUT_FORMAT="$2"; shift ;;
            --verbose)         VERBOSE=true ;;
            --no-color)        COLOR=never ;;
            --help|-h)         usage; exit 0 ;;
            *)                 err "Unknown option: $1"; usage; exit 1 ;;
        esac
        shift
    done

    setup_colors

    if ! has_cmd docker; then
        err "docker CLI not found in PATH"
        exit 1
    fi

    START_TIME=$(date +%s)

    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        print_tap_header
    else
        echo ""
        echo -e "${BOLD}Docker Smoke Tests${RESET}"
        echo -e "Image: ${TEST_IMAGE}  Port: ${SMOKE_PORT}"
        echo -e "Time:  $(date -u +%Y-%m-%dT%H:%M:%SZ)"
    fi

    section "Daemon & Infrastructure"
    test_daemon_running
    test_api_responsive
    test_socket_accessible

    section "Image Operations"
    test_image_pull
    test_image_build

    section "Container Operations"
    test_container_lifecycle
    test_port_binding
    test_container_dns

    section "Storage"
    test_volume_mount
    test_disk_space

    section "Networking"
    test_network_create

    section "Resource Limits"
    test_resource_limits

    section "Compose"
    test_compose_stack

    # ── Results ──
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        print_tap_footer
    else
        print_summary
    fi

    [[ $FAIL -eq 0 ]] && exit 0 || exit 1
}

main "$@"