#!/bin/bash
# Post-install configuration:
#   - hostname / locale / keyboard / timezone / hwclock
#   - users (root + moze) + sudo
#   - services (NetworkManager, lightdm, dbus, polkitd, docker, bluetoothd,
#     acpid, tlp, sshd[disabled], dhcpcd[disabled in favor of NM])
#   - zram swap (zramen)
#   - NVIDIA PRIME render-offload setup
#   - SSH config copy
#   - Nix bootstrap (nix-daemon service + first-boot user package install)

# shellcheck source=common.sh
source "$(dirname "${BASH_SOURCE[0]}")/common.sh"

configure_system() {
    step "System configuration"
    local TARGET="${TARGET:-/mnt}"

    # ----- hostname -----
    echo "$HOSTNAME" > "$TARGET/etc/hostname"
    cat > "$TARGET/etc/hosts" <<EOF
127.0.0.1   localhost
::1         localhost
127.0.1.1   $HOSTNAME.localdomain $HOSTNAME
EOF

    # ----- rc.conf -----
    cat > "$TARGET/etc/rc.conf" <<EOF
KEYMAP="$KEYMAP"
HARDWARECLOCK="$HARDWARECLOCK"
TIMEZONE="$TIMEZONE"
EOF

    # ----- timezone (also create symlink for systems checking it) -----
    run_chroot "ln -sf /usr/share/zoneinfo/$TIMEZONE /etc/localtime"

    # ----- locales (glibc) -----
    if [[ -f "$TARGET/etc/default/libc-locales" ]]; then
        sed -i "s/^#\($LOCALE.*\)/\1/" "$TARGET/etc/default/libc-locales"
        run_chroot "xbps-reconfigure -f glibc-locales"
    fi
    cat > "$TARGET/etc/locale.conf" <<EOF
LANG=$LANG
LC_ALL=$LANG
EOF

    # ----- vconsole (for early TTY keymap) -----
    cat > "$TARGET/etc/vconsole.conf" <<EOF
KEYMAP=$KEYMAP
EOF

    ok "locale / keymap / tz configured"
}

configure_users() {
    step "Creating users"
    local TARGET="${TARGET:-/mnt}"

    # ----- root password (chpasswd via stdin — no shell expansion of $) -----
    set_chroot_password root "$ROOT_PASSWORD"

    # ----- ensure groups exist -----
    run_chroot "groupadd -f docker"
    run_chroot "groupadd -f plugdev"

    # ----- create user (idempotent: tolerate 'already exists', fail on real errors) -----
    if ! run_chroot "id -u $USERNAME >/dev/null 2>&1"; then
        run_chroot "useradd -m -u $USER_UID -G $USER_GROUPS -s $DEFAULT_SHELL -c '$USER_FULLNAME' $USERNAME"
    else
        log "user $USERNAME already exists — skipping useradd"
        run_chroot "usermod -G $USER_GROUPS -s $DEFAULT_SHELL $USERNAME"
    fi
    set_chroot_password "$USERNAME" "$USER_PASSWORD"

    # ----- sudoers: wheel group -----
    mkdir -p "$TARGET/etc/sudoers.d"
    if [[ "${TEST_MODE:-0}" == "1" ]]; then
        # Test harness needs passwordless sudo to run smoke checks via SSH.
        cat > "$TARGET/etc/sudoers.d/10-wheel" <<'EOF'
%wheel ALL=(ALL:ALL) NOPASSWD: ALL
Defaults env_keep += "EDITOR"
EOF
    else
        cat > "$TARGET/etc/sudoers.d/10-wheel" <<'EOF'
%wheel ALL=(ALL:ALL) ALL
Defaults env_keep += "EDITOR"
EOF
    fi
    chmod 440 "$TARGET/etc/sudoers.d/10-wheel"

    # git: GUI askpass so prompts work without a controlling terminal
    install -d -m 0755 "$TARGET/usr/local/bin"
    cat > "$TARGET/usr/local/bin/git-askpass" <<'EOF'
#!/bin/sh
for cmd in zenity qarma; do
    command -v "$cmd" >/dev/null 2>&1 || continue
    case "$1" in
        *[Uu]sername*) exec "$cmd" --entry   --title="Git Credentials" --text="$1" ;;
        *)             exec "$cmd" --password --title="Git Credentials" --text="$1" ;;
    esac
