#!/usr/bin/env bash
# Safe end-to-end command lifecycle canary for reboot/shutdown contract v1.
# Verifies ack flow: accepted -> execution_started -> completed/failed.

set -euo pipefail

SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"

BLUE='\033[0;34m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
NC='\033[0m'

ENV_FILE="$PROJECT_ROOT/.env"
if [[ -f "$ENV_FILE" ]]; then
  while IFS='=' read -r key value; do
    key="${key//[$'\t\r\n']}"
    key="$(echo "$key" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//')"
    [[ -z "$key" ]] && continue
    [[ "$key" =~ ^# ]] && continue
    [[ "$key" =~ ^[A-Z_][A-Z0-9_]*$ ]] || continue
    value="${value%%#*}"
    value="${value//[$'\t\r\n']}"
    value="$(echo "$value" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//')"
    export "$key=$value"
  done < "$ENV_FILE"
fi

BROKER="${MQTT_BROKER:-localhost}"
PORT="${MQTT_PORT:-1883}"
MQTT_USERNAME="${MQTT_USERNAME:-}"
MQTT_PASSWORD="${MQTT_PASSWORD:-}"
MQTT_TLS_ENABLED="${MQTT_TLS_ENABLED:-0}"
MQTT_TLS_CA_CERT="${MQTT_TLS_CA_CERT:-}"
MQTT_TLS_CERT="${MQTT_TLS_CERT:-}"
MQTT_TLS_KEY="${MQTT_TLS_KEY:-}"
MQTT_TLS_INSECURE="${MQTT_TLS_INSECURE:-0}"
CLIENT_UUID_FILE="$PROJECT_ROOT/src/config/client_uuid.txt"
LAST_COMMAND_STATE_FILE="$PROJECT_ROOT/src/config/last_command_state.json"

if [[ ! -f "$CLIENT_UUID_FILE" ]]; then
  echo -e "${RED}client UUID file missing: $CLIENT_UUID_FILE${NC}"
  exit 1
fi

if ! command -v mosquitto_pub >/dev/null 2>&1 || ! command -v mosquitto_sub >/dev/null 2>&1; then
  echo -e "${RED}mosquitto_pub/sub not found. Install mosquitto-clients.${NC}"
  exit 1
fi

CLIENT_UUID="$(tr -d '[:space:]' < "$CLIENT_UUID_FILE")"
COMMAND_TOPIC="infoscreen/${CLIENT_UUID}/commands"
COMMAND_TOPIC_ALIAS="infoscreen/${CLIENT_UUID}/command"
ACK_TOPIC="infoscreen/${CLIENT_UUID}/commands/ack"
ACK_TOPIC_ALIAS="infoscreen/${CLIENT_UUID}/command/ack"

MQTT_AUTH_ARGS=()
MQTT_TLS_ARGS=()

if [[ -n "$MQTT_USERNAME" ]]; then
  MQTT_AUTH_ARGS+=( -u "$MQTT_USERNAME" )
fi
if [[ -n "$MQTT_PASSWORD" ]]; then
  MQTT_AUTH_ARGS+=( -P "$MQTT_PASSWORD" )
fi
if [[ "$MQTT_TLS_ENABLED" == "1" || "$MQTT_TLS_ENABLED" == "true" || "$MQTT_TLS_ENABLED" == "yes" ]]; then
  [[ -n "$MQTT_TLS_CA_CERT" ]] && MQTT_TLS_ARGS+=( --cafile "$MQTT_TLS_CA_CERT" )
  [[ -n "$MQTT_TLS_CERT" ]] && MQTT_TLS_ARGS+=( --cert "$MQTT_TLS_CERT" )
  [[ -n "$MQTT_TLS_KEY" ]] && MQTT_TLS_ARGS+=( --key "$MQTT_TLS_KEY" )
  if [[ "$MQTT_TLS_INSECURE" == "1" || "$MQTT_TLS_INSECURE" == "true" || "$MQTT_TLS_INSECURE" == "yes" ]]; then
    MQTT_TLS_ARGS+=( --insecure )
  fi
fi

ACTION="${1:-reboot_host}"
MODE="${2:-success}"  # success | failed
TOPIC_MODE="${3:-canonical}" # canonical | alias
WAIT_SEC="${4:-25}"

if [[ "$ACTION" != "reboot_host" && "$ACTION" != "shutdown_host" ]]; then
  echo -e "${RED}invalid action '$ACTION' (expected reboot_host|shutdown_host)${NC}"
  exit 1
fi
if [[ "$MODE" != "success" && "$MODE" != "failed" ]]; then
  echo -e "${RED}invalid mode '$MODE' (expected success|failed)${NC}"
  exit 1
fi
if [[ "$TOPIC_MODE" != "canonical" && "$TOPIC_MODE" != "alias" ]]; then
  echo -e "${RED}invalid topic mode '$TOPIC_MODE' (expected canonical|alias)${NC}"
  exit 1
fi
if ! [[ "$WAIT_SEC" =~ ^[0-9]+$ ]] || [[ "$WAIT_SEC" -lt 1 ]]; then
  echo -e "${RED}invalid wait seconds '$WAIT_SEC' (expected positive integer)${NC}"
  exit 1
fi

if [[ "$TOPIC_MODE" == "alias" ]]; then
  COMMAND_TOPIC="$COMMAND_TOPIC_ALIAS"
fi

COMMAND_ID="$(python3 - <<'PY'
import uuid
print(uuid.uuid4())
PY
)"
ISSUED_AT="$(date -u +"%Y-%m-%dT%H:%M:%SZ")"
EXPIRES_EPOCH="$(( $(date +%s) + 240 ))"
EXPIRES_AT="$(date -u -d "@$EXPIRES_EPOCH" +"%Y-%m-%dT%H:%M:%SZ")"

PAYLOAD="$(cat <<EOF
{
  "schema_version": "1.0",
  "command_id": "$COMMAND_ID",
  "client_uuid": "$CLIENT_UUID",
  "action": "$ACTION",
  "issued_at": "$ISSUED_AT",
  "expires_at": "$EXPIRES_AT",
  "requested_by": 1,
  "reason": "canary_test"
}
EOF
)"

