#!/bin/bash
set -euo pipefail

BANNER=$(cat << "EOF"
           :^7J5GB##&&##GPY?~:            
        ^75B&@@@@@@&&&@@@@@@@#GJ~:        
       5&@@@&B5?7~^^^^^~!7YP#@@@@#!       
       Y##P7^               :~JB#B!       
        ::                      :         
  7PP?:          :^~!!~^:          :?PP7  
 :B@@B:        !5B&@@@@&B5!        :#@@B: 
  :!!:       ^G@@@&BPPB@@@@G^       :!!:  
            :B@@@5^    ^5@@@B:            
:7J7:       !@@@#       :&@@@~       :?J7:
J@@@5       :#@@@Y:    :Y@@@B:       5@@@J
!@@@&^       ~B@@@&G55G&@@@B~       ~&@@@~
 5@@@G:       :7P#@@@@@@#P7:       :B@@@Y 
 :P@@@B~         :~!77!~:         ~B@@@P  
   Y@@@&Y^                      ^5@@@@J   
    !G@@@&P7^                ^7P&@@@G~    
      !P&@@@&B?     ::     ?B&@@@&P!      
        ^75#&&Y   :P&&5:   5&&B57^        
           :^^    :P&&5:    ^^:           
                    ^^                    

                [nullpoint]
EOF
)

TANG_SERVERS=(
#    "https://tang1.example.com your-thumbprint-1"
#    "https://tang2.example.com your-thumbprint-2"
)
TPM_ENABLED=true
TPM_PCR_BANK="sha256"
TPM_PCR_IDS="0,1,2,3,4,5,6,7,8"
ALMA_USER="null"
ENABLE_MOTD=true
# REQUIRED: Set your SSH public key here - installation will fail without it!
SSH_KEY="ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIOkoTn2NreAXMriOUqzyj3YoFW6jMo9B5B+3R5k8yrMi dodox@dodox-ProArt"

########################################################
# Config End
########################################################

set -euo pipefail

echo -e "\n$BANNER"
echo -e "\n[+] Starting post-installation configuration..."

# Check for SSH key
if [ -z "${SSH_KEY:-}" ]; then
    echo "ERROR: No SSH key configured!"
    echo "You must set SSH_KEY variable at the top of this script."
    exit 1
fi

# Check for TPM
echo "[+] Checking for TPM..."
if [ ! -d "/sys/class/tpm/tpm0" ]; then
    echo "WARNING: No TPM detected!"
    TPM_ENABLED=false
else
    echo "TPM detected."
    TPM_ENABLED=true
fi

# Install basic packages first
echo "[+] Installing basic packages..."
dnf install -y epel-release || exit 1
dnf config-manager --set-enabled crb || exit 1
dnf install -y zsh git wget curl || exit 1

# Create user and add SSH key
echo "[+] Creating user ${ALMA_USER}..."
useradd -m -G wheel -s /bin/zsh ${ALMA_USER} || exit 1
mkdir -p /home/${ALMA_USER}/.ssh
echo "${SSH_KEY}" > /home/${ALMA_USER}/.ssh/authorized_keys
chmod 700 /home/${ALMA_USER}/.ssh
chmod 600 /home/${ALMA_USER}/.ssh/authorized_keys
chown -R ${ALMA_USER}:${ALMA_USER} /home/${ALMA_USER}/.ssh

# Configure passwordless sudo
echo "[+] Configuring passwordless sudo for ${ALMA_USER}..."
echo "${ALMA_USER} ALL=(ALL) NOPASSWD: ALL" > /etc/sudoers.d/99-${ALMA_USER}
chmod 440 /etc/sudoers.d/99-${ALMA_USER}

# Install oh-my-zsh and powerlevel10k
echo "[+] Installing oh-my-zsh and powerlevel10k for ${ALMA_USER} and root..."

# Install for user
su - ${ALMA_USER} -c 'export RUNZSH=no CHSH=no KEEP_ZSHRC=yes && bash -c "$(wget -O- https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)"' 2>/dev/null || echo "WARNING: user oh-my-zsh installation failed"
su - ${ALMA_USER} -c 'git clone --depth=1 https://github.com/romkatv/powerlevel10k.git ~/.oh-my-zsh/custom/themes/powerlevel10k' 2>/dev/null || echo "WARNING: user powerlevel10k installation failed"

# Install for root
export RUNZSH=no CHSH=no KEEP_ZSHRC=yes && bash -c "$(wget -O- https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" 2>/dev/null || echo "WARNING: root oh-my-zsh installation failed"
git clone --depth=1 https://github.com/romkatv/powerlevel10k.git /root/.oh-my-zsh/custom/themes/powerlevel10k 2>/dev/null || echo "WARNING: root powerlevel10k installation failed"