done
printf '%s' "$1" >&2; read -r answer; printf '%s\n' "$answer"
EOF
    chmod 0755 "$TARGET/usr/local/bin/git-askpass"
    cat > "$TARGET/etc/gitconfig" <<'EOF'
[core]
	askPass = /usr/local/bin/git-askpass
EOF

    ok "user '$USERNAME' created and added to: $USER_GROUPS"
}

configure_ssh_config() {
    step "Installing SSH config for $USERNAME"
    local TARGET="${TARGET:-/mnt}"
    local src="$SSH_SOURCE_DIR"
    local dst="$TARGET/home/$USERNAME/$SSH_TARGET_DIR_REL"

    if [[ ! -d "$src" ]]; then
        warn "no SSH source dir at $src — skipping"
        return 0
    fi

    install -d -m 0700 "$dst"
    cp -a "$src"/. "$dst/"
    # Tighten perms.
    find "$dst" -type d -exec chmod 700 {} +
    find "$dst" -type f -exec chmod 600 {} +
    find "$dst" -type f -name '*.pub' -exec chmod 644 {} +
    [[ -f "$dst/known_hosts"     ]] && chmod 644 "$dst/known_hosts"
    [[ -f "$dst/known_hosts.old" ]] && chmod 644 "$dst/known_hosts.old"
    [[ -f "$dst/config"          ]] && chmod 600 "$dst/config"

    run_chroot "chown -R $USERNAME:$USERNAME /home/$USERNAME/$SSH_TARGET_DIR_REL"
    ok "SSH config installed at /home/$USERNAME/$SSH_TARGET_DIR_REL"
}

configure_nvidia_prime() {
    step "Configuring NVIDIA PRIME render-offload"
    local TARGET="${TARGET:-/mnt}"

    # 1) Xorg config: Intel as primary, NVIDIA as PRIME provider.
    install -d -m 0755 "$TARGET/etc/X11/xorg.conf.d"
    cat > "$TARGET/etc/X11/xorg.conf.d/10-intel.conf" <<'EOF'
Section "OutputClass"
    Identifier "intel"
    MatchDriver "i915"
    Driver "modesetting"
EndSection
EOF
    cat > "$TARGET/etc/X11/xorg.conf.d/20-nvidia.conf" <<'EOF'
Section "OutputClass"
    Identifier "nvidia"
    MatchDriver "nvidia-drm"
    Driver "nvidia"
    Option "AllowEmptyInitialConfiguration"
    Option "PrimaryGPU" "no"
    ModulePath "/usr/lib/nvidia/xorg"
    ModulePath "/usr/lib/xorg/modules"
EndSection
EOF

    # 2) Blacklist nouveau so it doesn't grab the GPU before nvidia loads.
    install -d -m 0755 "$TARGET/etc/modprobe.d"
    cat > "$TARGET/etc/modprobe.d/blacklist-nouveau.conf" <<'EOF'
blacklist nouveau
options nouveau modeset=0
EOF

    # btusb USB autosuspend causes firmware download failures on Intel BT.
    cat > "$TARGET/etc/modprobe.d/btusb-quirks.conf" <<'EOF'
options btusb enable_autosuspend=0
EOF

    # 3) Modules to load early for KMS.
    install -d -m 0755 "$TARGET/etc/modules-load.d"
    cat > "$TARGET/etc/modules-load.d/nvidia.conf" <<'EOF'
nvidia
nvidia_modeset
nvidia_uvm
nvidia_drm
EOF

    # 4) Wrapper script + desktop file: run any app on NVIDIA via `prime-run`.
    install -d -m 0755 "$TARGET/usr/local/bin"
    cat > "$TARGET/usr/local/bin/prime-run" <<'EOF'
#!/bin/sh
# Run a program on the NVIDIA dGPU via PRIME render offload.
exec env __NV_PRIME_RENDER_OFFLOAD=1 \
         __VK_LAYER_NV_optimus=NVIDIA_only \
         __GLX_VENDOR_LIBRARY_NAME=nvidia \
         "$@"
EOF
    chmod 0755 "$TARGET/usr/local/bin/prime-run"

    # 5) dracut: include nvidia modules in initramfs; explicitly omit nouveau.
    install -d -m 0755 "$TARGET/etc/dracut.conf.d"
    cat > "$TARGET/etc/dracut.conf.d/10-nvidia.conf" <<'EOF'
add_drivers+=" nvidia nvidia_modeset nvidia_uvm nvidia_drm "
omit_drivers+=" nouveau "
EOF

    ok "NVIDIA PRIME offload configured (use 'prime-run <app>')"
}

