#!/usr/bin/env bash
# sovran-install-headless.sh — Non-interactive remote installer for Sovran_SystemsOS

usage() {
    cat <<'USAGE'
Usage: sovran-install-headless.sh [OPTIONS]

Options:
  --disk /dev/sda              Target OS disk (required)
  --data-disk /dev/sdb         Data disk for Bitcoin (optional)
  --role server|desktop|node   Installation role (default: server)
  --deploy-key "ssh-ed25519 AAAA..."  SSH pubkey for remote access after install
  --headscale-server URL       Headscale login server for post-install Tailnet
  --headscale-key KEY          Headscale pre-auth key for the installed OS

VM notes:
  Desktop test installs still need a 256 GB OS disk (thin-provisioned is OK).
  Node/server installs need a separate 2 TB data disk. UEFI is preferred; UEFI
  VMs avoid NVRAM boot-entry writes, and legacy BIOS VMs use GRUB.
USAGE
}

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────────────
DISK=""
DATA_DISK=""
ROLE="server"
DEPLOY_KEY=""
HEADSCALE_SERVER=""
HEADSCALE_KEY=""
DATA_DISK_HAS_TIMECHAIN=false
BOOT_MODE=""
VIRTUALIZATION=""

FLAKE="/etc/sovran/flake"
LOG="/tmp/sovran-headless-install.log"

BYTES_256GB=$((256 * 1024 * 1024 * 1024))
BYTES_2TB=$((2 * 1000 * 1000 * 1000 * 1000))

# ── Logging ───────────────────────────────────────────────────────────────────
exec > >(tee -a "$LOG") 2>&1

log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*"
}

die() {
    log "ERROR: $*"
    exit 1
}

# ── Argument parsing ─────────────────────────────────────────────────────────
while [[ $# -gt 0 ]]; do
    case "$1" in
        --disk)        DISK="$2";        shift 2 ;;
        --data-disk)   DATA_DISK="$2";   shift 2 ;;
        --role)        ROLE="$2";        shift 2 ;;
        --deploy-key)  DEPLOY_KEY="$2";  shift 2 ;;
        --headscale-server) HEADSCALE_SERVER="$2"; shift 2 ;;
        --headscale-key)    HEADSCALE_KEY="$2";    shift 2 ;;
        -h|--help)
            usage
            exit 0
            ;;
        *) die "Unknown argument: $1" ;;
    esac
done

# ── Validate required arguments ───────────────────────────────────────────────
[[ -n "$DISK" ]] || die "--disk is required"

case "$ROLE" in
    server|desktop|node) ;;
    *) die "--role must be one of: server, desktop, node" ;;
esac

# ── Detect firmware / VM environment ─────────────────────────────────────────
if [[ -d /sys/firmware/efi ]]; then
    BOOT_MODE="uefi"
else
    BOOT_MODE="bios"
fi

if command -v systemd-detect-virt >/dev/null 2>&1; then
    if VIRT_RESULT=$(systemd-detect-virt --vm 2>/dev/null); then
        [[ "$VIRT_RESULT" != "none" ]] && VIRTUALIZATION="$VIRT_RESULT"
    fi
fi

log "Boot mode: ${BOOT_MODE} ($([[ "$BOOT_MODE" == uefi ]] && echo systemd-boot || echo GRUB))"
[[ -n "$VIRTUALIZATION" ]] && log "Virtual machine detected: ${VIRTUALIZATION}"

# ── Validate disk existence and size ─────────────────────────────────────────
log "=== Validating disks ==="

[[ -b "$DISK" ]] || die "OS disk not found: $DISK"

disk_size_bytes() {
    local dev="$1"
    lsblk -b -dno SIZE "$dev" 2>/dev/null || echo 0
}

OS_SIZE=$(disk_size_bytes "$DISK")
log "OS disk $DISK: $OS_SIZE bytes"
[[ "$OS_SIZE" -ge "$BYTES_256GB" ]] \
    || die "OS disk $DISK is too small ($(( OS_SIZE / 1024 / 1024 / 1024 )) GB). Minimum is 256 GB."

