#!/bin/bash -eu
# =============================================================================
# 02-migrate-peers.sh — Run as Linux user after 01
# 1. Shuts down legacy WireGuard
# 2. Injects legacy private key into new wg0.conf
# 3. Starts new WireGuard
# 4. Verifies public key is identical
# 5. Migrates peers via API
# 6. Destroys private key
# Auto-rollback: restores legacy if anything fails
# =============================================================================
if [ "$EUID" -eq 0 ]; then
    echo "[X] Run as your Linux user, not root"
    exit 1
fi

RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; CYAN='\033[0;36m'; RESET='\033[0m'

BACKUP_DIR="$HOME/.wg-migration"
LOG="$BACKUP_DIR/migration.log"
OPERATOR_DIR="$HOME/operator-docker"
WG_CONF="$OPERATOR_DIR/services/wireguard/config/wg_confs/wg0.conf"
WG_CONF_BAK=""

ok()   { echo -e "${GREEN}[OK]${RESET} $1";  echo "[OK] $1"  >> "$LOG"; }
info() { echo -e "${CYAN}[i]${RESET} $1";   echo "[i] $1"   >> "$LOG"; }
warn() { echo -e "${YELLOW}[!]${RESET} $1"; echo "[!] $1"   >> "$LOG"; }
fail() { echo -e "${RED}[X]${RESET} $1";    echo "[X] $1"   >> "$LOG"; exit 1; }
sep()  { echo -e "\n${CYAN}------------------------------------${RESET}"; }

rollback_02() {
    echo ""
    warn "Error detected — starting rollback..."
    echo "[ROLLBACK-02] $(date '+%Y-%m-%d %H:%M:%S')" >> "$LOG"

    [ -n "$WG_CONF_BAK" ] && [ -f "$WG_CONF_BAK" ] && \
        cp "$WG_CONF_BAK" "$WG_CONF" && warn "wg0.conf restored."

    cd "$OPERATOR_DIR" 2>/dev/null || true
    docker compose stop wireguard 2>/dev/null || true

    LEGACY_COMPOSE_DIR=$(cat "$BACKUP_DIR/legacy_compose_dir" 2>/dev/null || true)
    if [ -n "$LEGACY_COMPOSE_DIR" ] && [ -d "$LEGACY_COMPOSE_DIR" ]; then
        sudo modprobe ip6table_filter 2>/dev/null || true
        cd "$LEGACY_COMPOSE_DIR"
        docker compose up -d 2>/dev/null || true
        sleep 5
        LEGACY_UP=$(docker ps --format '{{.Names}}' | grep -i wireguard | head -1 || true)
        [ -n "$LEGACY_UP" ] && ok "Legacy restored: $LEGACY_UP" || \
            warn "Legacy failed to start. Run manually: cd $LEGACY_COMPOSE_DIR && docker compose up -d"
    fi
    warn "Rollback complete. Check: cat $LOG"
}
trap rollback_02 ERR

echo ""
echo "========================================="
echo "  02 - Identity and peer migration"
echo "  $(date '+%Y-%m-%d %H:%M:%S')"
echo "========================================="
echo "[START-02] $(date '+%Y-%m-%d %H:%M:%S')" >> "$LOG"

# ─── PREREQUISITES ────────────────────────────────────────────────────────────
sep
info "Checking prerequisites..."

docker info &>/dev/null || fail "No Docker access. Run: sudo usermod -aG docker $USER && newgrp docker"

[ -f "$BACKUP_DIR/private_key" ]        || fail "No private key backup found."
[ -f "$BACKUP_DIR/wg0-legacy.conf" ]    || fail "No peer backup found."
[ -f "$BACKUP_DIR/public_key" ]         || fail "No public key backup found."
[ -f "$BACKUP_DIR/legacy_compose_dir" ] || fail "No legacy snapshot found."
ok "Backups found."

[ -d "$OPERATOR_DIR" ] || fail "operator-docker not found. Run 01 first."
[ -f "$WG_CONF" ]      || fail "wg0.conf not found: $WG_CONF"
ok "operator-docker found."

for CONTAINER in operator-api mysql; do
    STATUS=$(docker inspect -f '{{.State.Running}}' "$CONTAINER" 2>/dev/null || echo "false")
    [ "$STATUS" = "true" ] || fail "Container '$CONTAINER' is not running."
    ok "Running: $CONTAINER"
done

# ─── LOAD CREDENTIALS ─────────────────────────────────────────────────────────
sep
info "Loading credentials..."
set -a; source "$OPERATOR_DIR/.env"; set +a