configure_zram() {
    [[ "${ZRAM_ENABLE:-yes}" == "yes" ]] || return 0
    step "Configuring zram (zramen)"
    local TARGET="${TARGET:-/mnt}"
    install -d -m 0755 "$TARGET/etc/default"
    cat > "$TARGET/etc/default/zramen" <<EOF
ALGO=zstd
PERCENT=$ZRAM_SIZE_PCT
PRIORITY=100
EOF
    ok "zram configured at $ZRAM_SIZE_PCT% RAM"
}

configure_nix() {
    [[ "${ENABLE_NIX:-yes}" == "yes" ]] || return 0
    step "Configuring nix multi-user (daemon mode)"
    local TARGET="${TARGET:-/mnt}"
    # Void's `nix` xbps package installs to /usr and ships a runit service.
    install -d -m 0755 "$TARGET/etc/nix"
    cat > "$TARGET/etc/nix/nix.conf" <<'EOF'
experimental-features = nix-command flakes
build-users-group = nixbld
auto-optimise-store = true
sandbox = true
EOF

    # First-boot script: as $USERNAME, install user packages.
    install -d -m 0755 "$TARGET/usr/local/libexec"
    cat > "$TARGET/usr/local/libexec/first-boot-nix.sh" <<EOF
#!/bin/bash
set -e
mark=/var/lib/first-boot-nix.done
[[ -f "\$mark" ]] && exit 0

# Wait for nix-daemon to be available.
for _ in \$(seq 1 60); do
    [[ -S /nix/var/nix/daemon-socket/socket ]] && break
    sleep 2
done

if [[ ! -S /nix/var/nix/daemon-socket/socket ]]; then
    echo "nix-daemon not available; aborting first-boot nix install" >&2
    exit 0
fi

su - $USERNAME -c '
    set -e
    . /etc/profile.d/nix.sh 2>/dev/null || true
    # google-chrome / spotify / discord are unfree -> need allow-unfree + --impure.
    export NIXPKGS_ALLOW_UNFREE=1
    nix profile install --impure ${NIX_USER_PACKAGES[*]} || true
'

mkdir -p "\$(dirname "\$mark")"
touch "\$mark"
EOF
    chmod 0755 "$TARGET/usr/local/libexec/first-boot-nix.sh"

    # Persistent nixpkgs config so the installed user can install unfree packages
    # without needing to export NIXPKGS_ALLOW_UNFREE=1 every time.
    install -d -m 0755 "$TARGET/home/$USERNAME/.config/nixpkgs"
    echo '{ allowUnfree = true; }' > "$TARGET/home/$USERNAME/.config/nixpkgs/config.nix"
    run_chroot "chown -R $USERNAME:$USERNAME /home/$USERNAME/.config/nixpkgs"

    # runit one-shot service.
    install -d -m 0755 "$TARGET/etc/sv/first-boot-nix"
    cat > "$TARGET/etc/sv/first-boot-nix/run" <<'EOF'
#!/bin/sh
exec 2>&1
/usr/local/libexec/first-boot-nix.sh
exec chpst -b first-boot-nix pause
EOF
    chmod 0755 "$TARGET/etc/sv/first-boot-nix/run"
    cat > "$TARGET/etc/sv/first-boot-nix/finish" <<'EOF'
#!/bin/sh
sv down first-boot-nix
EOF
    chmod 0755 "$TARGET/etc/sv/first-boot-nix/finish"

    ok "Nix configured; user packages will install on first boot"
}

