#!/bin/bash
# scripts/deploy_to_node.sh
# OAuth-First Guided Onboarding for Jetson Orin Nano
# Flow: get_mesh_auth.py → manifest.json → rsync + --exclude="venv" --exclude=".venv" --exclude="__pycache__" --exclude=".git" --copy-unsafe-links SSH headless provisioning → polling loop

set -euo pipefail

# ============================================================================
# Configuration
# ============================================================================

TARGET_USER="howsa"
TARGET_HOSTNAME="192.168.x.x"
SOURCE_DIR="$HOME/engram/"
TARGET_DIR="~/engram/"
MANIFEST_PATH="$HOME/.engram/manifest.json"
ENV_FILE="$HOME/engram/.env.tailscale"

# Color codes
RED='\033[0;31m'
GREEN='\033[0;32m'
BLUE='\033[0;34m'
YELLOW='\033[0;33m'
CYAN='\033[0;36m'
NC='\033[0m'

# Logging functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_pass() {
    echo -e "${GREEN}[✓]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[!]${NC} $1"
}

log_progress() {
    echo -e "${CYAN}[→]${NC} $1"
}

log_error() {
    echo -e "${RED}[✗]${NC} $1"
    exit 1
}

# ============================================================================
# Pre-Flight Checks
# ============================================================================

log_info "═══════════════════════════════════════════════════════════════"
log_info "Engram OAuth-First Deployment — Jetson Orin Nano"
log_info "═══════════════════════════════════════════════════════════════"
echo ""

log_progress "Pre-flight checks..."

# Check manifest exists
if [[ ! -f "$MANIFEST_PATH" ]]; then
    log_error "Manifest not found at $MANIFEST_PATH"
    log_error "Run: python3 scripts/get_mesh_auth.py"
fi

log_pass "Manifest found: $MANIFEST_PATH"

# Check .env.tailscale exists
if [[ ! -f "$ENV_FILE" ]]; then
    log_error "Auth key not found at $ENV_FILE"
    log_error "Run: python3 scripts/get_mesh_auth.py"
fi

log_pass "Auth key file found: $ENV_FILE"

# Source the auth key
source "$ENV_FILE"
if [[ -z "${TAILSCALE_AUTHKEY:-}" ]]; then
    log_error "TAILSCALE_AUTHKEY is empty"
fi

log_pass "Auth key loaded"

# Check SSH connectivity
log_progress "Testing SSH connectivity to Jetson ($TARGET_HOSTNAME)..."
if ! ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=no "${TARGET_USER}@${TARGET_HOSTNAME}" exit 2>/dev/null; then
    log_error "Cannot reach Jetson at $TARGET_HOSTNAME (SSH failed)"
fi

log_pass "SSH connectivity verified"

# Extract worker_id from manifest
WORKER_ID=$(jq -r '.worker_id' "$MANIFEST_PATH")
log_pass "Worker ID: $WORKER_ID"

echo ""

# ============================================================================
# Phase 1: Rsync Source Code + Manifest
# ============================================================================

log_info "Phase 1: Transferring Source Code & Manifest"
log_progress "Syncing source code to Jetson..."

rsync -aqz \
  --exclude='venv' \
  --exclude='.venv' \
  --exclude='__pycache__' \
  --exclude='.git' \
  --exclude='REFLECTIONS' \
  --exclude='MEMORIES' \
  --exclude='node_modules' \
  --exclude='.env' \
  --exclude='.env.tailscale' \
  --exclude='REPORTS' \
  --exclude='.claude' \
  --exclude='.mypy_cache' \
  --exclude='.pytest_cache' \
  --exclude='.ruff_cache' \
  --exclude='logs' \
  --exclude='artifacts' \
  --exclude='synapse_data' \
  --exclude='matrix' \
  "$SOURCE_DIR" "${TARGET_USER}@${TARGET_HOSTNAME}:${TARGET_DIR}" || log_error "rsync failed"

log_pass "Source code transferred"

# Copy manifest.json to Jetson
log_progress "Transferring manifest.json to Jetson..."
ssh -o StrictHostKeyChecking=no "${TARGET_USER}@${TARGET_HOSTNAME}" "mkdir -p ~/.engram" || log_warning "Could not create ~/.engram on Jetson"

