setup

#!/usr/bin/env bash

set -e

# Master identity pubkeys for each group
declare -A group_pubkeys

# TPM-protected identity paths (can be overridden in private/setup-env)
AGE_TPM_IDENTITY="${AGE_TPM_IDENTITY:-${HOME}/.age/tpm-identity.txt}"
AGE_MASTER_KEY_TPM="${AGE_MASTER_KEY_TPM:-${HOME}/.age/master-key.age}"
# Runtime decryption path - must match defaults in age.nix module
if [[ "$(uname -s)" == "Darwin" ]]; then
    AGE_RUNTIME_KEY="/tmp/age-master-key"
else
    AGE_RUNTIME_KEY="/dev/shm/age-master-key"
fi
AGE_DECRYPTED_KEY=""  # Set during rekey

THISHOST="$(hostname -s)"
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
ROOTS=/nix/var/nix/gcroots/my-builds

# Source private environment if it exists
[[ -f "$SCRIPT_DIR/private/setup-env" ]] && source "$SCRIPT_DIR/private/setup-env"

# Options
HOSTS=()
DO_REFRESH=false
DO_SWITCH=false
DO_REKEY=false
DO_GIT=true
DO_PUSH=true
DO_CACHE_PUSH=true
DO_LIST=false
DO_ALLOW_SHORT_HOST=false
DO_OFFLINE=false
MESSAGE=""
VERBOSE=false

# Cache for host metadata
declare -A HOST_META

usage() {
    cat <<EOF
Usage: $(basename "$0") [OPTIONS] [HOST...]

Build and deploy NixOS/nix-darwin configurations.

If no HOST is specified, builds the current host (${THISHOST}).

Options:
  -s, --switch      Switch to the new configuration after building
  -k, --rekey       Rekey age secrets for the specified hosts before building
  -r, --refresh     Update flake inputs before building
  -n, --no-git      Skip git commit and push
  -np, --no-push    Commit locally but skip git push
  -nc, --no-cache   Skip pushing closures to attic cache (still inhibits per-build hooks)
  -m, --message=MSG Custom git commit message (default: "autosave (<host> / <os>)")
  -l, --list        List known hosts and exit
  --allow-short-host
                    Allow using short host name for remote switch when FQDN is missing
  --offline         Pass --offline to nix build (no network fetches)
  -v, --verbose     Enable verbose output (set -x)
  -h, --help        Show this help message

Host specifiers:
  HOST                 Individual host name (e.g., o1, pavel-fw)
  +GROUP               All hosts in a group (e.g., +infra, +pavel)

Examples:
  $(basename "$0")                  # Build current host
  $(basename "$0") -s               # Build and switch current host
  $(basename "$0") -k -s            # Rekey, build, and switch current host
  $(basename "$0") o1 o2            # Build o1 and o2
  $(basename "$0") +infra           # Build all hosts in 'infra' group
  $(basename "$0") -k +infra        # Rekey and build all infra hosts
  $(basename "$0") -r -s            # Update flake, build, and switch
  $(basename "$0") pavel-fw -s --allow-short-host
                              # Remote switch via short host name when no FQDN is defined
  $(basename "$0") -n               # Build without git operations
EOF
    exit 0
}

load_host_meta() {
    echo "Loading host metadata from flake..."
    add_all
    local json
    local errfile
    errfile="$(mktemp)"
    if ! json=$(nix eval ".?submodules=1#hostMeta" --json 2>"$errfile"); then
        echo "Error: nix eval failed for host metadata:" >&2
        cat "$errfile" >&2
        rm -f "$errfile"
        exit 1
    fi
    rm -f "$errfile"

    # Parse JSON and populate HOST_META
    while IFS="=" read -r key value; do
        HOST_META["$key"]="$value"
    done < <(echo "$json" | jq -r 'to_entries[] | "\(.key)=\(.value | @json)"')
}

get_host_attr() {
    local host="$1"
    local attr="$2"
    local meta="${HOST_META[$host]}"
    if [[ -z "$meta" ]]; then
        echo ""
        return
    fi
    echo "$meta" | jq -r ".$attr // empty"
}

get_hosts_by_group() {
    local target_group="$1"
    local hosts=()
    for host in "${!HOST_META[@]}"; do
        local group
        group=$(get_host_attr "$host" "group")
        if [[ "$group" == "$target_group" ]]; then
            hosts+=("$host")
        fi
    done
    echo "${hosts[@]}"
}

