#!/usr/bin/env bash

set -euo pipefail

SCRIPT_DIR=$(cd $(dirname ${BASH_SOURCE[0]});pwd)
source "${SCRIPT_DIR}/common.sh"

export PLAYGROUND_TAG=${PLAYGROUND_TAG:-meta-failover}
export PLAYGROUND_PORT_OFFSET=${PLAYGROUND_PORT_OFFSET:-40000}
export TICI_META_COUNT=${TICI_META_COUNT:-2}
export TICI_WORKER_COUNT=${TICI_WORKER_COUNT:-1}
export TICI_READER_COUNT=${TICI_READER_COUNT:-1}
export PD_ADDR=${PD_ADDR:-127.0.0.1:$((DEFAULT_PD_PORT + PLAYGROUND_PORT_OFFSET))}
export LEADER_SWITCH_TIMEOUT_SECS=${LEADER_SWITCH_TIMEOUT_SECS:-30}
export REGISTRY_CLEANUP_TIMEOUT_SECS=${REGISTRY_CLEANUP_TIMEOUT_SECS:-90}
export PLAYGROUND_COMPONENT_VERSION=${PLAYGROUND_COMPONENT_VERSION:-v1.16.2-feature.fts}

PROBE_BIN=${PROBE_BIN:-"${PROJECT_ROOT}/target/debug/meta_discovery_probe"}
PLAYGROUND_DATA_DIR="${TIUP_HOME:-$HOME/.tiup}/data/${PLAYGROUND_TAG}"
PLAYGROUND_LOG_FILE="${TEST_OUTPUT_DIR}/log/playground.log"

playground_display() {
    tiup "playground:${PLAYGROUND_COMPONENT_VERSION}" display --tag "${PLAYGROUND_TAG}" 2>/dev/null || true
}

