#!/bin/bash
set -euo pipefail

# -----------------------------------------------------------------------------
# Configuration & Context Setup
# -----------------------------------------------------------------------------
# Force working directory to the script's location
cd "$(dirname "$(readlink -f "$0")")"

# Colors for logging
GREEN="\033[0;32m"
YELLOW="\033[1;33m"
RED="\033[0;31m"
BLUE="\033[0;34m"
NC="\033[0m"

log_info()  { echo -e "${BLUE}[INFO]${NC} $1"; }
log_success() { echo -e "${GREEN}[OK]${NC} $1"; }
log_warn()  { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

# Elevate privileges early and keep sudo alive
sudo --validate
while true; do sudo --non-interactive true; sleep 60; kill -0 "$$" || exit; done 2>/dev/null &

# Safety Check: Prevent running as root directly
if [ "$EUID" -eq 0 ]; then
    log_error "Do not run this script directly as root or with 'sudo ./setup.sh'."
    log_error "Run it as your normal user. Sudo prompts will be issued as needed."
    exit 1
fi

USER_HOME="$HOME"
USER_NAME="$(whoami)"

log_info "Starting Arch Linux post-installation configuration for user: ${USER_NAME}..."

# -----------------------------------------------------------------------------
# 1. Package Management & System Updates
# -----------------------------------------------------------------------------
log_info "Step 1: Configuring Mirrors & Packages..."

# Update mirrorlist first
log_info "Updating pacman mirrorlist via Reflector..."
sudo pacman -Sy --noconfirm --needed reflector
sudo reflector --country 'Austria' --latest 10 --sort rate --save /etc/pacman.d/mirrorlist

# Bootstrap Paru AUR helper if not installed
if ! command -v paru >/dev/null 2>&1; then
    log_info "Paru not found. Installing dependencies and bootstrapping Paru..."
    sudo pacman -Syu --noconfirm --needed base-devel git
    
    git submodule update --init --recursive aur.archlinux.org/paru
    pushd aur.archlinux.org/paru >/dev/null
    git pull origin master || git pull
    makepkg -si --noconfirm
    popd >/dev/null
else
    log_info "Paru is already installed."
fi

# Install packages listed in 'all'
if [ -f "./all" ]; then
    log_info "Installing packages from list ('all')..."
    paru -S --needed --noconfirm - < ./all
else
    log_warn "Package list file './all' not found. Skipping bulk package installation."
fi

# Remove orphan packages safely
log_info "Cleaning up orphan packages..."
paru -Qtdq | xargs -r paru -Rns --noconfirm || true

# -----------------------------------------------------------------------------
# 2. System Configuration (/etc & Security Overrides)
# -----------------------------------------------------------------------------
log_info "Step 2: Deploying /etc System Configurations & Security Hardening..."

# Copy basic configuration files safely
sudo cp -v ./etc/keyd/keyd.conf /etc/keyd/
sudo cp -v ./etc/pacman.conf /etc/
sudo cp -v ./etc/kernel/cmdline /etc/kernel/
sudo cp -v ./etc/modprobe.d/blacklist-security.conf /etc/modprobe.d/
sudo cp -v ./etc/pam.d/system-local-login /etc/pam.d/
sudo cp -v ./etc/pam.d/system-login /etc/pam.d/
sudo cp -v ./etc/pam.d/greetd /etc/pam.d/
sudo cp -v ./etc/login.defs /etc/
sudo cp -v ./etc/updatedb.conf /etc/
sudo visudo -cf ./etc/sudoers.d/00-hardening && sudo cp -v ./etc/sudoers.d/00-hardening /etc/sudoers.d/
sudo cp -v ./etc/systemd/coredump.conf /etc/systemd/
sudo cp -v ./etc/sysctl.d/99-security.conf /etc/sysctl.d/
sudo cp -v ./etc/cups/cupsd.conf /etc/cups/

# Append tmpfs to /etc/fstab securely using tee
if [ -f "./etc/fstab" ]; then
    if ! grep -q tmpfs /etc/fstab; then
        log_info "Appending tmpfs configurations to /etc/fstab..."
        sudo tee -a /etc/fstab < ./etc/fstab > /dev/null
    else
        log_info "tmpfs entry already present in /etc/fstab. Skipping."
    fi
fi

# Configure UFW Firewall
log_info "Configuring UFW Firewall..."
sudo ufw default deny incoming
sudo ufw default allow outgoing
sudo ufw default deny routed
sudo ufw limit ssh
sudo ufw allow from 192.168.1.0/24 to any port 8080 proto tcp
sudo ufw logging low
sudo ufw --force enable
sudo systemctl enable --now ufw

# Configure Fail2ban
sudo cp -v ./etc/fail2ban/jail.d/sshd.conf /etc/fail2ban/jail.d/

# Mkinitcpio presets
sudo cp -v ./etc/mkinitcpio.d/linux-hardened.preset /etc/mkinitcpio.d/ 2>/dev/null || true
sudo cp -v ./etc/mkinitcpio.d/linux-lts.preset /etc/mkinitcpio.d/ 2>/dev/null || true
sudo cp -v ./etc/mkinitcpio.d/linux.preset /etc/mkinitcpio.d/ 2>/dev/null || true

# Systemd Service Hardening Drop-ins
log_info "Applying systemd service drop-in hardening overrides..."
for svc in sshd fail2ban cups; do
    sudo mkdir -p "/etc/systemd/system/${svc}.service.d"
    if [ -f "./etc/systemd/system/${svc}.service.d/override.conf" ]; then
        sudo cp -v "./etc/systemd/system/${svc}.service.d/override.conf" "/etc/systemd/system/${svc}.service.d/"
    fi
done

# CUPS Log directory permissions
sudo mkdir -p /var/log/cups
sudo chown root:lp /var/log/cups
sudo chmod 755 /var/log/cups

# -----------------------------------------------------------------------------
# 3. Kernel Parameters, Rebuilds & Systemd Services
# -----------------------------------------------------------------------------
log_info "Step 3: Reloading System Configuration & Enabling Services..."

sudo sysctl --system
sudo mkinitcpio -P
sudo systemctl daemon-reload

# Enable User-Level Socket/Service Units
log_info "Enabling user-level systemd services..."
systemctl --user enable gnome-keyring-daemon.socket
systemctl --user enable ssh-agent.socket

# Enable System-Level Services
log_info "Enabling system-level services and timers..."
SYSTEM_SERVICES=(
    "apparmor.service"
    "auditd.service"
    "reflector.timer"
    "paccache.timer"
    "keyd.service"
    "proton.VPN.service"
    "snapper-cleanup.timer"
    "snapper-timeline.timer"
    "fail2ban.service"
    "fstrim.timer"
)

for svc in "${SYSTEM_SERVICES[@]}"; do
    sudo systemctl enable --now "$svc"
done

# Set all AppArmor profiles to enforce mode
log_info "Enforcing AppArmor profiles..."
if command -v aa-enforce >/dev/null 2>&1; then
    sudo aa-enforce /etc/apparmor.d/* 2>/dev/null || true
else
    log_warn "aa-enforce tool not found. AppArmor profiles will load in default mode via apparmor.service."
fi
# Restart active core services if already running
log_info "Restarting active target services..."
for svc in sshd fail2ban cups; do
    if systemctl is-active --quiet "$svc"; then
        log_info "Restarting active service: $svc"
        sudo systemctl restart "$svc"
    else
        log_info "Skipping restart for inactive service: $svc"
    fi
done

# -----------------------------------------------------------------------------
# 4. Deploy Dotfiles & User Configuration
# -----------------------------------------------------------------------------
log_info "Step 4: Deploying Dotfiles..."

# User dotfiles
cp -v dotfiles/.bash_profile "$USER_HOME/"
cp -v dotfiles/.git-prompt.sh "$USER_HOME/"
cp -v dotfiles/.bashrc "$USER_HOME/"
cp -v dotfiles/.inputrc "$USER_HOME/"

mkdir -pv "$USER_HOME/.config/" "$USER_HOME/.local/bin/" "$USER_HOME/Pictures/"

rm -rf "$USER_HOME/.config/fastfetch"
cp -rv dotfiles/.config/fastfetch "$USER_HOME/.config/"

rm -rf "$USER_HOME/.config/nvim"
cp -rv dotfiles/.config/nvim "$USER_HOME/.config/"

if [ -d "Pictures" ]; then
    cp -v Pictures/* "$USER_HOME/Pictures/" 2>/dev/null || true
fi

# Safely handle Go directory structure
mkdir -p "$USER_HOME/Projects"
if [ -d "$USER_HOME/go" ] && [ "$USER_HOME/go" != "$USER_HOME/Projects/go" ]; then
    log_info "Merging ~/go into ~/Projects/go..."
    mkdir -p "$USER_HOME/Projects/go"
    cp -a "$USER_HOME/go/." "$USER_HOME/Projects/go/"
    rm -rf "$USER_HOME/go"
fi

if command -v go >/dev/null 2>&1; then
    go env -w GOPATH="$USER_HOME/Projects/go"
fi

# Root dotfiles deployment
log_info "Deploying root user dotfiles..."
sudo cp -v dotfiles/.bash_profile /root/
sudo cp -v dotfiles/.git-prompt.sh /root/
sudo cp -v rootfiles/.bashrc /root/
sudo mkdir -pv /root/.config/
sudo cp -rv dotfiles/.config/nvim /root/.config/

log_success "Arch Linux post-installation setup completed successfully!"

# -----------------------------------------------------------------------------
# 5. System Reboot Prompt
# -----------------------------------------------------------------------------
echo ""
read -r -p "Reboot now to apply all kernel parameters and LSM changes? [y/N] " response
case "$response" in
    [yY][eE][sS]|[yY])
        log_info "Rebooting system..."
        sudo reboot
        ;;
    *)
        log_info "Reboot skipped. Remember to reboot manually later to activate AppArmor and auditd kernel settings."
        ;;
esac