#!/bin/bash

# Ubuntu 24.04 Unattended Installation Script for Hetzner Dedicated Servers
# Uses btrfs with RAID 1 over 2 disks
# Based on Hetzner installimage system

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

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

warn() {
    echo -e "${YELLOW}[$(date +'%Y-%m-%d %H:%M:%S')] WARNING: $1${NC}"
}

error() {
    echo -e "${RED}[$(date +'%Y-%m-%d %H:%M:%S')] ERROR: $1${NC}"
    exit 1
}

# Check if running in rescue system
INSTALLIMAGE_PATH=""
if [ -f /root/.oldroot/nfs/install/installimage ]; then
    INSTALLIMAGE_PATH="/root/.oldroot/nfs/install/installimage"
elif [ -f /usr/bin/installimage ]; then
    INSTALLIMAGE_PATH="/usr/bin/installimage"
else
    error "This script must be run from Hetzner rescue system (installimage not found)"
fi

# Additional check for rescue system
if [ ! -f /etc/hostname ] || [ "$(cat /etc/hostname 2>/dev/null)" != "rescue" ]; then
    warn "Hostname is not 'rescue' - make sure you're running this in Hetzner rescue system"
fi

log "Found installimage at: $INSTALLIMAGE_PATH"

# Detect available NVMe drives
DRIVES=($(lsblk -d -n -o NAME | grep nvme | head -2))
if [ ${#DRIVES[@]} -lt 2 ]; then
    error "At least 2 drives are required for RAID 1 setup. Found: ${#DRIVES[@]}"
fi

DRIVE1="/dev/${DRIVES[0]}"
DRIVE2="/dev/${DRIVES[1]}"

log "Detected drives: $DRIVE1, $DRIVE2"

# Get drive information
DRIVE1_MODEL=$(lsblk -d -n -o MODEL $DRIVE1 2>/dev/null || echo "Unknown")
DRIVE2_MODEL=$(lsblk -d -n -o MODEL $DRIVE2 2>/dev/null || echo "Unknown")
DRIVE1_SERIAL=$(lsblk -d -n -o SERIAL $DRIVE1 2>/dev/null || echo "Unknown")
DRIVE2_SERIAL=$(lsblk -d -n -o SERIAL $DRIVE2 2>/dev/null || echo "Unknown")

log "Drive 1: $DRIVE1 - Model: $DRIVE1_MODEL, Serial: $DRIVE1_SERIAL"
log "Drive 2: $DRIVE2 - Model: $DRIVE2_MODEL, Serial: $DRIVE2_SERIAL"

# Create the installimage configuration file
# Note: installimage expects /autosetup to be a file, not a directory
cat > /autosetup << EOF
## ======================================================
##  Hetzner Online GmbH - installimage - Ubuntu 24.04 with btrfs RAID
## ======================================================

## ====================
##  HARD DISK DRIVE(S):
## ====================

# Device Model: $DRIVE1_MODEL, Serial Number: $DRIVE1_SERIAL
DRIVE1 $DRIVE1
# Device Model: $DRIVE2_MODEL, Serial Number: $DRIVE2_SERIAL
DRIVE2 $DRIVE2

## ===============
##  SOFTWARE RAID:
## ===============

## activate software RAID?  < 0 | 1 >
SWRAID 1

## Choose the level for the software RAID < 0 | 1 | 10 >
SWRAIDLEVEL 1

## ==========
##  HOSTNAME:
## ==========

## which hostname should be set?
HOSTNAME Ubuntu-2404-noble-amd64-btrfs

## ================
##  NETWORK CONFIG:
## ================

# IPV4_ONLY no

## =============
##  MISC CONFIG:
## =============

USE_KERNEL_MODE_SETTING yes

## ==========================
##  PARTITIONS / FILESYSTEMS:
## ==========================

## Using btrfs with subvolumes for better snapshot and management capabilities
## RAID 1 provides redundancy across both drives

# Boot partition (cannot be on btrfs RAID)
PART /boot ext4 1024M

# Swap partition
PART swap swap 4G

# Main btrfs volume for root filesystem
PART btrfs.main btrfs all

## btrfs subvolume definitions:
## Using @ prefix convention for subvolumes
SUBVOL btrfs.main @           /
SUBVOL btrfs.main @home       /home
SUBVOL btrfs.main @var        /var
SUBVOL btrfs.main @var/log    /var/log
SUBVOL btrfs.main @tmp        /tmp
SUBVOL btrfs.main @opt        /opt
SUBVOL btrfs.main @srv        /srv
SUBVOL btrfs.main @snapshots  /.snapshots

## ========================
##  OPERATING SYSTEM IMAGE:
## ========================

## Ubuntu 24.04 LTS Noble Numbat
IMAGE /root/.oldroot/nfs/install/../images/Ubuntu-2404-noble-amd64-base.tar.gz
EOF

log "Created installimage configuration at /autosetup"

# Create autosetup directory for additional scripts
mkdir -p /autosetup_scripts

# Create post-installation script for btrfs optimization
cat > /autosetup_scripts/post_install.sh << 'EOF'
#!/bin/bash

# Post-installation script for btrfs optimization

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

log "Starting post-installation btrfs optimization..."

# Install btrfs-progs if not already installed
if ! command -v btrfs &> /dev/null; then
    log "Installing btrfs-progs..."
    apt-get update
    apt-get install -y btrfs-progs
fi

# Set btrfs mount options for better performance and features
log "Configuring btrfs mount options..."

# Backup original fstab
cp /etc/fstab /etc/fstab.backup

# Update fstab with optimized btrfs mount options
sed -i 's/btrfs\s\+defaults/btrfs defaults,noatime,compress=zstd:3,space_cache=v2,autodefrag/' /etc/fstab

# Create btrfs maintenance scripts
mkdir -p /etc/cron.weekly

# Weekly balance script
cat > /etc/cron.weekly/btrfs-balance << 'BALANCE_EOF'
#!/bin/bash
# Weekly btrfs balance to optimize space usage
/usr/bin/btrfs balance start -dusage=50 -musage=50 / 2>/dev/null || true
BALANCE_EOF

chmod +x /etc/cron.weekly/btrfs-balance

# Weekly scrub script for data integrity
cat > /etc/cron.weekly/btrfs-scrub << 'SCRUB_EOF'
#!/bin/bash
# Weekly btrfs scrub for data integrity check
/usr/bin/btrfs scrub start / 2>/dev/null || true
SCRUB_EOF

chmod +x /etc/cron.weekly/btrfs-scrub

# Install and configure snapper for automatic snapshots
log "Installing and configuring snapper for automatic snapshots..."
apt-get install -y snapper

# Create snapper config for root
snapper -c root create-config /

# Configure snapper for reasonable snapshot retention
snapper -c root set-config TIMELINE_CREATE=yes
snapper -c root set-config TIMELINE_CLEANUP=yes
snapper -c root set-config NUMBER_CLEANUP=yes
snapper -c root set-config NUMBER_MIN_AGE=1800
snapper -c root set-config NUMBER_LIMIT=50
snapper -c root set-config NUMBER_LIMIT_IMPORTANT=10

# Enable snapper timer
systemctl enable snapper-timeline.timer
systemctl enable snapper-cleanup.timer

log "Post-installation btrfs optimization completed"
EOF

chmod +x /autosetup_scripts/post_install.sh

log "Created post-installation script at /autosetup_scripts/post_install.sh"

# Create a script to monitor RAID status
cat > /autosetup_scripts/raid_monitor.sh << 'EOF'
#!/bin/bash

# RAID monitoring script for btrfs

check_btrfs_raid() {
    echo "=== Btrfs RAID Status ==="
    btrfs filesystem show
    echo
    echo "=== Btrfs Device Stats ==="
    btrfs device stats /
    echo
    echo "=== Btrfs Filesystem Usage ==="
    btrfs filesystem usage /
}

# Check if btrfs tools are available
if command -v btrfs &> /dev/null; then
    check_btrfs_raid
else
    echo "btrfs-progs not installed. Install with: apt-get install btrfs-progs"
fi
EOF

chmod +x /autosetup_scripts/raid_monitor.sh

log "Created RAID monitoring script at /autosetup_scripts/raid_monitor.sh"

# Verify configuration
log "Verifying configuration..."
if [ ! -f /autosetup ]; then
    error "Failed to create /autosetup configuration file"
fi

log "Configuration files created successfully:"
log "  - /autosetup (main installation config)"

# Run installimage with the configuration
log "Starting Ubuntu 24.04 installation with btrfs RAID 1..."
log "This will DESTROY ALL DATA on $DRIVE1 and $DRIVE2"

read -p "Are you sure you want to continue? (yes/no): " confirm
if [ "$confirm" != "yes" ]; then
    log "Installation cancelled by user"
    exit 0
fi

# Execute installimage
log "Executing installimage..."
$INSTALLIMAGE_PATH -a -c /autosetup

if [ $? -eq 0 ]; then
    log "Installation completed successfully!"
    log "The system will reboot automatically."
    log "After reboot, run /autosetup_scripts/post_install.sh to optimize btrfs settings"
    log "Use /autosetup_scripts/raid_monitor.sh to monitor RAID status"
else
    error "Installation failed. Check the logs for details."
fi

log "Installation process completed."
EOF