#!/usr/bin/env bash
# orioninstall — OrionOS installer (terminal + GUI-backend mode)
#
# Terminal mode:  sudo orioninstall
# GUI mode:       sudo orioninstall --gui /dev/sdX HOSTNAME USERNAME PASSWORD
#   Emits PROGRESS:N:message lines for the GUI to display.

set -euo pipefail

LOGFILE=/tmp/orioninstall.log
exec > >(tee -a "$LOGFILE") 2>&1

# ── mode detection ────────────────────────────────────────────────────
GUI_MODE=0
if [[ "${1:-}" == "--gui" ]]; then
    GUI_MODE=1
    DISK="${2:?missing disk argument}"
    HOSTNAME="${3:?missing hostname}"
    USERNAME="${4:?missing username}"
    PASSWORD="${5:?missing password}"
    TIMEZONE="${6:-UTC}"
    KEYBOARD="${7:-us}"
    shift 7 2>/dev/null || true
fi

# ── output helpers ────────────────────────────────────────────────────
RED='\033[0;31m'; YEL='\033[1;33m'; GRN='\033[0;32m'
BLU='\033[1;34m'; BOLD='\033[1m'; NC='\033[0m'

progress() { # N "message"
    if [[ $GUI_MODE -eq 1 ]]; then
        echo "PROGRESS:${1}:${2}"
    else
        echo -e "${BLU}[${1}%] ${2}${NC}"
    fi
}
info() {
    if [[ $GUI_MODE -eq 1 ]]; then echo ">>> $*";
    else echo -e "${BLU}>>> $*${NC}"; fi
}
ok()   {
    if [[ $GUI_MODE -eq 1 ]]; then echo "  ✓ $*";
    else echo -e "${GRN}    ✓ $*${NC}"; fi
}
warn() {
    if [[ $GUI_MODE -eq 1 ]]; then echo "  ! $*";
    else echo -e "${YEL}    ! $*${NC}"; fi
}
die()  {
    if [[ $GUI_MODE -eq 1 ]]; then echo "ERROR: $*";
    else echo -e "${RED}ERROR: $*${NC}"; fi
    exit 1
}

[[ $EUID -ne 0 ]] && die "Run with sudo: sudo orioninstall"

# ── terminal mode: interactive prompts ───────────────────────────────
if [[ $GUI_MODE -eq 0 ]]; then
    clear
    echo -e "${BOLD}"
    cat <<'BANNER'
 ╔════════════════════════════════════════════╗
 ║           OrionOS Installer                 ║
 ╚════════════════════════════════════════════╝
BANNER
    echo -e "${NC}"
    echo "This will install OrionOS to a disk on this computer."
    echo "The target disk will be ERASED."
    echo ""

    info "Available disks:"
    lsblk -d -o NAME,SIZE,MODEL --noheadings | grep -v "^loop" | \
        awk '{printf "  /dev/%-10s %6s   %s\n", $1, $2, $3}'
    echo ""

    read -rp "$(echo -e "${BOLD}Target disk (e.g. sda, nvme0n1): /dev/${NC}")" DISK_NAME
    DISK="/dev/${DISK_NAME}"
    [[ -b "$DISK" ]] || die "Not a block device: $DISK"

    echo ""
    echo -e "${RED}${BOLD}WARNING: ALL DATA ON $DISK WILL BE ERASED!${NC}"
    read -rp "Type YES to confirm: " CONFIRM
    [[ "$CONFIRM" == "YES" ]] || die "Aborted."

    echo ""
    read -rp "Hostname [skyos]: " HOSTNAME;   HOSTNAME="${HOSTNAME:-skyos}"
    read -rp "Username [sky]: "   USERNAME;   USERNAME="${USERNAME:-sky}"
    while true; do
        read -rsp "Password for $USERNAME: " PASSWORD;  echo
        read -rsp "Confirm password: "        PASSWORD2; echo
        [[ "$PASSWORD" == "$PASSWORD2" ]] && break
        warn "Passwords don't match, try again."
    done
fi

# ── detect boot mode ──────────────────────────────────────────────────
if [[ -d /sys/firmware/efi ]]; then
    BOOT_MODE="uefi"
else
    BOOT_MODE="bios"
fi

# ── partition names helper ────────────────────────────────────────────
part_name() { # disk suffix(1 or 2)
    local disk="$1" n="$2"
    if [[ "$disk" == *nvme* || "$disk" == *mmcblk* ]]; then
        echo "${disk}p${n}"
    else
        echo "${disk}${n}"
    fi
}