TMP_ACK_LOG="$(mktemp)"
cleanup() {
  [[ -n "${SUB_PID_1:-}" ]] && kill "$SUB_PID_1" >/dev/null 2>&1 || true
  [[ -n "${SUB_PID_2:-}" ]] && kill "$SUB_PID_2" >/dev/null 2>&1 || true
  rm -f "$TMP_ACK_LOG"
}
trap cleanup EXIT

echo -e "${BLUE}================================================${NC}"
echo -e "${BLUE}Command Lifecycle Canary${NC}"
echo -e "${BLUE}================================================${NC}"
echo " Broker      : $BROKER:$PORT"
echo " Client UUID : $CLIENT_UUID"
echo " Command ID  : $COMMAND_ID"
echo " Action      : $ACTION"
echo " Mode        : $MODE"
echo " Cmd Topic   : $COMMAND_TOPIC"
echo " Ack Topics  : $ACK_TOPIC , $ACK_TOPIC_ALIAS"
echo ""
echo -e "${YELLOW}IMPORTANT${NC}: to avoid real reboot/shutdown, run simclient with"
echo "  COMMAND_HELPER_PATH=$PROJECT_ROOT/scripts/mock-command-helper.sh"
echo ""

# Subscribe first to avoid missing retained/non-retained race windows.
mosquitto_sub -h "$BROKER" -p "$PORT" "${MQTT_AUTH_ARGS[@]}" "${MQTT_TLS_ARGS[@]}" -q 1 -v -t "$ACK_TOPIC" >> "$TMP_ACK_LOG" &
SUB_PID_1=$!
mosquitto_sub -h "$BROKER" -p "$PORT" "${MQTT_AUTH_ARGS[@]}" "${MQTT_TLS_ARGS[@]}" -q 1 -v -t "$ACK_TOPIC_ALIAS" >> "$TMP_ACK_LOG" &
SUB_PID_2=$!
sleep 0.5

if [[ "$MODE" == "failed" ]]; then
  echo -e "${YELLOW}If simclient was started with MOCK_COMMAND_HELPER_FORCE_FAIL=1, expected terminal status is failed.${NC}"
fi

