#!/usr/bin/env bash
# NixOS Deployment Helper Script
# Remote NixOS deployment from Hermes container to target hosts.
#
# Usage: ./deploy.sh <hostname> [branch] [action]
#
# Actions:
#   switch       Activate configuration now (default)
#   boot         Activate on next reboot
#   test         Activate without switching generations
#   build        Build locally only, no remote activation
#   dry-activate Show what would change without applying
#
# Examples:
#   ./deploy.sh lazyworkhorse                  # deploy master/switch to lazyworkhorse
#   ./deploy.sh cyt-pi feat/test boot          # deploy feat/test branch, activate on boot
#   ./deploy.sh uConsole master build          # just build, don't deploy
#   NO_BUILD_CHECK=1 ./deploy.sh uConsole      # skip the pre-flight nix build
#
# Environment variables:
#   SSH_USER      SSH user (default: auto-detected per host)
#   SSH_PORT      SSH port (default: auto-detected per host)
#   SSH_KEY       SSH identity file
#   BUILD_HOST    Build flake for this host (default: same as target host)
#   NO_BUILD_CHECK  Set to 1 to skip local nix build before deployment

set -euo pipefail

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

info()  { echo -e "${BLUE}[INFO]${NC}  $*"; }
ok()    { echo -e "${GREEN}[OK]${NC}    $*"; }
warn()  { echo -e "${YELLOW}[WARN]${NC}  $*"; }
error() { echo -e "${RED}[ERROR]${NC} $*" >&2; }
step()  { echo -e "\n${CYAN}━━━ $* ━━━${NC}"; }

# ── Cleanup trap ───────────────────────────────────────────────────────
cleanup() {
    local ec=$?
    if [ $ec -ne 0 ]; then
        error "Deployment failed with exit code $ec"
    fi
    exit $ec
}
trap cleanup EXIT

# ── Usage / Help ───────────────────────────────────────────────────────
show_usage() {
    cat <<EOF
Usage: $0 <hostname> [branch] [action]

Remote NixOS deployment from Hermes container to target hosts.

HOSTNAME (required):
    lazyworkhorse     x86_64 main server
    cyt-pi            aarch64 Pi Zero 2 W (via reverse tunnel)
    uConsole          aarch64 ClockworkPi

BRANCH (optional, default: master):
    Git branch or tag to deploy. Fetched from origin.

ACTION (optional, default: switch):
    switch           Activate configuration now (default)
    boot             Activate on next reboot
    test             Activate without switching generations
    build            Build locally only, skip remote deployment
    dry-activate     Show what would change without applying

Environment variables:
    SSH_USER       SSH username override
    SSH_PORT       SSH port override
    SSH_KEY        SSH identity file path
    BUILD_HOST     Build flake hostname (default: same as HOSTNAME)
    NO_BUILD_CHECK Skip local nix build validation (set to 1)

Examples:
    $0 lazyworkhorse                       # deploy master/switch
    $0 cyt-pi feat/test boot               # deploy feature branch, boot
    $0 uConsole master build               # just build, no remote
    NO_BUILD_CHECK=1 $0 uConsole           # skip build check

EOF
    exit 0
}

# ── Argument parsing ───────────────────────────────────────────────────
HOSTNAME="${1:-}"
BRANCH="${2:-master}"
ACTION="${3:-switch}"
NO_BUILD_CHECK="${NO_BUILD_CHECK:-0}"

if [ "$HOSTNAME" = "--help" ] || [ "$HOSTNAME" = "-h" ] || [ -z "$HOSTNAME" ]; then
    show_usage
fi

# ── Host configuration ─────────────────────────────────────────────────
case "$HOSTNAME" in
    lazyworkhorse)
        DEFAULT_SSH_USER="ai-worker"
        DEFAULT_SSH_PORT="2424"
        ARCH="x86_64-linux"
        ;;
    cyt-pi)
        DEFAULT_SSH_USER="gortium"
        DEFAULT_SSH_PORT="19999"
        ARCH="aarch64-linux"
        ;;
    uConsole)
        DEFAULT_SSH_USER="gortium"
        DEFAULT_SSH_PORT="22"
        ARCH="aarch64-linux"
        ;;
    *)
        error "Unknown host: $HOSTNAME"
        echo "Supported hosts: lazyworkhorse, cyt-pi, uConsole"
        exit 1
        ;;
esac

SSH_USER="${SSH_USER:-$DEFAULT_SSH_USER}"
SSH_PORT="${SSH_PORT:-$DEFAULT_SSH_PORT}"
SSH_KEY="${SSH_KEY:-/opt/data/home/.ssh/id_hermes_gitea}"
BUILD_HOST="${BUILD_HOST:-$HOSTNAME}"

SSH_OPTS="-p $SSH_PORT -i $SSH_KEY -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null"
SSH_TARGET="${SSH_USER}@${HOSTNAME}"
export GIT_SSH_COMMAND="ssh -i $SSH_KEY -p 2222 -o StrictHostKeyChecking=no"
export PATH="/nix/var/nix/profiles/default/bin:$PATH"