install_vscode_real() {
    # Install official Microsoft VS Code (the real proprietary build), NOT the
    # `vscode` xbps package which is actually code-oss and ships `code-oss`.
    step "Installing official Microsoft VS Code"
    local TARGET="${TARGET:-/mnt}"
    local url="https://update.code.visualstudio.com/latest/linux-x64/stable"
    local tmp="$TARGET/tmp/vscode.tar.gz"

    install -d -m 0755 "$TARGET/opt" "$TARGET/usr/local/bin"

    if ! run_chroot "curl -fsSL --retry 3 -o /tmp/vscode.tar.gz '$url'"; then
        warn "failed to download VS Code; skipping (install manually later)"
        rm -f "$tmp"
        return 0
    fi

    # Tarball extracts to VSCode-linux-x64/. Move it to /opt/vscode.
    rm -rf "$TARGET/opt/vscode"
    if ! run_chroot "tar -xzf /tmp/vscode.tar.gz -C /opt && mv /opt/VSCode-linux-x64 /opt/vscode"; then
        warn "failed to extract VS Code tarball; skipping"
        rm -f "$tmp"
        return 0
    fi
    rm -f "$tmp"

    # `code` shim on PATH.
    ln -sf /opt/vscode/bin/code "$TARGET/usr/local/bin/code"

    # Desktop entry so it shows up in the Cinnamon menu.
    install -d -m 0755 "$TARGET/usr/local/share/applications"
    cat > "$TARGET/usr/local/share/applications/code.desktop" <<'EOF'
[Desktop Entry]
Name=Visual Studio Code
Comment=Code Editing. Redefined.
GenericName=Text Editor
Exec=/opt/vscode/bin/code %F
Icon=/opt/vscode/resources/app/resources/linux/code.png
Type=Application
StartupNotify=false
StartupWMClass=Code
Categories=TextEditor;Development;IDE;
MimeType=text/plain;inode/directory;application/x-code-workspace;
Actions=new-empty-window;
Keywords=vscode;

[Desktop Action new-empty-window]
Name=New Empty Window
Exec=/opt/vscode/bin/code --new-window %F
Icon=/opt/vscode/resources/app/resources/linux/code.png
EOF

    ok "VS Code installed at /opt/vscode (use 'code' on PATH)"
}

enable_services() {
    step "Enabling services (runit)"
    local TARGET="${TARGET:-/mnt}"
    local svdir="$TARGET/etc/runit/runsvdir/default"
    install -d -m 0755 "$svdir"

    local svc
    local enabled=(
        dbus
        NetworkManager
        polkitd
        docker
        acpid
        tlp
        elogind
        chronyd
        nix-daemon
        first-boot-nix
        zramen
        cupsd
        cups-browsed
    )

    # Display manager: greetd for wayland/niri, lightdm for cinnamon.
    if [[ "${DESKTOP:-cinnamon}" == "niri" ]]; then
        enabled+=(greetd bluetoothd)
    else
        enabled+=(lightdm bluetoothd)
    fi

    [[ "${SSHD_ENABLE:-no}" == "yes" ]] && enabled+=(sshd)

    for svc in "${enabled[@]}"; do
        if [[ -d "$TARGET/etc/sv/$svc" ]]; then
            ln -sf "/etc/sv/$svc" "$svdir/$svc"
            log "enabled $svc"
        else
            warn "no service dir /etc/sv/$svc — skipping"
        fi
    done

    # Disable dhcpcd if it's there (NetworkManager handles it).
    rm -f "$svdir/dhcpcd" 2>/dev/null || true

    ok "services enabled"
}

reconfigure_all() {
    step "Reconfiguring all packages (initramfs + grub artifacts)"
    run_chroot "xbps-reconfigure -fa"
    ok "reconfigure complete"
}