[ -z "${OPERATOR_API_TOKEN:-}" ] && fail "OPERATOR_API_TOKEN not found in .env"
[ -z "${BASE_DOMAIN:-}" ]        && fail "BASE_DOMAIN not found in .env"
[ -z "${API_SUBDOMAIN:-}" ]      && fail "API_SUBDOMAIN not found in .env"

echo "$BASE_DOMAIN" | grep -qE '[/;|&<> ]|^[0-9]+-' && \
    fail "Invalid BASE_DOMAIN in .env: '$BASE_DOMAIN'"

API_URL="https://${API_SUBDOMAIN}.${BASE_DOMAIN}"
ok "API URL: $API_URL"

# ─── STOP LEGACY WIREGUARD ────────────────────────────────────────────────────
sep
info "Stopping legacy WireGuard..."
LEGACY_COMPOSE_DIR=$(cat "$BACKUP_DIR/legacy_compose_dir")
[ -d "$LEGACY_COMPOSE_DIR" ] || fail "Legacy directory not found: $LEGACY_COMPOSE_DIR"

cd "$LEGACY_COMPOSE_DIR"
docker compose stop wireguard 2>/dev/null || true
sleep 3
STILL_UP=$(docker compose ps --services --filter "status=running" 2>/dev/null | grep -i wireguard || true)
[ -n "$STILL_UP" ] && fail "Legacy WireGuard is still running."
ok "Legacy WireGuard stopped."

# ─── BACKUP NEW CONF ──────────────────────────────────────────────────────────
sep
info "Backing up new wg0.conf..."
WG_CONF_BAK="$WG_CONF.bak.$(date +%s)"
cp "$WG_CONF" "$WG_CONF_BAK"
ok "Backup: $WG_CONF_BAK"

# ─── INJECT PRIVATE KEY ───────────────────────────────────────────────────────
sep
info "Injecting legacy private key..."
LEGACY_PRIVATE_KEY=$(cat "$BACKUP_DIR/private_key" | tr -d '[:space:]')
LEGACY_PUBLIC_KEY=$(cat "$BACKUP_DIR/public_key" | tr -d '[:space:]')
[ -z "$LEGACY_PRIVATE_KEY" ] && fail "Private key is empty."
[ -z "$LEGACY_PUBLIC_KEY" ]  && fail "Public key is empty."

sed -i "s|^PrivateKey = .*|PrivateKey = $LEGACY_PRIVATE_KEY|" "$WG_CONF"
ok "Private key written to wg0.conf."

WG_SERVER_KEY="$OPERATOR_DIR/services/wireguard/config/server/privatekey-server"
WG_SERVER_PUB="$OPERATOR_DIR/services/wireguard/config/server/publickey-server"
mkdir -p "$(dirname "$WG_SERVER_KEY")"
echo "$LEGACY_PRIVATE_KEY" > "$WG_SERVER_KEY"
echo "$LEGACY_PUBLIC_KEY"  > "$WG_SERVER_PUB"
ok "Private key written to server keyfile."

# ─── START NEW WIREGUARD ──────────────────────────────────────────────────────
sep
info "Starting new WireGuard with legacy identity..."
echo "[WIREGUARD-START] $(date '+%Y-%m-%d %H:%M:%S')" >> "$LOG"

cd "$OPERATOR_DIR"
sudo modprobe ip6table_filter 2>/dev/null || true
docker compose up -d wireguard
info "Waiting for startup..."
sleep 8

STATUS=$(docker inspect -f '{{.State.Running}}' wireguard 2>/dev/null || echo "false")
[ "$STATUS" = "true" ] || fail "New WireGuard failed to start. Check: docker logs wireguard"
ok "New WireGuard running."

# ─── VERIFY PUBLIC KEY ────────────────────────────────────────────────────────
sep
info "Verifying server identity..."

NEW_PUBLIC_KEY=$(docker exec wireguard cat /config/server/publickey-server 2>/dev/null | \
    tr -d '[:space:]' || true)
if [ -z "$NEW_PUBLIC_KEY" ]; then
    NEW_PUBLIC_KEY=$(docker exec wireguard wg show wg0 2>/dev/null | \
        grep "public key" | awk '{print $NF}' | tr -d '[:space:]' || true)
fi
[ -z "$NEW_PUBLIC_KEY" ] && fail "Could not retrieve public key from new WireGuard."