progress 2 "Detected boot mode: $BOOT_MODE"
info "Installing OrionOS to $DISK (boot: $BOOT_MODE)"

# ── 1. Partition ──────────────────────────────────────────────────────
progress 5 "Partitioning $DISK…"
info "Partitioning $DISK..."

if [[ "$BOOT_MODE" == "uefi" ]]; then
    parted -s "$DISK" \
        mklabel gpt \
        mkpart ESP  fat32 1MiB 513MiB \
        set 1 esp on \
        mkpart root ext4 513MiB 100%
    sleep 1; partprobe "$DISK" 2>/dev/null || true; sleep 1
    EFI_PART="$(part_name "$DISK" 1)"
    ROOT_PART="$(part_name "$DISK" 2)"
else
    parted -s "$DISK" \
        mklabel msdos \
        mkpart primary ext4 1MiB 100% \
        set 1 boot on
    sleep 1; partprobe "$DISK" 2>/dev/null || true; sleep 1
    ROOT_PART="$(part_name "$DISK" 1)"
    EFI_PART=""
fi
ok "Partitioned."

# ── 2. Format ─────────────────────────────────────────────────────────
progress 10 "Formatting partitions…"
info "Formatting..."
if [[ "$BOOT_MODE" == "uefi" ]]; then
    mkfs.fat -F32 -n ORIONEFI "$EFI_PART"
    ok "EFI: FAT32 ($EFI_PART)"
fi
mkfs.ext4 -L OrionOS "$ROOT_PART"
ok "Root: ext4 ($ROOT_PART)"

# ── 3. Mount ──────────────────────────────────────────────────────────
progress 14 "Mounting partitions…"
info "Mounting..."
mount "$ROOT_PART" /mnt
if [[ "$BOOT_MODE" == "uefi" ]]; then
    mkdir -p /mnt/boot
    mount "$EFI_PART" /mnt/boot
fi
ok "Mounted."

# ── 4. Copy live system ───────────────────────────────────────────────
progress 18 "Copying live system (this takes a few minutes)…"
info "Copying live system to disk..."

# On archiso live systems the squashfs is already mounted at
# /run/archiso/airootfs. Copy from there: it's the clean base system
# (no live-session overlay junk) and it's much faster and more reliable
# than rsyncing the overlayfs root.
SFS_ROOT="/run/archiso/airootfs"
if [[ -d "$SFS_ROOT" ]] && [[ -f "$SFS_ROOT/etc/os-release" ]]; then
    info "Copying from squashfs mount at $SFS_ROOT..."
    rsync -axHAXS "$SFS_ROOT/" /mnt/
else
    info "Squashfs not found; falling back to live root copy..."
    rsync -axHAXS \
        --exclude=/proc/ \
        --exclude=/sys/ \
        --exclude=/dev/ \
        --exclude=/run/ \
        --exclude=/mnt/ \
        --exclude=/media/ \
        --exclude=/tmp/ \
        --exclude=/lost+found \
        / /mnt/
fi
ok "Files copied."

# ── 5. fstab ─────────────────────────────────────────────────────────
progress 75 "Writing fstab…"
info "Writing fstab..."
mkdir -p /mnt/etc
genfstab -U /mnt > /mnt/etc/fstab
ok "fstab written."

# ── 6. Chroot configuration ───────────────────────────────────────────
progress 80 "Configuring installed system…"
info "Configuring system..."

arch-chroot /mnt /bin/bash -s "$HOSTNAME" "$USERNAME" "$PASSWORD" \
                                "$BOOT_MODE" "$ROOT_PART" \
                                "${TIMEZONE:-UTC}" "${KEYBOARD:-us}" \
    <<'CHROOT'
HOSTNAME="$1"; USERNAME="$2"; PASSWORD="$3"
BOOT_MODE="$4"; ROOT_PART="$5"
TIMEZONE="${6:-UTC}"; KEYBOARD="${7:-us}"

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

# Locale
sed -i 's/#en_US.UTF-8 UTF-8/en_US.UTF-8 UTF-8/' /etc/locale.gen 2>/dev/null || true
locale-gen
echo "LANG=en_US.UTF-8" > /etc/locale.conf

# Console keyboard layout
echo "KEYMAP=${KEYBOARD:-us}" > /etc/vconsole.conf