if [[ -n "$DATA_DISK" ]]; then
    [[ -b "$DATA_DISK" ]] || die "Data disk not found: $DATA_DISK"
    [[ "$DATA_DISK" != "$DISK" ]] || die "OS disk and data disk cannot be the same device"
    DATA_SIZE=$(disk_size_bytes "$DATA_DISK")
    log "Data disk $DATA_DISK: $DATA_SIZE bytes"
    [[ "$DATA_SIZE" -ge "$BYTES_2TB" ]] \
        || die "Data disk $DATA_DISK is too small ($(( DATA_SIZE / 1024 / 1024 / 1024 )) GB). Minimum is 2 TB."
fi

# ── Helper: partition suffix ──────────────────────────────────────────────────
part_suffix() {
    local dev="$1" n="$2"
    local base
    base=$(basename "$dev")
    if [[ "$base" == nvme* || "$base" == mmcblk* || "$base" == loop* || "$base" == md* || "$base" =~ [0-9]$ ]]; then
        echo "${dev}p${n}"
    else
        echo "${dev}${n}"
    fi
}

nix_string() {
    local value="$1"
    value="${value//\\/\\\\}"
    value="${value//\"/\\\"}"
    value="${value//$'\n'/\\n}"
    value="${value//$'\r'/\\r}"
    value="${value//$'\t'/\\t}"
    value="${value//\$\{/\\\$\{}"
    printf '"%s"' "$value"
}

# ── Detect existing Bitcoin timechain data on data disk ───────────────────────
if [[ -n "$DATA_DISK" ]]; then
    DATA_P1=$(part_suffix "$DATA_DISK" 1)
    if [[ -b "$DATA_P1" ]]; then
        DATA_LABEL=$(lsblk -no LABEL "$DATA_P1" 2>/dev/null | head -n1 || true)
        if [[ -z "$DATA_LABEL" ]]; then
            DATA_LABEL=$(blkid -o value -s LABEL "$DATA_P1" 2>/dev/null || true)
        fi

        if [[ "$DATA_LABEL" == "BTCEcoandBackup" ]]; then
            CHECK_MOUNT=$(mktemp -d /tmp/sovran-data-check.XXXXXX)
            if mount -o ro "$DATA_P1" "$CHECK_MOUNT" 2>/dev/null; then
                if [[ -d "$CHECK_MOUNT/BTCEcoandBackup/Bitcoin_Node" && -d "$CHECK_MOUNT/BTCEcoandBackup/Electrs_Data" ]]; then
                    DATA_DISK_HAS_TIMECHAIN=true
                    log "Existing Bitcoin timechain detected on data drive — preserving data"
                fi
                umount "$CHECK_MOUNT" || true
            fi
            rmdir "$CHECK_MOUNT" 2>/dev/null || true
        fi
    fi
fi

# ── Step 1: Wipe disks ────────────────────────────────────────────────────────
log "=== Wiping disk(s) ==="

sgdisk --zap-all "$DISK"
wipefs --all --force "$DISK"

if [[ -n "$DATA_DISK" && "$DATA_DISK_HAS_TIMECHAIN" != true ]]; then
    sgdisk --zap-all "$DATA_DISK"
    wipefs --all --force "$DATA_DISK"
fi

partprobe "$DISK"
[[ -n "$DATA_DISK" && "$DATA_DISK_HAS_TIMECHAIN" != true ]] && partprobe "$DATA_DISK"
sleep 2

# ── Step 2: Partition OS disk ─────────────────────────────────────────────────
if [[ "$BOOT_MODE" == "uefi" ]]; then
    log "=== Partitioning OS disk (UEFI) ==="
    sgdisk \
        -n "1:1M:+512M" -t "1:EF00" -c "1:ESP" \
        -n "2:0:0"      -t "2:8300" -c "2:root" \
        "$DISK"