if [ "$NEW_PUBLIC_KEY" != "$LEGACY_PUBLIC_KEY" ]; then
    warn "Public key MISMATCH!"
    warn "  Expected : $LEGACY_PUBLIC_KEY"
    warn "  Got      : $NEW_PUBLIC_KEY"
    fail "Server identity was not preserved."
fi

ok "Public key verified: $NEW_PUBLIC_KEY"
ok "Existing users will reconnect automatically."

# ─── VERIFY API ───────────────────────────────────────────────────────────────
sep
info "Verifying API..."
HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" --max-time 15 \
    -H "Authorization: Bearer $OPERATOR_API_TOKEN" \
    "$API_URL/wireguard/peers" 2>/dev/null || echo "000")
[ "$HTTP_CODE" = "200" ] || fail "API not responding (HTTP $HTTP_CODE)."
ok "API responding."

# ─── MIGRATE PEERS ────────────────────────────────────────────────────────────
sep
info "Migrating peers to new API..."

LEGACY_CONF=$(cat "$BACKUP_DIR/wg0-legacy.conf")
MIGRATED=0
FAILED=0
CURRENT_PUBKEY=""
CURRENT_IPS=""

migrate_peer() {
    local PUBKEY="$1"
    local IPS="$2"

    IPS_JSON=$(echo "$IPS" | tr ',' '\n' | \
        sed 's/^[[:space:]]*//' | sed 's/[[:space:]]*$//' | \
        awk 'NF {print "\"" $0 "\""}' | paste -sd ',' | \
        sed 's/^/[/' | sed 's/$/]/')

    HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" --max-time 15 \
        -X POST "$API_URL/wireguard/peers" \
        -H "Authorization: Bearer $OPERATOR_API_TOKEN" \
        -H "Content-Type: application/json" \
        -d "{\"public_key\": \"$PUBKEY\", \"allowed_ips\": $IPS_JSON}" 2>/dev/null || echo "000")

    if [ "$HTTP_CODE" = "200" ]; then
        ok "Peer migrated: ${PUBKEY:0:20}... → $IPS"
        MIGRATED=$((MIGRATED + 1))
    else
        warn "Peer failed: ${PUBKEY:0:20}... (HTTP $HTTP_CODE)"
        FAILED=$((FAILED + 1))
    fi
}

while IFS= read -r line; do
    if [[ "$line" == "[Peer]" ]]; then
        if [ -n "$CURRENT_PUBKEY" ] && [ -n "$CURRENT_IPS" ]; then
            migrate_peer "$CURRENT_PUBKEY" "$CURRENT_IPS"
        fi
        CURRENT_PUBKEY=""
        CURRENT_IPS=""
    elif [[ "$line" =~ ^PublicKey\ =\ (.+)$ ]]; then
        CURRENT_PUBKEY="${BASH_REMATCH[1]}"
    elif [[ "$line" =~ ^AllowedIPs\ =\ (.+)$ ]]; then
        CURRENT_IPS="${BASH_REMATCH[1]}"
    fi
done <<< "$LEGACY_CONF"

[ -n "$CURRENT_PUBKEY" ] && [ -n "$CURRENT_IPS" ] && migrate_peer "$CURRENT_PUBKEY" "$CURRENT_IPS"

# ─── VERIFY PEERS ─────────────────────────────────────────────────────────────
sep
info "Verifying peers in WireGuard..."
REGISTERED=$(docker exec wireguard wg show wg0 | grep -c "^  peer:" 2>/dev/null || echo 0)
ok "$REGISTERED peers active in wg0."

# ─── DESTROY PRIVATE KEY ──────────────────────────────────────────────────────
sep
info "Destroying private key..."
command -v shred &>/dev/null && \
    shred -u "$BACKUP_DIR/private_key" 2>/dev/null || \
    rm -f "$BACKUP_DIR/private_key"
ok "Private key deleted from disk."

trap - ERR
echo "[END-02] $(date '+%Y-%m-%d %H:%M:%S')" >> "$LOG"

echo ""
echo "========================================="
echo "  Migration complete"
echo "========================================="
echo ""
echo "  Public key     : $NEW_PUBLIC_KEY"
echo "  API            : $API_URL"
echo "  Peers migrated : $MIGRATED"
if [ "$FAILED" -gt 0 ]; then
    echo -e "  ${RED}Peers failed   : $FAILED${RESET}"
    echo "  Check: cat $LOG"
else
    echo "  Errors         : none"
fi
echo ""
echo "  Existing users reconnect automatically"
echo "  without new configuration."
echo ""
echo "  Log: cat $LOG"
echo "========================================="