# Install dotfiles from git repo (cloning needed as we're in chroot)
echo "[+] Installing dotfiles for ${ALMA_USER} and root..."

# Install for user
su - ${ALMA_USER} -c '
    cd && 
    git clone https://git.dominik-roth.eu/dodox/nullpoint.git /tmp/nullpoint-dotfiles && 
    cd /tmp/nullpoint-dotfiles/dotfiles && 
    for file in .*; do
        if [ -f "$file" ] && [ "$file" != "." ] && [ "$file" != ".." ]; then
            cp "$file" ~/ 2>/dev/null || true
        fi
    done &&
    cd && rm -rf /tmp/nullpoint-dotfiles
' || echo "WARNING: user dotfiles installation failed"

# Install for root
cd /root && 
git clone https://git.dominik-roth.eu/dodox/nullpoint.git /tmp/nullpoint-dotfiles-root && 
cd /tmp/nullpoint-dotfiles-root/dotfiles && 
for file in .*; do
    if [ -f "$file" ] && [ "$file" != "." ] && [ "$file" != ".." ]; then
        cp "$file" /root/ 2>/dev/null || true
    fi
done &&
cd /root && rm -rf /tmp/nullpoint-dotfiles-root || echo "WARNING: root dotfiles installation failed"

# Set zsh as default shell for root
echo "[+] Setting zsh as default shell for root..."
chsh -s /bin/zsh root

# Set up MOTD
if [ "$ENABLE_MOTD" = true ]; then
    echo "[+] Setting up MOTD..."
    cat > /etc/motd << MOTD
$BANNER
MOTD
fi

# Install additional packages
echo "[+] Installing additional packages..."
dnf install -y \
  clevis clevis-luks tpm2-tools tpm2-tss \
  tmux neovim python3-pip \
  tree gcc make autoconf automake tar bzip2 || exit 1


# Install dropbear for early boot SSH
echo "[+] Installing dropbear for early boot SSH..."
dnf install -y dropbear dracut-network || exit 1

# Install lsd and bat
echo "[+] Installing lsd and bat..."
# Install using fixed versions that should work
LSD_VERSION="1.0.0"
BAT_VERSION="0.24.0"

# Download and install lsd
if curl -sL "https://github.com/lsd-rs/lsd/releases/download/v${LSD_VERSION}/lsd-v${LSD_VERSION}-x86_64-unknown-linux-musl.tar.gz" | tar xz -C /tmp; then
    mv /tmp/lsd-*/lsd /usr/local/bin/ 2>/dev/null || true
    chmod +x /usr/local/bin/lsd 2>/dev/null || true
fi

# Download and install bat
if curl -sL "https://github.com/sharkdp/bat/releases/download/v${BAT_VERSION}/bat-v${BAT_VERSION}-x86_64-unknown-linux-musl.tar.gz" | tar xz -C /tmp; then
    mv /tmp/bat-*/bat /usr/local/bin/ 2>/dev/null || true
    chmod +x /usr/local/bin/bat 2>/dev/null || true
fi

# Create batman script for fancy man pages
cat > /usr/local/bin/batman << 'BATMAN'
#!/bin/bash
export MANPAGER="sh -c 'col -bx | bat -l man -p'"
export MANROFFOPT="-c"
man "$@"
BATMAN
chmod +x /usr/local/bin/batman

# Create .tmp directory for user
mkdir -p /home/${ALMA_USER}/.tmp
chown ${ALMA_USER}:${ALMA_USER} /home/${ALMA_USER}/.tmp