echo -e "${YELLOW}Publishing command payload...${NC}"
mosquitto_pub -h "$BROKER" -p "$PORT" "${MQTT_AUTH_ARGS[@]}" "${MQTT_TLS_ARGS[@]}" -q 1 -t "$COMMAND_TOPIC" -m "$PAYLOAD"

EXPECTED_TERMINAL="completed"
if [[ "$MODE" == "failed" ]]; then
  EXPECTED_TERMINAL="failed"
fi

EXPECT_RECOVERY_COMPLETION=0
if [[ "$ACTION" == "reboot_host" && "$MODE" == "success" ]]; then
  EXPECTED_TERMINAL=""
  EXPECT_RECOVERY_COMPLETION=1
fi

DEADLINE=$(( $(date +%s) + WAIT_SEC ))
SEEN_ACCEPTED=0
SEEN_STARTED=0
SEEN_TERMINAL=0

while [[ $(date +%s) -lt $DEADLINE ]]; do
  if grep -q '"command_id"' "$TMP_ACK_LOG" 2>/dev/null; then
    if grep -q "\"command_id\": \"$COMMAND_ID\"" "$TMP_ACK_LOG"; then
      grep -q '"status": "accepted"' "$TMP_ACK_LOG" && SEEN_ACCEPTED=1 || true
      grep -q '"status": "execution_started"' "$TMP_ACK_LOG" && SEEN_STARTED=1 || true
      if [[ -n "$EXPECTED_TERMINAL" ]]; then
        grep -q "\"status\": \"$EXPECTED_TERMINAL\"" "$TMP_ACK_LOG" && SEEN_TERMINAL=1 || true
      fi
    fi
  fi

  if [[ $SEEN_ACCEPTED -eq 1 && $SEEN_STARTED -eq 1 ]]; then
    if [[ -z "$EXPECTED_TERMINAL" || $SEEN_TERMINAL -eq 1 ]]; then
      break
    fi
  fi
  sleep 1
done

echo ""
echo -e "${BLUE}Ack stream (filtered by command_id):${NC}"
python3 - <<'PY' "$TMP_ACK_LOG" "$COMMAND_ID"
import json
import sys

path, command_id = sys.argv[1], sys.argv[2]
with open(path, "r", encoding="utf-8", errors="ignore") as f:
    for line in f:
        line = line.strip()
        if not line:
            continue
        parts = line.split(" ", 1)
        payload = parts[1] if len(parts) == 2 else parts[0]
        try:
            obj = json.loads(payload)
        except Exception:
            continue
        if obj.get("command_id") == command_id:
            print(json.dumps(obj, indent=2))
PY

if [[ $SEEN_ACCEPTED -eq 1 && $SEEN_STARTED -eq 1 ]]; then
  if [[ -z "$EXPECTED_TERMINAL" ]]; then
    echo -e "${GREEN}PASS${NC}: observed accepted -> execution_started"
    echo -e "${YELLOW}NOTE${NC}: completed for reboot_host is expected only after client reconnect/recovery."
  elif [[ $SEEN_TERMINAL -eq 1 ]]; then
    echo -e "${GREEN}PASS${NC}: observed accepted -> execution_started -> $EXPECTED_TERMINAL"
  else
    echo -e "${RED}FAIL${NC}: missing expected terminal state $EXPECTED_TERMINAL for command_id=$COMMAND_ID"
    exit 1
  fi
else
  echo -e "${RED}FAIL${NC}: missing expected lifecycle states for command_id=$COMMAND_ID"
  if [[ -n "$EXPECTED_TERMINAL" ]]; then
    echo " observed: accepted=$SEEN_ACCEPTED execution_started=$SEEN_STARTED terminal($EXPECTED_TERMINAL)=$SEEN_TERMINAL"
  else
    echo " observed: accepted=$SEEN_ACCEPTED execution_started=$SEEN_STARTED"
  fi
  exit 1
fi

if [[ -f "$LAST_COMMAND_STATE_FILE" ]]; then
  echo ""
  echo -e "${BLUE}Last command state:${NC}"
  python3 -m json.tool "$LAST_COMMAND_STATE_FILE" || cat "$LAST_COMMAND_STATE_FILE"
fi