scp -o StrictHostKeyChecking=no "$MANIFEST_PATH" "${TARGET_USER}@${TARGET_HOSTNAME}:~/.engram/manifest.json" || log_error "scp manifest failed"

log_pass "Manifest transferred to ~/.engram/manifest.json"

echo ""

# ============================================================================
# Phase 2: Remote Provisioning (Headless Mode)
# ============================================================================

log_info "Phase 2: Remote Provisioning (Headless Mode)"
log_progress "Triggering prep_node.sh on Jetson with auth key..."

# Execute prep_node.sh remotely in headless mode
# Using nohup to handle potential reboot disconnections
ssh -o StrictHostKeyChecking=no "${TARGET_USER}@${TARGET_HOSTNAME}" \
    "nohup sudo TAILSCALE_AUTHKEY='${TAILSCALE_AUTHKEY}' bash ~/engram/scripts/prep_node.sh > /tmp/prep_node.log 2>&1 &" \
    || log_warning "SSH returned non-zero (this is expected if Jetson reboots)"

log_progress "Provisioning script started on Jetson (may include reboot)"

echo ""

# ============================================================================
# Phase 3: Polling Loop — Wait for Jetson to Join Mesh
# ============================================================================

log_info "Phase 3: Waiting for Jetson to Join Tailscale Mesh"
log_progress "Polling tailscale status for $WORKER_ID (up to 120 seconds)..."

POLL_TIMEOUT=120
POLL_INTERVAL=5
ELAPSED=0

JETSON_FOUND=false

while [[ $ELAPSED -lt $POLL_TIMEOUT ]]; do
    if tailscale status 2>/dev/null | grep -q "$WORKER_ID"; then
        JETSON_FOUND=true
        break
    fi

    echo -ne "\r${CYAN}[→]${NC} Waiting... ${ELAPSED}s / ${POLL_TIMEOUT}s"
    sleep $POLL_INTERVAL
    ELAPSED=$((ELAPSED + POLL_INTERVAL))
done

echo ""

if [[ "$JETSON_FOUND" != "true" ]]; then
    log_warning "Jetson not found in tailscale status after ${POLL_TIMEOUT}s"
    log_warning "Jetson may still be provisioning. Check manually:"
    echo "  tailscale status | grep $WORKER_ID"
    exit 1
fi

log_pass "Jetson appeared in mesh: $WORKER_ID"

echo ""

# ============================================================================
# Phase 4: Verification
# ============================================================================

log_info "Phase 4: Verification"

# Get Jetson's Tailscale IP
JETSON_IP=$(tailscale status 2>/dev/null | grep "$WORKER_ID" | awk '{print $1}' || echo "")

if [[ -z "$JETSON_IP" ]]; then
    log_warning "Could not extract Jetson's Tailscale IP"
    JETSON_IP="<check manually>"
fi

log_progress "Jetson Tailscale IP: $JETSON_IP"

# Ping over Tailscale
if [[ "$JETSON_IP" != "<check manually>" ]]; then
    log_progress "Pinging Jetson over Tailscale mesh..."
    if ping -c 1 -W 2 "$JETSON_IP" &>/dev/null; then
        log_pass "Jetson is reachable at $JETSON_IP"
    else
        log_warning "Ping failed (Jetson may still be booting)"
    fi
fi

echo ""

# ============================================================================
# Phase 5: Cleanup
# ============================================================================

log_info "Phase 5: Cleanup"

log_progress "Removing .env.tailscale (auth key no longer needed)..."
rm -f "$ENV_FILE"
log_pass "Auth key cleaned up"

echo ""

# ============================================================================
# Final Summary
# ============================================================================

log_info "═══════════════════════════════════════════════════════════════"
log_pass "Deployment Complete"
log_info "═══════════════════════════════════════════════════════════════"

echo ""
echo -e "${CYAN}Status Summary:${NC}"
echo "  Worker ID:              $WORKER_ID"
echo "  Tailscale IP:           $JETSON_IP"
echo "  Status:                 Online (check with: tailscale status)"
echo ""

echo -e "${CYAN}Next Steps:${NC}"
echo "  1. Verify heartbeat: ssh $JETSON_IP"
echo "  2. Check logs: tailscale status"
echo "  3. View agent status on dashboard"
echo ""

log_pass "✓ Jetson is ready for work"
echo ""