# X11 keyboard layout
mkdir -p /etc/X11/xorg.conf.d
cat > /etc/X11/xorg.conf.d/10-keyboard.conf <<KBEOF
Section "InputClass"
    Identifier "All Keyboards"
    MatchIsKeyboard "on"
    Driver "libinput"
    Option "XkbLayout"  "${KEYBOARD:-us}"
    Option "XkbOptions" "terminate:ctrl_alt_bksp"
EndSection
KBEOF

# Timezone
ln -sf "/usr/share/zoneinfo/${TIMEZONE:-UTC}" /etc/localtime
hwclock --systohc 2>/dev/null || true

# Remove live-only archiso package to prevent mkinitcpio confusion
pacman -Rns --noconfirm archiso 2>/dev/null || true

# Create user
if ! id "$USERNAME" &>/dev/null; then
    useradd -m -G wheel,audio,video,storage,input,network -s /bin/bash "$USERNAME"
fi
echo "${USERNAME}:${PASSWORD}" | chpasswd
echo "root:${PASSWORD}" | chpasswd

# sudoers
echo "%wheel ALL=(ALL) NOPASSWD: ALL" > /etc/sudoers.d/wheel
chmod 440 /etc/sudoers.d/wheel

# Remove live autologin (installed system should prompt login)
rm -f /etc/systemd/system/getty@tty1.service.d/autologin.conf

# Auto-startx on tty1
HOME_DIR="/home/$USERNAME"
cat > "$HOME_DIR/.bash_profile" <<'BP'
[[ -z $DISPLAY && $XDG_VTNR -eq 1 ]] && exec startx
BP
chown "$USERNAME:$USERNAME" "$HOME_DIR/.bash_profile"

# Rebuild initramfs with standard hooks (no archiso hooks on installed system)
cat > /etc/mkinitcpio.conf <<'MKCONF'
MODULES=()
BINARIES=()
FILES=()
HOOKS=(base udev autodetect microcode modconf kms keyboard keymap consolefont block filesystems fsck)
MKCONF
mkinitcpio -P

# Bootloader
if [[ "$BOOT_MODE" == "uefi" ]]; then
    bootctl --esp-path=/boot install 2>/dev/null || true
    ROOT_UUID="$(blkid -s UUID -o value "$ROOT_PART")"
    mkdir -p /boot/loader/entries
    cat > /boot/loader/loader.conf <<'LC'
default orionos.conf
timeout 5
editor  no
LC
    # Build initrd lines — only include microcode if the image exists
    INITRD_LINES="initrd  /initramfs-linux.img"
    [[ -f /boot/amd-ucode.img   ]] && INITRD_LINES="initrd  /amd-ucode.img\n${INITRD_LINES}"
    [[ -f /boot/intel-ucode.img ]] && INITRD_LINES="initrd  /intel-ucode.img\n${INITRD_LINES}"
    printf "title   OrionOS\nlinux   /vmlinuz-linux\n%b\noptions root=UUID=%s rw quiet loglevel=3\n" \
        "$INITRD_LINES" "$ROOT_UUID" > /boot/loader/entries/orionos.conf
else
    # Find the whole disk (strip trailing partition number)
    DISK="${ROOT_PART%[0-9]}"
    # Handle nvme/mmcblk: strip trailing 'p' + digit
    [[ "$DISK" == *p ]] && DISK="${DISK%p}"
    grub-install --target=i386-pc "$DISK" || true
    grub-mkconfig -o /boot/grub/grub.cfg
fi

systemctl enable NetworkManager 2>/dev/null || true
systemctl enable bluetooth 2>/dev/null || true

echo "Chroot configuration complete."
CHROOT

ok "System configured."

# ── 7. Unmount ────────────────────────────────────────────────────────
progress 97 "Unmounting…"
info "Unmounting..."
sync
[[ "$BOOT_MODE" == "uefi" ]] && umount /mnt/boot 2>/dev/null || true
umount /mnt
ok "Done."

progress 100 "Installation complete!"

if [[ $GUI_MODE -eq 0 ]]; then
    echo ""
    echo -e "${GRN}${BOLD}"
    cat <<'DONE'
 ╔════════════════════════════════════════════╗
 ║   OrionOS installed successfully!            ║
 ║   Remove the USB drive and reboot.         ║
 ╚════════════════════════════════════════════╝
DONE
    echo -e "${NC}"
    echo "Log saved to: $LOGFILE"
    echo ""
    read -rp "Reboot now? [y/N] " REBOOT
    [[ "$REBOOT" =~ ^[Yy]$ ]] && reboot
fi