# ── Banner ─────────────────────────────────────────────────────────────
echo "╔══════════════════════════════════════════════╗"
echo "║         NixOS Remote Deployment              ║"
echo "╚══════════════════════════════════════════════╝"
info "Host:   $HOSTNAME ($ARCH)"
info "Branch: $BRANCH"
info "Action: $ACTION"
info "SSH:    ${SSH_USER}@${HOSTNAME}:${SSH_PORT}"
echo ""

# ── Pre-flight checks ─────────────────────────────────────────────────
step "Pre-flight checks"

# 1. Check required tools
for cmd in nix git ssh; do
    if ! command -v "$cmd" &>/dev/null; then
        error "Required tool not found: $cmd"
        exit 1
    fi
done
ok "Required tools available (nix, git, ssh)"

# 2. Check infra repo
INFRA_DIR="$(cd "$(dirname "$0")/.." && pwd)"
if [ ! -d "$INFRA_DIR/.git" ]; then
    error "Not a git repository: $INFRA_DIR"
    exit 1
fi
ok "Infra repo found at $INFRA_DIR"

# 3. Check SSH connectivity (skip for build-only actions)
if [ "$ACTION" != "build" ]; then
    if ssh $SSH_OPTS -o ConnectTimeout=5 "$SSH_TARGET" "echo connected" &>/dev/null; then
        ok "SSH connectivity to $HOSTNAME verified"
    else
        warn "Cannot reach $HOSTNAME via SSH — deployment step will fail later"
    fi
fi

# ── Git sync ───────────────────────────────────────────────────────────
step "Git sync"

cd "$INFRA_DIR"

# Stash local changes if any
if ! git diff --quiet HEAD; then
    warn "Local changes detected, stashing..."
    git stash push -m "auto-stash before deploy $(date -Iseconds)"
    STASHED=1
else
    STASHED=0
fi

# Fetch and checkout
git fetch origin "$BRANCH" 2>/dev/null || git fetch origin master
if git rev-parse --verify "origin/$BRANCH" &>/dev/null 2>&1; then
    # Remote branch exists — fast-forward merge
    git checkout -B "$BRANCH" "origin/$BRANCH"
elif git rev-parse --verify "$BRANCH" &>/dev/null 2>&1; then
    # Local branch or tag
    git checkout "$BRANCH"
else
    error "Branch/tag not found: $BRANCH"
    exit 1
fi
ok "Checked out $BRANCH ($(git rev-parse --short HEAD))"

# Update submodules
if [ -f .gitmodules ]; then
    git submodule update --init --recursive
    ok "Submodules updated"
fi

# ── Build validation ──────────────────────────────────────────────────
if [ "$NO_BUILD_CHECK" != "1" ]; then
    step "Build validation"
    info "Building nixosConfigurations.$BUILD_HOST (no link)..."

    if nix build --no-link --print-build-logs \
        ".#nixosConfigurations.${BUILD_HOST}.config.system.build.toplevel" 2>&1; then
        ok "Build succeeded for $BUILD_HOST"
    else
        error "Build failed for $BUILD_HOST"
        exit 1
    fi
else
    warn "Build check skipped (NO_BUILD_CHECK=1)"
fi

# ── Deployment ─────────────────────────────────────────────────────────
if [ "$ACTION" = "build" ]; then
    step "Build complete (no deployment)"
    info "Use one of: switch, boot, test, dry-activate to deploy"
    exit 0
fi

step "Deployment ($ACTION)"

# Build the nixos-rebuild command
case "$ACTION" in
    switch|boot|test)
        nixos-rebuild "$ACTION" \
            --flake ".#$HOSTNAME" \
            --target-host "$SSH_TARGET" \
            --build-host "localhost" \
            --use-remote-sudo \
            --max-jobs 4
        ;;
    dry-activate)
        nixos-rebuild dry-activate \
            --flake ".#$HOSTNAME" \
            --target-host "$SSH_TARGET" \
            --build-host "localhost" \
            --use-remote-sudo
        ;;
    *)
        error "Unknown action: $ACTION"
        echo "Valid actions: switch, boot, test, build, dry-activate"
        exit 1
        ;;
esac

# ── Check result ───────────────────────────────────────────────────────
DEPLOY_EXIT=$?
if [ $DEPLOY_EXIT -eq 0 ]; then
    echo ""
    ok "Deployment to $HOSTNAME ($ACTION) completed successfully"
    case "$ACTION" in
        switch|test)
            info "Configuration is now active"
            ;;
        boot)
            info "Configuration will activate on next reboot"
            ;;
        dry-activate)
            info "Dry-run complete — no changes applied"
            ;;
    esac
else
    error "Deployment failed with exit code $DEPLOY_EXIT"
    exit $DEPLOY_EXIT
fi

echo ""
echo "╔══════════════════════════════════════════════╗"
echo "║         Deployment Complete                   ║"
echo "╚══════════════════════════════════════════════╝"
info "Host:   $HOSTNAME"
info "Branch: $BRANCH ($(git rev-parse --short HEAD))"
info "Action: $ACTION"
info "Time:   $(date -Iseconds)"