# Configure Clevis for automatic unlock
if [ ${#TANG_SERVERS[@]} -gt 0 ] || [ "$TPM_ENABLED" = true ]; then
    echo "[+] Configuring Clevis for automatic unlock..."
    
    # Find LUKS devices
    LUKS_DEVICES=$(lsblk -o NAME,FSTYPE -nr | grep crypto_LUKS | cut -d' ' -f1)
    
    for device in $LUKS_DEVICES; do
        DEVICE_PATH="/dev/${device}"
        echo "Configuring Clevis for ${DEVICE_PATH}..."
        
        if [ "$TPM_ENABLED" = true ] && [ ${#TANG_SERVERS[@]} -eq 0 ]; then
            # TPM only
            clevis luks bind -d "$DEVICE_PATH" tpm2 "{\"pcr_bank\":\"$TPM_PCR_BANK\",\"pcr_ids\":\"$TPM_PCR_IDS\"}"
        elif [ "$TPM_ENABLED" = false ] && [ ${#TANG_SERVERS[@]} -gt 0 ]; then
            # Tang only
            for server in "${TANG_SERVERS[@]}"; do
                read -r url thumbprint <<< "$server"
                clevis luks bind -d "$DEVICE_PATH" tang "{\"url\":\"$url\",\"thp\":\"$thumbprint\"}"
            done
        elif [ "$TPM_ENABLED" = true ] && [ ${#TANG_SERVERS[@]} -gt 0 ]; then
            # Both TPM and Tang (require both)
            CONFIG="{\"t\":2,\"pins\":{"
            CONFIG+="\"tpm2\":{\"pcr_bank\":\"$TPM_PCR_BANK\",\"pcr_ids\":\"$TPM_PCR_IDS\"},"
            CONFIG+="\"tang\":{\"t\":1,\"tang\":["
            for server in "${TANG_SERVERS[@]}"; do
                read -r url thumbprint <<< "$server"
                CONFIG+="{\"url\":\"$url\",\"thp\":\"$thumbprint\"},"
            done
            CONFIG="${CONFIG%,}]}}}"
            clevis luks bind -d "$DEVICE_PATH" sss "$CONFIG"
        fi
    done
fi

# Enable Clevis for early boot (only needed for AlmaLinux < 8.7)
echo "[+] Configuring Clevis for early boot..."
OS_VERSION=$(cat /etc/redhat-release | grep -oE '[0-9]+\.[0-9]+' | head -1)
if [[ "$(echo "$OS_VERSION < 8.7" | bc)" -eq 1 ]]; then
    echo "  - Enabling clevis-luks-askpass.path for AlmaLinux $OS_VERSION"
    systemctl enable clevis-luks-askpass.path || true
else
    echo "  - AlmaLinux $OS_VERSION: clevis-luks-askpass.path not needed"
fi

# Configure dropbear for remote unlock
echo "[+] Configuring dropbear SSH for remote unlock..."

# Create custom dracut module for dropbear SSH
mkdir -p /usr/lib/dracut/modules.d/60dropbear-ssh

# Create the module setup script
cat > /usr/lib/dracut/modules.d/60dropbear-ssh/module-setup.sh << 'EOF'
#!/bin/bash

check() {
    require_binaries dropbear dbclient dropbearkey dropbearconvert || return 1
    return 0
}

depends() {
    echo network
}

install() {
    inst_multiple dropbear dbclient dropbearkey dropbearconvert
    
    # Create directories
    inst_dir /etc/dropbear
    inst_dir /var/log
    inst_dir /root/.ssh
    
    # Copy authorized keys if they exist
    if [ -f /etc/dropbear/authorized_keys ]; then
        inst /etc/dropbear/authorized_keys /root/.ssh/authorized_keys
    fi
    
    # Install ED25519 host key only
    keyfile="/etc/dropbear/dropbear_ed25519_host_key"
    if [ ! -f "$keyfile" ]; then
        dropbearkey -t ed25519 -f "$keyfile" 2>/dev/null
    fi
    [ -f "$keyfile" ] && inst "$keyfile"
    
    # Install the service
    inst_simple "$moddir/dropbear.service" /etc/systemd/system/dropbear.service
    systemctl -q --root "$initdir" enable dropbear.service
    
    # Install unlock helper
    inst_simple "$moddir/unlock-luks.sh" /usr/bin/unlock-luks
    chmod 755 "$initdir/usr/bin/unlock-luks"
}
EOF

# Create systemd service for dropbear
cat > /usr/lib/dracut/modules.d/60dropbear-ssh/dropbear.service << 'EOF'
[Unit]
Description=Dropbear SSH Server
After=network-online.target
Wants=network-online.target

[Service]
Type=forking
ExecStart=/usr/sbin/dropbear -R -E -p 22
ExecReload=/bin/kill -HUP $MAINPID
KillMode=process
Restart=on-failure

[Install]
WantedBy=sysinit.target
EOF

# Create unlock helper script
cat > /usr/lib/dracut/modules.d/60dropbear-ssh/unlock-luks.sh << 'EOF'
#!/bin/bash
echo "=== LUKS Remote Unlock Helper ==="
echo ""
echo "Checking for encrypted devices..."

# Show block devices if available
if command -v lsblk >/dev/null 2>&1; then
    echo "Block devices:"
    lsblk -o NAME,SIZE,TYPE,FSTYPE 2>/dev/null || echo "  (lsblk not available)"
else
    echo "Block devices: (listing /dev/sd* and /dev/md*)"
    ls -la /dev/sd* /dev/md* 2>/dev/null || echo "  No standard devices found"
fi

echo ""
echo "Encrypted devices status:"
# Check for LUKS devices waiting to be unlocked
for dev in /dev/mapper/luks-*; do
    if [ -e "$dev" ]; then
        echo "  Found: $dev"
    fi
done

# Check systemd-ask-password files directly
if [ -d /run/systemd/ask-password ]; then
    echo ""
    echo "Password prompts waiting:"
    ls -la /run/systemd/ask-password/ 2>/dev/null
fi

echo ""
echo "Starting unlock process..."
echo "Enter your LUKS passphrase when prompted:"
echo ""

# Run the password agent
if command -v systemd-tty-ask-password-agent >/dev/null 2>&1; then
    systemd-tty-ask-password-agent
else
    echo "ERROR: systemd-tty-ask-password-agent not found!"
    echo "Try running: /lib/systemd/systemd-tty-ask-password-agent"
fi
EOF

chmod +x /usr/lib/dracut/modules.d/60dropbear-ssh/*.sh

# Setup dropbear authorized keys
mkdir -p /etc/dropbear
echo "${SSH_KEY}" > /etc/dropbear/authorized_keys
chmod 600 /etc/dropbear/authorized_keys

# Generate ED25519 host key only (most secure)
echo "[+] Generating ED25519 SSH host key..."

# Use system SSH key if available, otherwise generate dropbear key
openssh_key="/etc/ssh/ssh_host_ed25519_key"
dropbear_key="/etc/dropbear/dropbear_ed25519_host_key"

if [ -f "$openssh_key" ] && command -v dropbearconvert >/dev/null 2>&1; then
    echo "  Converting existing OpenSSH ED25519 key to dropbear format..."
    dropbearconvert openssh dropbear "$openssh_key" "$dropbear_key" 2>/dev/null || {
        echo "  Conversion failed, generating new dropbear key..."
        dropbearkey -t ed25519 -f "$dropbear_key" | grep -v "Generating" || true
    }
elif [ ! -f "$dropbear_key" ]; then
    echo "  Generating new ED25519 key..."
    dropbearkey -t ed25519 -f "$dropbear_key" | grep -v "Generating" || true
    
    # Also generate OpenSSH format to prevent key mismatch after boot
    if command -v ssh-keygen >/dev/null 2>&1; then
        echo "  Generating matching OpenSSH key..."
        mkdir -p /etc/ssh
        # Extract public key and generate OpenSSH private key
        dropbearkey -y -f "$dropbear_key" | grep "^ssh-" > "${openssh_key}.pub"
        # Note: Direct conversion from dropbear to openssh private key requires dropbearconvert
        # For now, we'll have different keys but document the solution
    fi
fi

# Display SHA256 fingerprint
if command -v ssh-keygen >/dev/null 2>&1; then
    fingerprint=$(dropbearkey -y -f "$dropbear_key" | ssh-keygen -lf - -E sha256 2>/dev/null | awk '{print $2}')
    if [ -n "$fingerprint" ]; then
        echo "  SHA256 fingerprint: $fingerprint"
        echo "  Note: This is the initramfs (rescue) SSH fingerprint."
        echo "  The normal system SSH may have a different fingerprint."
    fi
fi

# Configure dracut
cat > /etc/dracut.conf.d/60-dropbear-ssh.conf << 'EOF'
# Enable network and dropbear SSH
add_dracutmodules+=" network dropbear-ssh "
# Network configuration
kernel_cmdline="rd.neednet=1"
EOF

# Regenerate initramfs
echo "[+] Regenerating initramfs..."
dracut -f --regenerate-all

# Enable required services
echo "[+] Enabling services..."
# systemctl enable stratisd  # Not needed without Stratis
systemctl enable sshd

# Secure SSH configuration
echo "[+] Securing SSH..."
{
    # Disable root login
    sed -i 's/^#*PermitRootLogin.*/PermitRootLogin no/' /etc/ssh/sshd_config
    
    # Only allow SSH key authentication
    sed -i 's/^#*PasswordAuthentication.*/PasswordAuthentication no/' /etc/ssh/sshd_config
    sed -i 's/^#*ChallengeResponseAuthentication.*/ChallengeResponseAuthentication no/' /etc/ssh/sshd_config
    sed -i 's/^#*UsePAM.*/UsePAM no/' /etc/ssh/sshd_config
    
    # Only allow specific user
    echo "AllowUsers ${ALMA_USER}" >> /etc/ssh/sshd_config
    
    echo "  - Root login disabled"
    echo "  - Password authentication disabled"
    echo "  - Only user '${ALMA_USER}' allowed"
}

# Set SELinux to enforcing
echo "[+] Setting SELinux to enforcing..."
sed -i 's/^SELINUX=.*/SELINUX=enforcing/' /etc/selinux/config

echo "✅ Post-installation complete!"
echo ""
echo "IMPORTANT: The LUKS passphrase is set in install.conf"
echo "Save it securely for recovery purposes."
echo ""
echo "After reboot:"
echo "- SSH to port 22 for remote unlock: ssh root@<server-ip>"
echo "- Run 'unlock-luks' and follow the instructions to unlock LUKS"
echo "- Once unlocked, SSH to port 22 as user '${ALMA_USER}'"
echo "- LUKS passphrase: [see installer output]"