#!/usr/bin/env bash

SESSION="cs425-debug"
CMDS_WINDOW="cmds"

NUM_SERVERS=10

PORT=8989

WORKDIR="fa26_cs425_.group-mp_.team-01/MP2"
SCRIPT="peer.py"

SSH_OPTS="-o BatchMode=yes \
-o ConnectTimeout=2 \
-o ServerAliveInterval=1 \
-o ServerAliveCountMax=1"


# HELPERS


expand_targets() {

    if [ "$1" = "all" ]; then
        seq 1 "$NUM_SERVERS"
    else
        echo "$@"
    fi
}


get_pane() {

    local vm="$1"

    # VM1  -> pane 0
    # VM2  -> pane 1

    echo "$SESSION:$CMDS_WINDOW.$((vm - 1))"
}


is_running() {

    local vm="$1"

    ssh $SSH_OPTS "server$vm" \
        "pgrep -f '[p]ython3.*$SCRIPT' >/dev/null" \
        2>/dev/null
}



# START


start_vm() {

    local vm="$1"
    local pane

    pane=$(get_pane "$vm")

    if ! tmux has-session \
        -t "$SESSION" \
        2>/dev/null
    then
        echo "VM$vm: tmux session not running"
        return 1
    fi

    echo "starting VM$vm..."

    # Kill old process on remote VM

    ssh $SSH_OPTS "server$vm" \
        "fuser -k -9 $PORT/udp 2>/dev/null || true; \
         fuser -k -9 $PORT/tcp 2>/dev/null || true; \
         pkill -9 -f '[p]ython3.*$SCRIPT' 2>/dev/null || true" \
        >/dev/null 2>&1

    sleep 0.3

    # Reset pane

    tmux send-keys \
        -t "$pane" \
        C-c

    sleep 0.05

    tmux send-keys \
        -t "$pane" \
        "clear" \
        C-m

    # Start peer.py

    tmux send-keys \
        -t "$pane" \
        "ssh -t $SSH_OPTS server$vm 'cd \$HOME/$WORKDIR && MACHINE_INDEX=$vm python3 $SCRIPT'" \
        C-m
}



# KILL


kill_vm() {

    local vm="$1"

    echo "killing VM$vm..."

    ssh $SSH_OPTS "server$vm" \
        "fuser -k -9 $PORT/udp 2>/dev/null || true; \
         fuser -k -9 $PORT/tcp 2>/dev/null || true; \
         pkill -9 -f '[p]ython3.*$SCRIPT' 2>/dev/null || true" \
        >/dev/null 2>&1
}



# COMMAND


cmd_to_vm() {

    local vm="$1"
    shift

    local command="$*"
    local pane
    local pane_command

    pane=$(get_pane "$vm")

    if ! is_running "$vm"; then
        echo "VM$vm skipped (STOPPED)"
        return
    fi

    pane_command=$(
        tmux display-message \
            -p \
            -t "$pane" \
            '#{pane_current_command}' \
            2>/dev/null
    )

    if [ "$pane_command" != "ssh" ]; then
        echo "VM$vm skipped (pane not attached)"
        return
    fi

    echo "VM$vm <- $command"

    tmux send-keys \
        -t "$pane" \
        "$command" \
        C-m
}



# LIST


list_vm() {

    local vm="$1"

    echo
    echo "========== VM$vm =========="

    ssh $SSH_OPTS "server$vm" \
        "cat \$HOME/$WORKDIR/membership_list.txt" \
        2>/dev/null
}



# LOG


log_vm() {

    local vm="$1"

    echo
    echo "========== VM$vm =========="

    ssh $SSH_OPTS "server$vm" \
        "tail -n 30 \$HOME/$WORKDIR/machine.$vm.log" \
        2>/dev/null
}



# STATUS


status_vm() {

    local vm="$1"
    local pid

    pid=$(
        ssh $SSH_OPTS "server$vm" \
            "pgrep -f '[p]ython3.*$SCRIPT' | tr '\n' ' '" \
            2>/dev/null
    )

    if [ -n "$pid" ]; then
        echo "VM$vm RUNNING, PID(s): $pid"
    else
        echo "VM$vm STOPPED"
    fi
}



# MAIN


action="$1"
shift || true

case "$action" in


# START


start)

    if [ "$#" -eq 0 ]; then
        echo "Usage: ./cmd.sh start <all|N...>"
        exit 1
    fi

    targets=$(expand_targets "$@")

    for vm in $targets; do
        start_vm "$vm" &
    done

    wait

    sleep 0.5

    for vm in $targets; do
        status_vm "$vm" &
    done

    wait
    ;;



# KILL


kill)

    if [ "$#" -eq 0 ]; then
        echo "Usage: ./cmd.sh kill <all|N...>"
        exit 1
    fi

    targets=$(expand_targets "$@")

    for vm in $targets; do
        kill_vm "$vm" &
    done

    wait
    ;;



# CMD


cmd)

    target="$1"
    shift || true

    if [ -z "$target" ] || [ "$#" -eq 0 ]; then
        echo "Usage:"
        echo "./cmd.sh cmd <all|N> \"command\""
        exit 1
    fi

    if [ "$target" = "all" ]; then

        for vm in $(seq 1 "$NUM_SERVERS"); do
            cmd_to_vm "$vm" "$@" &
        done

        wait

    else

        cmd_to_vm "$target" "$@"

    fi
    ;;



# LIST


list)

    if [ "$#" -eq 0 ]; then
        targets=$(seq 1 "$NUM_SERVERS")
    else
        targets=$(expand_targets "$@")
    fi

    for vm in $targets; do
        list_vm "$vm" &
    done

    wait
    ;;



# LOG


log)

    if [ "$#" -eq 0 ]; then
        targets=$(seq 1 "$NUM_SERVERS")
    else
        targets=$(expand_targets "$@")
    fi

    for vm in $targets; do
        log_vm "$vm" &
    done

    wait
    ;;



# STATUS


status)

    if [ "$#" -eq 0 ]; then
        targets=$(seq 1 "$NUM_SERVERS")
    else
        targets=$(expand_targets "$@")
    fi

    for vm in $targets; do
        status_vm "$vm" &
    done

    wait
    ;;



# HELP


*)

    echo
    echo "Usage:"
    echo
    echo "  ./cmd.sh start  <all|N...>"
    echo "  ./cmd.sh kill   <all|N...>"
    echo "  ./cmd.sh cmd    <all|N> \"command\""
    echo "  ./cmd.sh list   [all|N...]"
    echo "  ./cmd.sh log    [all|N...]"
    echo "  ./cmd.sh status [all|N...]"
    echo
    echo "Examples:"
    echo
    echo "  ./cmd.sh start 1"
    echo "  ./cmd.sh start 2 3 4 5"
    echo "  ./cmd.sh start 6 7 8 9 10"
    echo
    echo "  ./cmd.sh cmd all \"switch suspect\""
    echo "  ./cmd.sh cmd all \"switch nosuspect\""
    echo "  ./cmd.sh cmd all \"set_drop_rate 10\""
    echo "  ./cmd.sh cmd all \"display_protocol\""
    echo
    echo "  ./cmd.sh kill 8 9"
    echo "  ./cmd.sh log 8 9"
    echo
    exit 1
    ;;
esac