get_current_platform() {
    case "$(uname -s)" in
        Darwin) echo "darwin" ;;
        Linux)  echo "linux" ;;
        *)      echo "unknown" ;;
    esac
}

expand_host_specifiers() {
    local current_platform
    current_platform=$(get_current_platform)
    local expanded=()

    for spec in "${HOSTS[@]}"; do
        if [[ "$spec" == +* ]]; then
            local group="${spec:1}"
            local group_hosts=()
            local skipped=()

            for host in "${!HOST_META[@]}"; do
                local host_group host_platform
                host_group=$(get_host_attr "$host" "group")
                host_platform=$(get_host_attr "$host" "platform")

                if [[ "$host_group" == "$group" ]]; then
                    if [[ "$host_platform" == "$current_platform" ]]; then
                        group_hosts+=("$host")
                    else
                        skipped+=("$host")
                    fi
                fi
            done

            if [[ ${#group_hosts[@]} -eq 0 ]]; then
                echo "Error: No buildable hosts found in group '$group' for platform '$current_platform'" >&2
                [[ ${#skipped[@]} -gt 0 ]] && echo "    (skipped incompatible: ${skipped[*]})" >&2
                echo "Use --list to see known hosts and their groups." >&2
                exit 1
            fi

            echo "    Expanding +$group: ${group_hosts[*]}"
            [[ ${#skipped[@]} -gt 0 ]] && echo "    (skipped incompatible: ${skipped[*]})"

            for h in "${group_hosts[@]}"; do
                expanded+=("$h")
            done
        else
            expanded+=("$spec")
        fi
    done
    HOSTS=("${expanded[@]}")
}

list_hosts() {
    load_host_meta
    echo "Known hosts:"
    printf "  %-16s %-8s %-8s %s\n" "HOST" "PLATFORM" "GROUP" "FQN"
    printf "  %-16s %-8s %-8s %s\n" "----" "--------" "-----" "---"
    for host in "${!HOST_META[@]}"; do
        local platform group fqn
        platform=$(get_host_attr "$host" "platform")
        group=$(get_host_attr "$host" "group")
        fqn=$(get_host_attr "$host" "fqn")
        [[ "$fqn" == "null" || -z "$fqn" ]] && fqn="(local)"
        printf "  %-16s %-8s %-8s %s\n" "$host" "$platform" "$group" "$fqn"
    done | sort
    exit 0
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            -s|--switch)
                DO_SWITCH=true
                shift
                ;;
            -k|--rekey)
                DO_REKEY=true
                shift
                ;;
            -r|--refresh)
                DO_REFRESH=true
                shift
                ;;
            -n|--no-git)
                DO_GIT=false
                shift
                ;;
            -np|--no-push)
                DO_PUSH=false
                shift
                ;;
            -nc|--no-cache)
                DO_CACHE_PUSH=false
                shift
                ;;
            -m|--message)
                MESSAGE="$2"
                shift 2
                ;;
            --message=*)
                MESSAGE="${1#*=}"
                shift
                ;;
            -m=*)
                MESSAGE="${1#*=}"
                shift
                ;;
            -l|--list)
                DO_LIST=true
                shift
                ;;
            --allow-short-host)
                DO_ALLOW_SHORT_HOST=true
                shift
                ;;
            --offline)
                DO_OFFLINE=true
                shift
                ;;
            -v|--verbose)
                VERBOSE=true
                shift
                ;;
            -h|--help)
                usage
                ;;
            -*)
                echo "Error: Unknown option: $1" >&2
                echo "Use --help for usage information." >&2
                exit 1
                ;;
            *)
                HOSTS+=("$1")
                shift
                ;;
        esac
    done

    # Default to current host if none specified
    if [[ ${#HOSTS[@]} -eq 0 ]]; then
        HOSTS=("$THISHOST")
    fi

    # Default commit message
    if [[ -z "$MESSAGE" ]]; then
        MESSAGE="autosave ($(hostname) / $(uname))"
    fi
}

validate_hosts() {
    for host in "${HOSTS[@]}"; do
        if [[ -z "${HOST_META[$host]}" ]]; then
            echo "Error: Unknown host '$host'" >&2
            echo "Use --list to see known hosts." >&2
            exit 1
        fi
    done
}

add_all() {
    git add .
    git submodule foreach git add .
}

commit_and_push() {
    if [[ "$DO_GIT" == "true" ]]; then
        git submodule foreach "git diff-index --quiet HEAD || (git add -A && git commit -m \"${MESSAGE}\")"
        git diff-index --quiet HEAD || (git add -A && git commit -m "${MESSAGE}")
        if [[ "$DO_PUSH" == "true" ]]; then
            git push --recurse-submodules=on-demand
        fi
    fi
}

do_refresh() {
    if [[ "$DO_REFRESH" == "true" ]]; then
        echo "==> Updating flake inputs..."
        nix flake update
    fi
}

decrypt_password_key() {
    local password_key="${HOME}/age-key.age"
    if [[ ! -f "$password_key" ]]; then
        return 1
    fi

    echo "    Decrypting master key (enter password)..." >&2
    if age -d -o "$AGE_DECRYPTED_KEY" "$password_key"; then
        chmod 600 "$AGE_DECRYPTED_KEY"
        echo "$AGE_DECRYPTED_KEY"
        return 0
    else
        echo "Error: Failed to decrypt master key" >&2
        rm -f "$AGE_DECRYPTED_KEY"
        return 1
    fi
}

decrypt_tpm_key() {
    if [[ ! -f "$AGE_MASTER_KEY_TPM" || ! -f "$AGE_TPM_IDENTITY" ]]; then
        return 1
    fi

    if ! command -v age-plugin-tpm &>/dev/null; then
        echo "    Warning: age-plugin-tpm not found, skipping TPM" >&2
        return 1
    fi

    echo "    Decrypting master key via TPM..." >&2
    if age -d -i "$AGE_TPM_IDENTITY" -o "$AGE_DECRYPTED_KEY" "$AGE_MASTER_KEY_TPM" 2>/dev/null; then
        chmod 600 "$AGE_DECRYPTED_KEY"
        echo "$AGE_DECRYPTED_KEY"
        return 0
    else
        echo "    Warning: TPM decryption failed, trying password key..." >&2
        rm -f "$AGE_DECRYPTED_KEY"
        return 1
    fi
}

get_master_identity() {
    # Returns path to the master identity file
    # Tries TPM first (Linux only), then falls back to password-protected key

    AGE_DECRYPTED_KEY="$AGE_RUNTIME_KEY"
    trap 'rm -f "$AGE_DECRYPTED_KEY"' EXIT

    # Linux: try TPM first
    if [[ "$(uname -s)" != "Darwin" ]]; then
        if decrypt_tpm_key; then
            return 0
        fi
    fi

    # Fallback (all platforms): password-protected key
    if decrypt_password_key; then
        return 0
    fi

    echo "Error: No usable master key found" >&2
    echo "Either run ./enroll-tpm for TPM protection," >&2
    echo "or create a password-protected key: age -p -o ~/age-key.age ~/age-key.txt" >&2
    return 1
}

cleanup_decrypted_key() {
    if [[ -n "$AGE_DECRYPTED_KEY" && -f "$AGE_DECRYPTED_KEY" ]]; then
        rm -f "$AGE_DECRYPTED_KEY"
        AGE_DECRYPTED_KEY=""
    fi
}

do_rekey() {
    if [[ "$DO_REKEY" == "true" ]]; then
        echo "==> Rekeying secrets for: ${HOSTS[*]}"
        add_all

        # Decrypt master key to the path expected by Nix config
        get_master_identity >/dev/null || exit 1

        local system
        system="$(uname -m)-$(uname -s | tr '[:upper:]' '[:lower:]')"

        for host in "${HOSTS[@]}"; do
            local group
            group=$(get_host_attr "$host" "group")

            local pubkey="${group_pubkeys[$group]}"
            if [[ -z "$pubkey" ]]; then
                echo "Warning: No pubkey configured for group '$group', skipping rekey for $host"
                continue
            fi

            echo "    Rekeying $host (group: $group)..."
            nix run ".?submodules=1#agenix-rekey-hosts.${host}.${system}.rekey"
        done

        cleanup_decrypted_key
    fi
}

attic_inhibit_start() {
    ATTIC_INHIBIT="/tmp/attic-push-inhibit"
    touch "$ATTIC_INHIBIT"
    trap 'rm -f "$ATTIC_INHIBIT"' EXIT
}

attic_push_closures() {
    if [[ "$DO_CACHE_PUSH" != "true" ]]; then
        echo "==> Skipping attic cache push (-nc)"
        return
    fi

    if command -v attic-push-closure &>/dev/null; then
        local -a closures=()
        for host in "${HOSTS[@]}"; do
            [[ -L "${ROOTS}/${host}" ]] && closures+=("$(readlink -f "${ROOTS}/${host}")")
        done
        if [[ ${#closures[@]} -gt 0 ]]; then
            echo "==> Pushing ${#closures[@]} closure(s) to attic..."
            attic-push-closure "${closures[@]}" || echo "    Warning: attic push failed (non-fatal)"
        fi
    fi
}

attic_inhibit_stop() {
    rm -f "$ATTIC_INHIBIT"
    trap - EXIT
}

do_build() {
    echo "==> Building: ${HOSTS[*]}"
    add_all
    mkdir -p "$ROOTS"

    local -a build_args=()
    if [[ "$DO_OFFLINE" == "true" ]]; then
        build_args+=(--offline)
    fi

    for host in "${HOSTS[@]}"; do
        local platform
        platform=$(get_host_attr "$host" "platform")

        echo "    Building $host ($platform)..."
        case "$platform" in
            darwin)
                nix build ".?submodules=1#darwinConfigurations.${host}.system" \
                    --out-link "${ROOTS}/${host}" \
                    "${build_args[@]}"
                ;;
            linux)
                nix build ".?submodules=1#nixosConfigurations.${host}.config.system.build.toplevel" \
                    --out-link "${ROOTS}/${host}" \
                    "${build_args[@]}"
                ;;
            *)
                echo "Error: Unknown platform '$platform' for host '$host'" >&2
                exit 1
                ;;
        esac
    done

    commit_and_push
}

do_switch() {
    if [[ "$DO_SWITCH" == "true" ]]; then
        echo "==> Switching: ${HOSTS[*]}"
        add_all

        for host in "${HOSTS[@]}"; do
            local platform fqn closure remote_target
            platform=$(get_host_attr "$host" "platform")
            fqn=$(get_host_attr "$host" "fqn")
            closure=$(readlink -f "${ROOTS}/${host}")

            # null fqn means local deployment
            [[ "$fqn" == "null" || -z "$fqn" ]] && fqn=""

            remote_target="$fqn"
            if [[ "$THISHOST" != "$host" && -z "$remote_target" ]]; then
                if [[ "$DO_ALLOW_SHORT_HOST" == "true" ]]; then
                    remote_target="$host"
                    echo "    Warning: '$host' has no FQDN; using short host name for remote switch"
                fi
            fi

            if [[ "$THISHOST" != "$host" && -z "$remote_target" ]]; then
                echo "Error: Host '$host' has no FQDN (smind.host.fqn/networking.domain), refusing local fallback switch" >&2
                echo "Use --allow-short-host to opt in to short-name remote switching for this run." >&2
                exit 1
            fi

            echo "    Switching $host..."
            case "$platform" in
                darwin)
                    if [[ "$THISHOST" == "$host" ]]; then
                        "${closure}/sw/bin/darwin-rebuild" switch --flake .
                    else
                        nix copy --to "ssh://${remote_target}" "${ROOTS}/${host}"
                        ssh -t "${remote_target}" sudo -H --preserve-env=PATH env nix-env -p /nix/var/nix/profiles/system --set "$closure"
                        ssh -t "${remote_target}" "sudo ${closure}/sw/bin/darwin-rebuild" activate
                    fi
                    ;;
                linux)
                    if [[ "$THISHOST" == "$host" ]]; then
                        sudo nix-env -p /nix/var/nix/profiles/system --set "$closure"
                        sudo /nix/var/nix/profiles/system/bin/switch-to-configuration switch
                    else
                        nix copy --to "ssh://root@${remote_target}" "${ROOTS}/${host}"
                        ssh "root@${remote_target}" nix-env -p /nix/var/nix/profiles/system --set "$closure"
                        ssh "root@${remote_target}" /nix/var/nix/profiles/system/bin/switch-to-configuration switch
                    fi
                    ;;
            esac
        done

        commit_and_push
    fi
}

main() {
    parse_args "$@"

    if [[ "$DO_LIST" == "true" ]]; then
        list_hosts
    fi

    if [[ "$VERBOSE" == "true" ]]; then
        set -x
    fi

    cd "$SCRIPT_DIR"
    load_host_meta
    expand_host_specifiers
    validate_hosts

    do_refresh
    do_rekey
    attic_inhibit_start
    do_build
    attic_push_closures
    do_switch
    attic_inhibit_stop

    echo "==> Done!"
}

main "$@"