else
    log "=== Partitioning OS disk (Legacy BIOS) ==="
    sgdisk \
        -n "1:1M:+1M" -t "1:EF02" -c "1:BIOS-boot" \
        -n "2:0:0"    -t "2:8300" -c "2:root" \
        "$DISK"
fi

partprobe "$DISK"
sleep 2

# ── Step 3: Partition data disk (if present) ──────────────────────────────────
if [[ -n "$DATA_DISK" && "$DATA_DISK_HAS_TIMECHAIN" != true ]]; then
    log "=== Partitioning data disk ==="
    sgdisk \
        -n "1:1M:0" -t "1:8300" -c "1:primary" \
        "$DATA_DISK"
    partprobe "$DATA_DISK"
    sleep 2
fi

# ── Step 4: Format partitions ─────────────────────────────────────────────────
log "=== Formatting partitions ==="

BOOT_P1=$(part_suffix "$DISK" 1)
BOOT_P2=$(part_suffix "$DISK" 2)

if [[ "$BOOT_MODE" == "uefi" ]]; then
    mkfs.vfat -F 32 "$BOOT_P1"
else
    log "Skipping ESP format for Legacy BIOS install"
fi
mkfs.ext4 -F -L sovran_systemsos "$BOOT_P2"

if [[ -n "$DATA_DISK" && "$DATA_DISK_HAS_TIMECHAIN" != true ]]; then
    DATA_P1=$(part_suffix "$DATA_DISK" 1)
    mkfs.ext4 -F -L BTCEcoandBackup "$DATA_P1"
fi

# ── Step 5: Mount filesystems ─────────────────────────────────────────────────
log "=== Mounting filesystems ==="

mount "$BOOT_P2" /mnt
if [[ "$BOOT_MODE" == "uefi" ]]; then
    mkdir -p /mnt/boot/efi
    mount -o umask=0077,defaults "$BOOT_P1" /mnt/boot/efi
else
    log "Legacy BIOS install: /boot/efi mount not required"
fi

if [[ -n "$DATA_DISK" ]]; then
    DATA_P1=$(part_suffix "$DATA_DISK" 1)
    mkdir -p /mnt/run/media/Second_Drive
    mount "$DATA_P1" /mnt/run/media/Second_Drive

    # ── Step 6: Create Bitcoin data directories ─────────────────────────────
    log "=== Creating Bitcoin data directories ==="
    mkdir -p /mnt/run/media/Second_Drive/BTCEcoandBackup/Bitcoin_Node
    mkdir -p /mnt/run/media/Second_Drive/BTCEcoandBackup/Electrs_Data
    mkdir -p /mnt/run/media/Second_Drive/BTCEcoandBackup/NixOS_Snapshot_Backup
fi

# ── Step 7: Generate hardware config ─────────────────────────────────────────
log "=== Generating hardware config ==="
nixos-generate-config --root /mnt

# ── Step 8: Copy flake source ─────────────────────────────────────────────────
log "=== Copying flake to /mnt ==="
cp /mnt/etc/nixos/hardware-configuration.nix /tmp/hardware-configuration.nix
rm -rf /mnt/etc/nixos/
mkdir -p /mnt/etc/nixos
cp -a "${FLAKE}/." /mnt/etc/nixos/
cp /tmp/hardware-configuration.nix /mnt/etc/nixos/hardware-configuration.nix

# ── Step 9: Write role-state.nix ─────────────────────────────────────────────
log "=== Writing role config ==="

case "$ROLE" in
    server)
        IS_SERVER=true; IS_DESKTOP=false; IS_NODE=false ;;
    desktop)
        IS_SERVER=false; IS_DESKTOP=true; IS_NODE=false ;;
    node)
        IS_SERVER=false; IS_DESKTOP=false; IS_NODE=true ;;
esac

