#!/usr/bin/env bash
# ==============================================================================
# NaxOS Automated & Interactive Installer Wizard
# Guides user through disk partitioning, ZFS pool creation or safe foreign pool
# import (TrueNAS / nixos-lukas), network setup, and initial credentials.
# ==============================================================================

set -euo pipefail

RED='\033[0;31m'
GREEN='\033[0;32m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m'

clear
cat << "EOF"
  _   _             ___  ____  
 | \ | | __ ___  __/ _ \/ ___| 
 |  \| |/ _` \ \/ / | | \___ \ 
 | |\  | (_| |>  <| |_| |___) |
 |_| \_|\__,_/_/\_\\___/|____/ 
 Declarative NixOS & ZFS Storage Appliance
================================================
EOF

echo -e "${BOLD}${CYAN}Welcome to the NaxOS Installation Wizard!${NC}\n"

# 1. Hardware Check
echo -e "${BOLD}[1/6] Scanning Hardware & Storage Devices...${NC}"
TOTAL_RAM_MB=$(free -m | awk '/^Mem:/{print $2}')
CPU_CORES=$(nproc)
echo "  Detected CPU Cores: $CPU_CORES"
echo "  Detected RAM: ${TOTAL_RAM_MB} MB"

echo -e "\nAvailable Physical Disks:"
lsblk -d -p -n -o NAME,SIZE,MODEL,TRAN | grep -v 'loop' || true

# 2. Storage Pool Setup (New vs Import)
echo -e "\n${BOLD}[2/6] Storage Architecture Configuration${NC}"
echo "1) Create a brand new ZFS storage pool (Formatted & optimized for NAS)"
echo "2) Import an EXISTING ZFS pool (Preserve existing data e.g. TrueNAS or nixos-lukas 'tank')"
echo "3) Skip storage configuration (Configure post-boot via Web Dashboard)"
read -rp "Select option [1-3] (default: 2): " STORAGE_OPTION
STORAGE_OPTION=${STORAGE_OPTION:-2}

POOL_NAME="tank"
IMPORTED_POOL=""
CREATE_POOL=false

if [ "$STORAGE_OPTION" = "1" ]; then
  CREATE_POOL=true
  read -rp "Enter new pool name (default: tank): " POOL_INPUT
  POOL_NAME=${POOL_INPUT:-tank}
  
  echo "Select pool layout: [1] Mirror, [2] RAID-Z1, [3] RAID-Z2, [4] Single/Stripe"
  read -rp "Layout (default: 1): " LAYOUT_INPUT
  LAYOUT_INPUT=${LAYOUT_INPUT:-1}

  read -rp "Enter disk devices space-separated (e.g. /dev/sdb /dev/sdc): " DISK_DEVICES
  
elif [ "$STORAGE_OPTION" = "2" ]; then
  echo -e "\n${BLUE}Scanning for unimported ZFS pools on connected drives...${NC}"
  zpool import || true
  
  read -rp "Enter the name of the existing pool to import (e.g. tank): " IMPORTED_POOL
  if [ -n "$IMPORTED_POOL" ]; then
    echo -e "${GREEN}Verifying pool '$IMPORTED_POOL' status safely without mounting...${NC}"
    zpool import -N -f "$IMPORTED_POOL" || echo "Warning: Pool could not be imported immediately. Will configure for boot-time import."
  fi
fi

# 3. Target System OS Disk Selection
echo -e "\n${BOLD}[3/6] Target OS Disk Selection${NC}"
echo "Enter the disk where NaxOS system files should be installed (WARNING: this disk will be formatted!):"
read -rp "OS Disk (e.g. /dev/sda or /dev/nvme0n1): " OS_DISK

if [ -z "$OS_DISK" ] || [ ! -b "$OS_DISK" ]; then
  echo -e "${RED}Error: Disk '$OS_DISK' does not exist! Aborting.${NC}"
  exit 1
fi

echo -e "${RED}${BOLD}WARNING: All data on $OS_DISK will be destroyed!${NC}"
read -rp "Are you sure you want to proceed with $OS_DISK? (yes/no): " CONFIRM
if [ "$CONFIRM" != "yes" ]; then
  echo "Installation canceled by user."
  exit 0
fi

# 4. Network Configuration
echo -e "\n${BOLD}[4/6] Network Configuration${NC}"
read -rp "Appliance Hostname (default: naxos): " HOSTNAME
HOSTNAME=${HOSTNAME:-naxos}

read -rp "Use DHCP for network? (yes/no, default: yes): " USE_DHCP
USE_DHCP=${USE_DHCP:-yes}

# 5. Admin User & Credentials
echo -e "\n${BOLD}[5/6] Administrator Account & Access${NC}"
read -rp "Admin username (default: admin): " ADMIN_USER
ADMIN_USER=${ADMIN_USER:-admin}

read -rp "Paste Admin SSH Public Key (optional): " ADMIN_SSH_KEY

# 6. Partitioning and Installation
echo -e "\n${BOLD}[6/6] Executing Partitioning & Installation...${NC}"

# Partition OS disk: 1GB EFI, remainder for root
wipefs -a "$OS_DISK"
parted -s "$OS_DISK" -- mklabel gpt
parted -s "$OS_DISK" -- mkpart ESP fat32 1MiB 1024MiB
parted -s "$OS_DISK" -- set 1 esp on
parted -s "$OS_DISK" -- mkpart primary ext4 1024MiB 100%

# Detect partition names (handles nvme0n1p1 vs sda1)
if [[ "$OS_DISK" =~ [0-9]$ ]]; then
  EFI_PART="${OS_DISK}p1"
  ROOT_PART="${OS_DISK}p2"
else
  EFI_PART="${OS_DISK}1"
  ROOT_PART="${OS_DISK}2"
fi

mkfs.vfat -F 32 -n NIXOS_BOOT "$EFI_PART"
mkfs.ext4 -F -L NIXOS_ROOT "$ROOT_PART"

# Mount filesystems
mount "$ROOT_PART" /mnt
mkdir -p /mnt/boot
mount "$EFI_PART" /mnt/boot

# Generate hardware configuration
mkdir -p /mnt/etc/nixos
nixos-generate-config --root /mnt

# If creating pool now
if [ "$CREATE_POOL" = true ] && [ -n "${DISK_DEVICES:-}" ]; then
  echo "Creating ZFS pool $POOL_NAME..."
  zpool create -f -o ashift=12 -O compression=lz4 -O acltype=posixacl -O xattr=sa "$POOL_NAME" $DISK_DEVICES
  zfs create "$POOL_NAME/media"
  zfs create "$POOL_NAME/backup"
  zfs create "$POOL_NAME/container"
fi

# Generate unique hostId
HOST_ID=$(head -c4 /dev/urandom | od -A none -t x4 | tr -d ' ')

# Write appliance flake and configuration
cat << NIXCONFIG > /mnt/etc/nixos/configuration.nix
{ config, pkgs, lib, ... }:

{
  imports = [
    ./hardware-configuration.nix
  ];

  boot.loader.systemd-boot.enable = true;
  boot.loader.efi.canTouchEfiVariables = true;
  boot.supportedFilesystems = [ "zfs" ];

  networking.hostName = "$HOSTNAME";
  networking.hostId = "$HOST_ID";
  networking.useDHCP = lib.mkDefault $([ "$USE_DHCP" = "yes" ] && echo "true" || echo "false");

  # Safe ZFS extra pools
  boot.zfs.extraPools = [ "$POOL_NAME" $([ -n "$IMPORTED_POOL" ] && echo "\"$IMPORTED_POOL\"") ];

  # Administrator user
  users.users.$ADMIN_USER = {
    isNormalUser = true;
    extraGroups = [ "wheel" "docker" "video" "render" ];
    openssh.authorizedKeys.keys = [ $([ -n "$ADMIN_SSH_KEY" ] && echo "\"$ADMIN_SSH_KEY\"") ];
  };
  security.sudo.wheelNeedsPassword = false;

  # SSH Access
  services.openssh.enable = true;

  # State version
  system.stateVersion = "24.11";
}
NIXCONFIG

echo -e "\n${GREEN}Beginning NixOS System Installation via nixos-install...${NC}"
nixos-install --no-root-passwd

echo -e "\n${BOLD}${GREEN}==============================================${NC}"
echo -e "${BOLD}${GREEN}  NaxOS Installation Complete!${NC}"
echo -e "${BOLD}${GREEN}  Remove the installation USB/ISO and reboot.${NC}"
echo -e "${BOLD}${GREEN}  Access Dashboard on: http://$HOSTNAME.local${NC}"
echo -e "${BOLD}${GREEN}==============================================${NC}"
EOF