leader_pid_for_addr() {
    local addr=$1
    local port=${addr##*:}
    lsof -tiTCP:${port} -sTCP:LISTEN | head -n1 || true
}

live_meta_process_count() {
    pgrep -af "${PLAYGROUND_DATA_DIR}/tici-meta-" | wc -l | tr -d ' '
}

cleanup_meta_failover() {
    local exit_code=${1:-0}
    echo "Cleaning up meta failover services..."

    set +e

    local pids=()
    while read -r pid; do
        if [[ "${pid}" =~ ^[0-9]+$ ]]; then
            pids+=("${pid}")
        fi
    done < <(playground_display | awk '$1 ~ /^[0-9]+$/ {print $1}')

    while read -r pid _; do
        if [[ "${pid}" =~ ^[0-9]+$ ]]; then
            pids+=("${pid}")
        fi
    done < <(pgrep -af "${PLAYGROUND_DATA_DIR}" || true)

    if [ -f "${TEST_OUTPUT_DIR}/minio.pid" ]; then
        local minio_pid
        minio_pid=$(cat "${TEST_OUTPUT_DIR}/minio.pid" 2>/dev/null || true)
        if [[ "${minio_pid}" =~ ^[0-9]+$ ]]; then
            pids+=("${minio_pid}")
        fi
    fi

    local deduped_pids=()
    while read -r pid; do
        if [[ "${pid}" =~ ^[0-9]+$ ]]; then
            deduped_pids+=("${pid}")
        fi
    done < <(printf '%s\n' "${pids[@]}" | awk '/^[0-9]+$/ && !seen[$0]++')
    pids=("${deduped_pids[@]}")

    if [ ${#pids[@]} -gt 0 ]; then
        kill -2 "${pids[@]}" >/dev/null 2>&1 || true
        local attempts=0
        while [ ${attempts} -lt 15 ]; do
            local alive=()
            local pid
            for pid in "${pids[@]}"; do
                if kill -0 "${pid}" >/dev/null 2>&1; then
                    alive+=("${pid}")
                fi
            done
            if [ ${#alive[@]} -eq 0 ]; then
                break
            fi
            echo "Waiting for services to stop..."
            sleep 2
            pids=("${alive[@]}")
            attempts=$((attempts + 1))
        done

        if [ ${#pids[@]} -gt 0 ]; then
            kill -9 "${pids[@]}" >/dev/null 2>&1 || true
        fi
    fi

    echo "Cleanup completed"
    return "${exit_code}"
}

trap 'dump_logs; cleanup_meta_failover' EXIT

probe_state() {
    "${PROBE_BIN}" --pd-addr "${PD_ADDR}"
}

extract_leader() {
    printf '%s\n' "$1" | awk -F= '/^leader=/{print $2}'
}

extract_peers() {
    printf '%s\n' "$1" | awk -F= '/^peer=/{print $2}'
}

contains_peer() {
    local state=$1
    local addr=$2
    extract_peers "${state}" | grep -Fxq "${addr}"
}

wait_for_initial_state() {
    local deadline=$((SECONDS + 30))
    while [ ${SECONDS} -lt ${deadline} ]; do
        if [ "$(live_meta_process_count)" -lt 2 ]; then
            sleep 1
            continue
        fi

        local state
        state=$(probe_state)
        local leader
        leader=$(extract_leader "${state}")
        local peer_count
        peer_count=$(extract_peers "${state}" | sed '/^$/d' | wc -l | tr -d ' ')
        if [ -n "${leader}" ] && [ "${peer_count}" -ge 2 ] && [ -n "$(leader_pid_for_addr "${leader}")" ]; then
            printf '%s\n' "${state}"
            return 0
        fi
        sleep 1
    done
    return 1
}

wait_for_leader_switch() {
    local old_leader=$1
    local deadline=$((SECONDS + LEADER_SWITCH_TIMEOUT_SECS))
    while [ ${SECONDS} -lt ${deadline} ]; do
        local state
        state=$(probe_state)
        local leader
        leader=$(extract_leader "${state}")
        if [ -n "${leader}" ] && [ "${leader}" != "${old_leader}" ]; then
            printf '%s\n' "${state}"
            return 0
        fi
        sleep 1
    done
    return 1
}

wait_for_peer_removal() {
    local removed_peer=$1
    local deadline=$((SECONDS + REGISTRY_CLEANUP_TIMEOUT_SECS))
    while [ ${SECONDS} -lt ${deadline} ]; do
        local state
        state=$(probe_state)
        if ! contains_peer "${state}" "${removed_peer}"; then
            printf '%s\n' "${state}"
            return 0
        fi
        sleep 1
    done
    return 1
}

if ! command -v lsof >/dev/null 2>&1; then
    log_error "lsof is required for this script"
    exit 1
fi

if [ ! -x "${PROBE_BIN}" ]; then
    log "Building meta discovery probe..."
    cargo build -p tici_integration_tests --bin meta_discovery_probe
else
    log "Reusing existing meta discovery probe at ${PROBE_BIN}"
fi

log "Cleaning previous playground state under ${PLAYGROUND_DATA_DIR}"
cleanup_meta_failover >/dev/null 2>&1 || true
rm -rf "${PLAYGROUND_DATA_DIR}"
rm -f "${PLAYGROUND_LOG_FILE}"

log "Preparing playground environment..."
bash "${PROJECT_ROOT}/tests/scripts/e2e-prepare.sh"
export PD_ADDR="127.0.0.1:$(get_pd_port)"
dump_configs

log "Waiting for two meta nodes and an elected leader..."
initial_state=$(wait_for_initial_state)
printf '%s\n' "${initial_state}"

old_leader=$(extract_leader "${initial_state}")
if [ -z "${old_leader}" ]; then
    log_error "Leader is empty after playground startup"
    exit 1
fi

leader_port=${old_leader##*:}
leader_pid=$(leader_pid_for_addr "${old_leader}")
if [ -z "${leader_pid}" ]; then
    log_error "Failed to find leader PID for ${old_leader}"
    playground_display || true
    exit 1
fi

log "Crashing current meta leader ${old_leader} (pid=${leader_pid})"
kill -9 "${leader_pid}"

log "Waiting for leader switch..."
leader_state=$(wait_for_leader_switch "${old_leader}") || {
    log_error "Timed out waiting for leader to switch away from ${old_leader}"
    probe_state || true
    exit 1
}
printf '%s\n' "${leader_state}"

new_leader=$(extract_leader "${leader_state}")
if [ -z "${new_leader}" ]; then
    log_error "Leader is empty after failover"
    exit 1
fi

log "Waiting for stale registry entry ${old_leader} to disappear..."
cleanup_state=$(wait_for_peer_removal "${old_leader}") || {
    log_error "Timed out waiting for stale registry entry ${old_leader} to disappear"
    probe_state || true
    exit 1
}
printf '%s\n' "${cleanup_state}"

log "Meta failover succeeded: old leader=${old_leader}, new leader=${new_leader}"