{
    cat <<EOF
# THIS FILE IS AUTO-GENERATED BY THE INSTALLER. DO NOT EDIT.
{ config, lib, ... }:
{
  sovran_systemsOS.roles.server_plus_desktop = lib.mkDefault ${IS_SERVER};
  sovran_systemsOS.roles.desktop = lib.mkDefault ${IS_DESKTOP};
  sovran_systemsOS.roles.node = lib.mkDefault ${IS_NODE};
EOF

    if [[ "$BOOT_MODE" == "bios" ]]; then
        GRUB_DEVICE=$(nix_string "$DISK")
        cat <<EOF

  # The installer ISO was booted in Legacy BIOS mode, so install GRUB to the
  # target disk instead of using the default UEFI systemd-boot configuration.
  boot.loader.systemd-boot.enable = lib.mkForce false;
  boot.loader.efi.canTouchEfiVariables = lib.mkForce false;
  boot.loader.grub.enable = lib.mkForce true;
  boot.loader.grub.device = lib.mkForce ${GRUB_DEVICE};
  fileSystems."/boot/efi".enable = lib.mkForce false;
EOF
    elif [[ -n "$VIRTUALIZATION" ]]; then
        cat <<EOF

  # VM UEFI firmware can reject or forget NVRAM boot-entry writes. Keep the
  # normal UEFI systemd-boot layout, but use the fallback removable path instead
  # of depending on EFI variable updates.
  boot.loader.efi.canTouchEfiVariables = lib.mkForce false;
EOF
    fi

    echo "}"
} > /mnt/etc/nixos/role-state.nix

# ── Step 10: Write custom.nix with deploy config ──────────────────────────────
log "=== Writing custom.nix ==="

if [[ -n "$DEPLOY_KEY" || -n "$HEADSCALE_SERVER" ]]; then
    {
        echo '{ config, lib, ... }:'
        echo '{'
        echo '  sovran_systemsOS.deploy = {'
        echo '    enable = true;'
        [[ -n "$DEPLOY_KEY" ]]        && echo "    authorizedKey = \"${DEPLOY_KEY}\";"
        [[ -n "$HEADSCALE_SERVER" ]]  && echo "    headscaleServer = \"${HEADSCALE_SERVER}\";"
        echo '  };'
        echo '}'
    } > /mnt/etc/nixos/custom.nix
else
    cp /mnt/etc/nixos/custom.template.nix /mnt/etc/nixos/custom.nix
    chmod 644 /mnt/etc/nixos/custom.nix
fi

# ── Write Headscale auth key if provided ─────────────────────────────────────
if [[ -n "$HEADSCALE_KEY" ]]; then
    mkdir -p /mnt/var/lib/secrets
    echo "$HEADSCALE_KEY" > /mnt/var/lib/secrets/headscale-authkey
    chmod 600 /mnt/var/lib/secrets/headscale-authkey
    log "Headscale auth key written to /mnt/var/lib/secrets/headscale-authkey"
fi

# ── Step 11: Copy configs to host for flake evaluation ───────────────────────
log "=== Copying config files to host /etc/nixos for flake evaluation ==="
mkdir -p /etc/nixos
cp /mnt/etc/nixos/role-state.nix /etc/nixos/role-state.nix
cp /mnt/etc/nixos/custom.nix /etc/nixos/custom.nix
cp /mnt/etc/nixos/hardware-configuration.nix /etc/nixos/hardware-configuration.nix

# ── Step 12: Run nixos-install ────────────────────────────────────────────────
log "=== Running nixos-install ==="
nixos-install \
    --root /mnt \
    --flake /mnt/etc/nixos#nixos \
    --no-root-password \
    --impure

log "=== Installation complete! ==="
log "You can now reboot into Sovran_SystemsOS."
log "After reboot, the machine will be accessible via SSH on port 22 (if --deploy-key was provided)."
[[ -n "$HEADSCALE_SERVER" ]] && \
    log "Tailscale will connect to Headscale at ${HEADSCALE_SERVER} on first boot."