#!/bin/bash

# The configurator normally runs under this small PTY proxy so Ctrl-T can work
# as a global tools shortcut even while gum owns the terminal in raw mode. The
# child configurator keeps stdout on the existing stream (so generated config and
# logs behave normally), while stderr and /dev/tty traffic are relayed through the
# PTY. Pressing Ctrl-T kills the current configurator screen, opens Tools, then
# starts the configurator over at keyboard selection.
if [[ -z "${OMARCHY_CONFIGURATOR_INNER:-}" && -z "${OMARCHY_CONFIGURATOR_NO_KEY_PROXY:-}" ]] && { [[ -t 0 ]] || [[ -t 2 ]]; }; then
  OMARCHY_CONFIGURATOR_PYTHON=$(command -v python3 2>/dev/null || command -v python 2>/dev/null || true)
  if [[ -n "$OMARCHY_CONFIGURATOR_PYTHON" ]]; then
    export OMARCHY_CONFIGURATOR_SCRIPT
    OMARCHY_CONFIGURATOR_SCRIPT=$(readlink -f "${BASH_SOURCE[0]}" 2>/dev/null || printf '%s' "${BASH_SOURCE[0]}")
    exec "$OMARCHY_CONFIGURATOR_PYTHON" - "$@" <<'PY'
import errno
import fcntl
import os
import pty
import select
import signal
import sys
import termios
import time

CTRL_T = b"\x14"
SCRIPT = os.environ["OMARCHY_CONFIGURATOR_SCRIPT"]
ARGS = sys.argv[1:]

try:
    tty_fd = os.open("/dev/tty", os.O_RDWR | os.O_NOCTTY)
except OSError:
    env = os.environ.copy()
    env["OMARCHY_CONFIGURATOR_INNER"] = "1"
    os.execvpe("bash", ["bash", SCRIPT] + ARGS, env)

stdout_fd = os.dup(1)
original_tty = termios.tcgetattr(tty_fd)
current_master = None
current_pgid = None


def restore_tty():
    try:
        termios.tcsetattr(tty_fd, termios.TCSADRAIN, original_tty)
    except termios.error:
        pass


def write_tty(data):
    if data:
        os.write(tty_fd, data)


def copy_window_size(master):
    try:
        size = fcntl.ioctl(tty_fd, termios.TIOCGWINSZ, b"\0" * 8)
        fcntl.ioctl(master, termios.TIOCSWINSZ, size)
    except OSError:
        pass


def on_winch(signum, frame):
    if current_master is not None:
        copy_window_size(current_master)
    if current_pgid is not None:
        try:
            os.killpg(current_pgid, signal.SIGWINCH)
        except OSError:
            pass


def on_termination(signum, frame):
    if current_pgid is not None:
        try:
            os.killpg(current_pgid, signum)
        except OSError:
            pass
    restore_tty()
    raise SystemExit(128 + signum)


def exit_code(status):
    if os.WIFEXITED(status):
        return os.WEXITSTATUS(status)
    if os.WIFSIGNALED(status):
        return 128 + os.WTERMSIG(status)
    return 1


def spawn_child(tools_only=False):
    global current_master, current_pgid

    pid, master = pty.fork()
    if pid == 0:
        os.dup2(stdout_fd, 1)
        try:
            os.close(stdout_fd)
        except OSError:
            pass

        env = os.environ.copy()
        env["OMARCHY_CONFIGURATOR_INNER"] = "1"
        env["OMARCHY_CONFIGURATOR_CTRL_T"] = "1"
        if tools_only:
            env["OMARCHY_CONFIGURATOR_TOOLS_ONLY"] = "1"
        else:
            env.pop("OMARCHY_CONFIGURATOR_TOOLS_ONLY", None)

        os.execvpe("bash", ["bash", SCRIPT] + ARGS, env)

    os.set_blocking(master, False)
    current_master = master
    current_pgid = pid
    copy_window_size(master)
    return pid, master


def wait_for_child(pid, timeout):
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            done, status = os.waitpid(pid, os.WNOHANG)
        except ChildProcessError:
            return 0
        if done == pid:
            return exit_code(status)
        time.sleep(0.05)
    return None


def terminate_child(pid):
    for sig, timeout in ((signal.SIGTERM, 0.5), (signal.SIGKILL, 0.2)):
        try:
            os.killpg(pid, sig)
        except ProcessLookupError:
            return 0
        except OSError:
            try:
                os.kill(pid, sig)
            except OSError:
                return 0

        code = wait_for_child(pid, timeout)
        if code is not None:
            return code
    return 130


def drain_master(master):
    while True:
        try:
            data = os.read(master, 4096)
        except BlockingIOError:
            return
        except OSError:
            return
        if not data:
            return
        write_tty(data)


def relay_child(tools_only=False):
    global current_master, current_pgid

    pid, master = spawn_child(tools_only)
    try:
        while True:
            try:
                done, status = os.waitpid(pid, os.WNOHANG)
            except ChildProcessError:
                drain_master(master)
                return 0
            if done == pid:
                drain_master(master)
                return exit_code(status)

            try:
                readable, _, _ = select.select([tty_fd, master], [], [], 0.1)
            except InterruptedError:
                continue

            if tty_fd in readable:
                try:
                    data = os.read(tty_fd, 1024)
                except OSError:
                    data = b""
                if not data:
                    continue

                if not tools_only and CTRL_T in data:
                    before = data.split(CTRL_T, 1)[0]
                    if before:
                        try:
                            os.write(master, before)
                        except OSError:
                            pass
                    terminate_child(pid)
                    drain_master(master)
                    return "tools"

                try:
                    os.write(master, data)
                except OSError:
                    pass

            if master in readable:
                try:
                    data = os.read(master, 4096)
                except BlockingIOError:
                    continue
                except OSError as exc:
                    if exc.errno != errno.EIO:
                        raise
                    data = b""
                if data:
                    write_tty(data)
    finally:
        current_master = None
        current_pgid = None
        try:
            os.close(master)
        except OSError:
            pass


signal.signal(signal.SIGWINCH, on_winch)
signal.signal(signal.SIGHUP, on_termination)
signal.signal(signal.SIGTERM, on_termination)

# Use raw-ish input so Ctrl-T/Ctrl-C can be intercepted and relayed, but keep
# the real terminal's output processing intact. Full raw mode disables ONLCR,
# causing gum/logo newlines to stair-step across the screen.
tty_attrs = termios.tcgetattr(tty_fd)
tty_attrs[0] &= ~(termios.BRKINT | termios.ICRNL | termios.INPCK | termios.ISTRIP | termios.IXON)
tty_attrs[2] |= termios.CS8
tty_attrs[3] &= ~(termios.ECHO | termios.ICANON | termios.IEXTEN | termios.ISIG)
tty_attrs[6][termios.VMIN] = 1
tty_attrs[6][termios.VTIME] = 0
termios.tcsetattr(tty_fd, termios.TCSADRAIN, tty_attrs)

try:
    while True:
        result = relay_child(tools_only=False)
        if result == "tools":
            relay_child(tools_only=True)
            continue
        raise SystemExit(result)
finally:
    restore_tty()
    try:
        os.close(stdout_fd)
        os.close(tty_fd)
    except OSError:
        pass
PY
  fi
fi
unset OMARCHY_CONFIGURATOR_PYTHON

# Minimal live-ISO presentation helpers. Keep these local to the ISO: the
# installed Omarchy package owns target setup helpers, while the ISO owns the
# installer UI that runs before any target packages exist.
OMARCHY_PATH="${OMARCHY_PATH:-/usr/share/omarchy}"
LOGO_PATH="${LOGO_PATH:-$OMARCHY_PATH/logo.txt}"

TERM_HEIGHT=${LINES:-24}
TERM_WIDTH=${COLUMNS:-80}
if [[ -e /dev/tty ]]; then
  TERM_SIZE=$(stty size 2>/dev/null </dev/tty || true)
  if [[ $TERM_SIZE =~ ^([0-9]+)[[:space:]]+([0-9]+)$ ]] && (( BASH_REMATCH[1] > 0 && BASH_REMATCH[2] > 0 )); then
    TERM_HEIGHT=${BASH_REMATCH[1]}
    TERM_WIDTH=${BASH_REMATCH[2]}
  fi
fi

if [[ -f $LOGO_PATH ]]; then
  LOGO_WIDTH=$(awk '{ if (length > max) max = length } END { print max+0 }' "$LOGO_PATH" 2>/dev/null || echo 0)
else
  LOGO_WIDTH=7
fi
(( LOGO_WIDTH > 0 )) || LOGO_WIDTH=7
PADDING_LEFT=$(((TERM_WIDTH - LOGO_WIDTH) / 2))
(( PADDING_LEFT < 0 )) && PADDING_LEFT=0
PADDING_LEFT_SPACES=$(printf "%*s" "$PADDING_LEFT" "")
PADDING="0 0 0 $PADDING_LEFT"

export TERM_HEIGHT TERM_WIDTH LOGO_PATH PADDING_LEFT PADDING_LEFT_SPACES PADDING
export GUM_CONFIRM_PROMPT_FOREGROUND="6"
export GUM_CONFIRM_SELECTED_FOREGROUND="0"
export GUM_CONFIRM_SELECTED_BACKGROUND="2"
export GUM_CONFIRM_UNSELECTED_FOREGROUND="7"
export GUM_CONFIRM_UNSELECTED_BACKGROUND="0"
export GUM_CHOOSE_PADDING="$PADDING"
export GUM_FILTER_PADDING="$PADDING"
export GUM_INPUT_PADDING="$PADDING"
export GUM_SPIN_PADDING="$PADDING"
export GUM_TABLE_PADDING="$PADDING"
export GUM_CONFIRM_PADDING="$PADDING"

clear_logo() {
  printf "\033[H\033[2J"
  if [[ -f $LOGO_PATH ]]; then
    gum style --foreground 2 --padding "1 0 0 $PADDING_LEFT" "$(<"$LOGO_PATH")"
  else
    gum style --foreground 2 --padding "1 0 0 $PADDING_LEFT" "Omarchy"
  fi
}

# Ensure tzupdate is available for guessing timezone
command -v tzupdate >/dev/null || sudo pacman -Sy --noconfirm tzupdate

abort() {
  gum style "${1:-Aborted installation}"
  echo
  gum style "You can retry later by running: ./.automated_script.sh"
  exit 1
}

tools_hint() {
  [[ -n "${OMARCHY_CONFIGURATOR_CTRL_T:-}" && -z "${OMARCHY_CONFIGURATOR_TOOLS_ONLY:-}" ]] || return 0

  # Draw the shortcut as a fixed footer instead of in the form body. Keep the
  # cursor where the active gum widget expects it; do not print a newline from
  # the bottom row or the VT will scroll.
  printf "\0337\033[%s;1H\033[2K%s\033[38;5;8mCtrl-T  Tools\033[0m\0338" "$TERM_HEIGHT" "$PADDING_LEFT_SPACES"
}

step() {
  clear_logo
  echo
  gum style --padding "0 0 0 $PADDING_LEFT" "$1"
  tools_hint
  echo
}

# Padded body text to line up with the logo and gum widgets.
say() {
  gum style --padding "0 0 0 $PADDING_LEFT" "$@"
}

notice() {
  clear_logo
  echo
  gum spin --spinner "pulse" --title "$1" -- sleep "${2:-2}"
  echo
}

# Tools menu — recovery-style actions kept off the main install path.
# Inlines its own disk picker so Ctrl-T can open it before disk_form is defined
# further down in the script.
tools_menu() {
  while true; do
    step "Tools"
    local choice
    choice=$(gum choose --header "Pick a tool" \
      "Open partition tool" \
      "Drop to shell" \
      "Reboot" \
      "Power off" \
      "Back to installer") || return 0

    case "$choice" in
      "Open partition tool")
        local d disks
        disks=$(lsblk -dpno NAME,TYPE 2>/dev/null | awk '$2=="disk"{print $1}' | grep -E '/dev/(sd|hd|vd|nvme|mmcblk|xv)' || true)
        if [[ -z "$disks" ]]; then
          notice "No suitable disks found" 2
          continue
        fi
        d=$(echo "$disks" | gum choose --header "Select disk") || continue
        clear
        cfdisk "$d" || true
        partprobe "$d" 2>/dev/null || true
        sync
        sleep 1
        ;;
      "Drop to shell")
        clear
        echo "Type 'exit' to return to the tools menu."
        bash --norc || true
        ;;
      "Reboot")
        reboot
        ;;
      "Power off")
        poweroff
        ;;
      "Back to installer")
        return 0
        ;;
    esac
  done
}

if [[ -n "${OMARCHY_CONFIGURATOR_TOOLS_ONLY:-}" ]]; then
  tools_menu
  exit 0
fi

# STEP 1: KEYBOARD LAYOUT

keyboard_form() {
  step "Let's setup your machine..."
  keyboards=$'Azerbaijani|azerty
Belarusian|by
Belgian|be-latin1
Bosnian|ba
Bulgarian|bg-cp1251
Croatian|croat
Czech|cz
Danish|dk-latin1
Dutch|nl
English (UK)|uk
English (US)|us
English (US, Dvorak)|dvorak
English (US, Colemak)|colemak
Estonian|et
Finnish|fi
French|fr
French (Canada)|cf
French (Switzerland)|fr_CH
Georgian|ge
German|de
German (Switzerland)|de_CH-latin1
Greek|gr
Hebrew|il
Hungarian|hu
Icelandic|is-latin1
Irish|ie
Italian|it
Japanese|jp106
Kazakh|kazakh
Khmer (Cambodia)|khmer
Kyrgyz|kyrgyz
Lao|la-latin1
Latvian|lv
Lithuanian|lt
Macedonian|mk-utf
Norwegian|no-latin1
Polish|pl
Portuguese|pt-latin1
Portuguese (Brazil)|br-abnt2
Romanian|ro
Russian|ru
Serbian|sr-latin
Slovak|sk-qwertz
Slovenian|slovene
Spanish|es
Spanish (Latin American)|la-latin1
Swedish|sv-latin1
Tajik|tj_alt-UTF8
Turkish|trq
Ukrainian|ua'
  choice=$(printf '%s\n' "$keyboards" | cut -d'|' -f1 | gum choose --height 10 --selected "English (US)" --header "Select keyboard layout") || abort
  keyboard=$(printf '%s\n' "$keyboards" | awk -F'|' -v c="$choice" '$1==c{print $2; exit}')

  # Only attempt to load keyboard layout if we're on a real console
  # loadkeys only works on Linux virtual consoles (tty), not in terminal emulators
  if [[ $(tty 2>/dev/null) == "/dev/tty"* ]]; then
    loadkeys "$keyboard" 2>/dev/null
  fi
}

keyboard_form

# STEP 2: USER

user_form() {
  step "Let's setup your user account..."

  while true; do
    username=$(gum input --placeholder "Alphanumeric without spaces (like dhh)" --prompt.foreground="#845DF9" --prompt "Username> ") || abort

    if [[ "$username" =~ ^[a-z_][a-z0-9_-]*[$]?$ ]]; then
      break
    else
      notice "Username must be alphanumeric with no spaces" 1
    fi
  done

  while true; do
    password=$(gum input --placeholder "Used for user + root, and disk encryption when enabled" --prompt.foreground="#845DF9" --password --prompt "Password> ") || abort
    password_confirmation=$(gum input --placeholder "Must match the password you just typed" --prompt.foreground="#845DF9" --password --prompt "Confirm> ") || abort

    if [[ -n "$password" && "$password" == "$password_confirmation" ]]; then
      break
    elif [[ -z "$password" ]]; then
      notice "Your password can't be blank!" 1
    else
      notice "Passwords didn't match!" 1
    fi
  done

  # Hash the password using yescrypt
  password_hash=$(printf '%s' "$password" | openssl passwd -6 -stdin)

  full_name=$(gum input --placeholder "Used for git authentication (hit return to skip)" --prompt.foreground="#845DF9" --prompt "Full name> ")
  email_address=$(gum input --placeholder "Used for git authentication (hit return to skip)" --prompt.foreground="#845DF9" --prompt "Email address> ")

  while true; do
    hostname=$(gum input --placeholder "Letters, digits, and dashes (or return for 'omarchy')" --prompt.foreground="#845DF9" --prompt "Hostname> ")

    if [[ "$hostname" =~ ^[A-Za-z0-9]([A-Za-z0-9-]{0,61}[A-Za-z0-9])?$ ]]; then
      break
    elif [[ -z $hostname ]]; then
      hostname="omarchy"
      break
    else
      notice "Hostname must be 1-63 letters, digits, or dashes, and cannot start or end with a dash" 1
    fi
  done

  # Pick timezone
  geo_guessed_timezone=$(tzupdate -p 2>/dev/null)

  if [[ -n $geo_guessed_timezone ]]; then
    timezone=$(timedatectl list-timezones | gum choose --height 10 --selected "$geo_guessed_timezone" --header "Timezone") || abort
  else
    timezone=$(timedatectl list-timezones | gum filter --height 10 --header "Timezone") || abort
  fi
}

user_form

while true; do
  # Add manual padding since gum table -p doesn't respect padding
  echo -e "Field,Value
Username,$username
Password,$(printf "%${#password}s" | tr ' ' '*')
Full name,${full_name:-[Skipped]}
Email address,${email_address:-[Skipped]}
Hostname,$hostname
Timezone,$timezone
Keyboard,$keyboard" |
    gum table -s "," -p | sed "s/^/${PADDING_LEFT_SPACES}/"

  echo
  if gum confirm --negative "No, change it" "Does this look right?"; then
    break
  else
    keyboard_form
    user_form
  fi
done

# STEP 3: DISK

get_disk_info() {
  local device="$1"
  local size vendor model label

  # Get basic disk information
  size=$(lsblk -dno SIZE "$device" 2>/dev/null)
  vendor=$(lsblk -dno VENDOR "$device" 2>/dev/null | sed 's/ *$//')
  model=$(lsblk -dno MODEL "$device" 2>/dev/null | sed 's/ *$//')

  # Combine vendor and model, avoiding duplication
  label=""
  if [[ -n $vendor && -n $model ]]; then
    if [[ $model == *$vendor* ]]; then
      label="$model"
    else
      label="$vendor $model"
    fi
  elif [[ -n $model ]]; then
    label="$model"
  elif [[ -n $vendor ]]; then
    label="$vendor"
  fi

  # Format display string
  local display="$device"
  [[ -n $size ]] && display="$display ($size)"
  [[ -n $label ]] && display="$display - $label"

  # Append compact partition summary
  local part_summary
  part_summary=$(lsblk -nro TYPE,NAME,FSTYPE,MOUNTPOINT "$device" 2>/dev/null | \
    awk '$1=="part" { printf "%s%s%s", s, ($3==""?"unknown":$3), ($4==""?"":"("$4")"); s=", " }')
  [[ -n $part_summary ]] && display+=" [$part_summary]"

  echo "$display"
}

get_root_disk() {
  local device="$1"
  local parent

  [[ -n $device ]] || return 1

  # /run/archiso/bootmnt usually resolves to the boot partition, but the
  # installer picker works with whole disks. Walk back to the parent disk so
  # the install media itself never shows up as a wipe target.
  device=$(readlink -f "$device" 2>/dev/null || printf "%s\n" "$device")

  while true; do
    parent=$(lsblk -dno PKNAME "$device" 2>/dev/null | tail -n1)
    [[ -n $parent ]] || break
    device="/dev/$parent"
  done

  if [[ $(lsblk -dno TYPE "$device" 2>/dev/null) == "disk" ]]; then
    printf "%s\n" "$device"
  fi
}

# Compute partition device path, handling NVMe/mmcblk's pN naming.
partition_path() {
  local _disk="$1" _num="$2"
  if [[ "$_disk" == *nvme* || "$_disk" == *mmcblk* ]]; then
    echo "${_disk}p${_num}"
  else
    echo "${_disk}${_num}"
  fi
}

# BitLocker writes the "-FVE-FS-" signature at offset 3 of the partition.
detect_bitlocker() {
  local p sig
  for p in $(lsblk -lnpo NAME,TYPE "$disk" | awk '$2=="part"{print $1}'); do
    sig=$(dd if="$p" bs=1 skip=3 count=8 status=none 2>/dev/null)
    if [[ "$sig" == "-FVE-FS-" ]]; then
      echo "$p"
      return 0
    fi
  done
  return 1
}

# Look for an existing Windows ESP (FAT partition containing EFI/Microsoft) on $disk.
detect_windows_esp() {
  local p tmp_mp
  for p in $(blkid -t TYPE=vfat -o device 2>/dev/null); do
    [[ "$p" == "$disk"* ]] || continue
    tmp_mp=$(mktemp -d)
    if mount -o ro "$p" "$tmp_mp" 2>/dev/null; then
      if [[ -d "$tmp_mp/EFI/Microsoft" ]]; then
        umount "$tmp_mp"; rmdir "$tmp_mp"
        echo "$p"
        return 0
      fi
      umount "$tmp_mp"
    fi
    rmdir "$tmp_mp" 2>/dev/null || true
  done
  return 1
}

to_gb() {
  awk -v bytes="$1" 'BEGIN {printf "%.1fGB", bytes/1024/1024/1024}'
}

# Shared error screen for every "can't fit Omarchy here" outcome of the
# partition install. Offers the partition tool as the meaningful second
# action so users have a way to actually fix the situation.
not_enough_space() {
  local available="${1:-0}"
  step "Not enough free space on $disk"
  say --foreground 1 "$disk has $(to_gb $available) of usable free space; Omarchy needs at least 32GB."
  say "Open the partition tool to free at least 32GB on $disk, then try again."
  echo
  if ! gum confirm --affirmative "Back" --negative "Open partition tool" "Return to installation mode?"; then
    open_partition_tool
  fi
  return 1
}

run_partition_install() {
  local dry=false
  [[ "${1:-}" == "dry" ]] && dry=true

  step "Checking for BitLocker on $disk"
  local bl_part
  bl_part=$(detect_bitlocker || true)
  if [[ -n "$bl_part" ]]; then
    say --foreground 1 "BitLocker signature detected on $bl_part."
    say "Suspend or disable BitLocker in Windows before installing alongside it."
    abort "Aborted: BitLocker is active on this disk."
  fi

  step "Looking for an existing Windows ESP on $disk"
  local shared_esp
  shared_esp=$(detect_windows_esp || true)

  if [[ -n "$shared_esp" ]]; then
    say "Found Windows ESP at $shared_esp — will install alongside."
    local tmp_mp esp_free_mib
    tmp_mp=$(mktemp -d)
    mount -o ro "$shared_esp" "$tmp_mp"
    esp_free_mib=$(df -BM --output=avail "$tmp_mp" | tail -1 | tr -dc '0-9')
    umount "$tmp_mp"; rmdir "$tmp_mp"
    if (( esp_free_mib < 100 )); then
      abort "Windows ESP $shared_esp has only ${esp_free_mib}MiB free — need at least 100MiB for Omarchy."
    fi
    say "Windows ESP has ${esp_free_mib}MiB free."
  else
    say "No Windows ESP detected on $disk — a new EFI partition will be created in free space."
  fi

  step "Analyzing free space on $disk"
  partprobe "$disk" 2>/dev/null || true
  sleep 1

  # An unlabeled disk has no partition table for parted to scan, so we
  # synthesize a whole-disk free region and mklabel gpt before mkpart.
  local FREE_SPACE_INFO needs_mklabel=false pt_type disk_size_b
  pt_type=$(lsblk -dno PTTYPE "$disk" 2>/dev/null)
  if [[ -z "$pt_type" ]]; then
    disk_size_b=$(lsblk -bdno SIZE "$disk" 2>/dev/null)
    if [[ -n "$disk_size_b" && "$disk_size_b" -gt 0 ]]; then
      local _mib=$((1024 * 1024))
      local _start=$_mib
      local _end=$(( (disk_size_b - _mib) / _mib * _mib ))
      FREE_SPACE_INFO="$_start $_end $((_end - _start))"
      needs_mklabel=true
    fi
  else
    # -s (script mode) keeps parted from prompting on stderr for warnings
    # like a stale GPT backup; without it the configurator hangs invisibly.
    FREE_SPACE_INFO=$(parted -ms "$disk" unit B print free </dev/null 2>/dev/null | \
      awk -F: '
        $NF ~ /free/ {
          start=$2; end=$3; size=$4;
          gsub(/B/, "", start);
          gsub(/B/, "", end);
          gsub(/B/, "", size);
          if (size+0 > max_size+0) {
            max_size=size; max_start=start; max_end=end;
          }
        }
        END {
          if (max_size+0 > 0) printf "%s %s %s\n", max_start, max_end, max_size;
        }')
  fi

  if [[ -z "$FREE_SPACE_INFO" ]]; then
    not_enough_space 0
    return $?
  fi

  local FREE_START_B FREE_END_B FREE_SIZE_B
  FREE_START_B=$(echo "$FREE_SPACE_INFO" | awk '{print $1}')
  FREE_END_B=$(echo "$FREE_SPACE_INFO"   | awk '{print $2}')
  FREE_SIZE_B=$(echo "$FREE_SPACE_INFO"  | awk '{print $3}')

  local mib=$((1024 * 1024)) gib=$((1024 * 1024 * 1024))
  local ALIGN=$mib
  local EFI_SIZE_B=$((2 * gib))
  local MIN_INSTALL_B=$((32 * gib))    # Minimum total Omarchy allocation

  local EFI_START_B=0 EFI_END_B=0 ROOT_START_B ROOT_END_B
  if [[ -z "$shared_esp" ]]; then
    EFI_START_B=$(( (FREE_START_B + ALIGN - 1) / ALIGN * ALIGN ))
    EFI_END_B=$((EFI_START_B + EFI_SIZE_B))
    ROOT_START_B=$(( (EFI_END_B + 1 + ALIGN - 1) / ALIGN * ALIGN ))
  else
    ROOT_START_B=$(( (FREE_START_B + ALIGN - 1) / ALIGN * ALIGN ))
  fi
  ROOT_END_B=$((FREE_END_B / ALIGN * ALIGN))

  if (( ROOT_END_B <= ROOT_START_B )); then
    not_enough_space "$FREE_SIZE_B"
    return $?
  fi

  local INSTALL_START_B INSTALL_MAX_B INSTALL_MAX_GIB INSTALL_SIZE_GIB
  if [[ -z "$shared_esp" ]]; then
    INSTALL_START_B=$EFI_START_B
  else
    INSTALL_START_B=$ROOT_START_B
  fi
  INSTALL_MAX_B=$((ROOT_END_B - INSTALL_START_B))

  if (( INSTALL_MAX_B < MIN_INSTALL_B )); then
    not_enough_space "$INSTALL_MAX_B"
    return $?
  fi

  INSTALL_MAX_GIB=$((INSTALL_MAX_B / gib))
  while true; do
    step "Choose Omarchy size"
    say "Free region: $(to_gb $FREE_START_B) — $(to_gb $FREE_END_B)  ($(to_gb $FREE_SIZE_B))"
    if [[ -z "$shared_esp" ]]; then
      say "The size includes a new 2GiB EFI partition; the rest becomes root."
    else
      say "Using existing ESP: $shared_esp"
    fi
    say "Omarchy can use 32–${INSTALL_MAX_GIB}GiB total. Leave the default to use all available space."
    echo

    INSTALL_SIZE_GIB=$(gum input \
      --value "$INSTALL_MAX_GIB" \
      --placeholder "32-${INSTALL_MAX_GIB}" \
      --prompt.foreground="#845DF9" \
      --prompt "Omarchy size in GiB> ") || return 1

    if [[ $INSTALL_SIZE_GIB =~ ^[0-9]+$ ]] && (( INSTALL_SIZE_GIB >= 32 && INSTALL_SIZE_GIB <= INSTALL_MAX_GIB )); then
      break
    fi

    notice "Enter a whole number from 32 to ${INSTALL_MAX_GIB}" 1
  done

  if [[ -z "$shared_esp" ]]; then
    ROOT_END_B=$((EFI_START_B + INSTALL_SIZE_GIB * gib))
  else
    ROOT_END_B=$((ROOT_START_B + INSTALL_SIZE_GIB * gib))
  fi

  # Shared Windows ESP can't host our LUKS layout, so encryption isn't
  # offered in that case. Otherwise default to encrypted and hide the
  # unencrypted path behind Ctrl+C, matching the full-disk flow.
  local mode="encrypted" affirmative confirm_status
  [[ -n "$shared_esp" ]] && mode="forced_unencrypted"

  while true; do
    step "Proposed layout on $disk"
    say "Free region: $(to_gb $FREE_START_B) — $(to_gb $FREE_END_B)  ($(to_gb $FREE_SIZE_B))"
    if [[ -z "$shared_esp" ]]; then
      say "  New EFI partition:  $(to_gb $EFI_START_B) — $(to_gb $EFI_END_B) (2GB FAT32)"
    else
      say "  Shared ESP:        $shared_esp (Windows, mounted at /efi)"
    fi
    say "  New root (Btrfs):  $(to_gb $ROOT_START_B) — $(to_gb $ROOT_END_B) ($(to_gb $((ROOT_END_B - ROOT_START_B))))"
    echo

    case $mode in
      encrypted)
        say --foreground 8 "Press Ctrl+C for unencrypted install."
        affirmative="Yes, install"
        ;;
      unencrypted)
        affirmative="Yes, install without encryption"
        ;;
      forced_unencrypted)
        say --foreground 8 "Shared Windows ESP cannot host encrypted Omarchy — installing unencrypted."
        affirmative="Yes, install"
        ;;
    esac
    echo

    gum confirm --affirmative "$affirmative" --negative "Cancel" \
      "Create Omarchy partitions in the free space on $disk?"
    confirm_status=$?

    case $confirm_status in
      0)
        [[ $mode == "encrypted" ]] && encrypt_installation=true || encrypt_installation=false
        break
        ;;
      1)
        return 1
        ;;
      130)
        case $mode in
          encrypted) mode="unencrypted" ;;
          unencrypted) mode="encrypted" ;;
        esac
        ;;
      *)
        abort
        ;;
    esac
  done

  local esp_mount_in_target
  if [[ "$encrypt_installation" == "true" ]]; then
    esp_mount_in_target=/boot
  else
    esp_mount_in_target=/efi
  fi

  # Detect T2 Mac and set appropriate kernel
  if lspci -nn 2>/dev/null | grep -q "106b:180[12]"; then
    kernel_choice="linux-t2"
  else
    kernel_choice="linux"
  fi

  local last_part efi_part_num root_part_num
  last_part=$(lsblk -n -o PARTN "$disk" 2>/dev/null | grep -E '^[0-9]+$' | sort -n | tail -1)
  if [[ -z "$shared_esp" ]]; then
    efi_part_num=$(( ${last_part:-0} + 1 ))
    root_part_num=$(( efi_part_num + 1 ))
  else
    efi_part_num=0   # not created by us
    root_part_num=$(( ${last_part:-0} + 1 ))
  fi

  local efi_dev root_partition_device root_mapper
  if [[ -n "$shared_esp" ]]; then
    efi_dev="$shared_esp"
  else
    efi_dev=$(partition_path "$disk" "$efi_part_num")
  fi
  root_partition_device=$(partition_path "$disk" "$root_part_num")
  if [[ "$encrypt_installation" == "true" ]]; then
    root_mapper="/dev/mapper/omarchy_root"
  else
    root_mapper="$root_partition_device"
  fi

  if $dry; then
    $needs_mklabel && say "[dry] would initialize GPT on $disk"
    say "[dry] would create root partition at $root_partition_device"
    [[ -z "$shared_esp" ]] && say "[dry] would create EFI partition at $efi_dev"
    say "[dry] encrypted=$encrypt_installation kernel=$kernel_choice"
  else
    step "Creating partitions on $disk"
    if $needs_mklabel; then
      parted --script "$disk" mklabel gpt
      partprobe "$disk" 2>/dev/null || true
      sleep 1
    fi
    if [[ -z "$shared_esp" ]]; then
      parted --script "$disk" mkpart primary fat32 "${EFI_START_B}B" "${EFI_END_B}B"
      parted --script "$disk" set "$efi_part_num" esp on
      parted --script "$disk" name "$efi_part_num" "OMARCHY_EFI"
    fi
    parted --script "$disk" mkpart primary btrfs "${ROOT_START_B}B" "${ROOT_END_B}B"
    parted --script "$disk" name "$root_part_num" "OMARCHY_ROOT"
    partprobe "$disk"
    sync
    sleep 2

    local i
    for i in 1 2 3 4 5; do
      if [[ -b "$root_partition_device" ]] && { [[ -n "$shared_esp" ]] || [[ -b "$efi_dev" ]]; }; then
        break
      fi
      sleep 1
    done

    if [[ "$encrypt_installation" == "true" ]]; then
      step "Setting up LUKS2 on $root_partition_device"
      printf "%s" "$password" | cryptsetup luksFormat --type luks2 --batch-mode "$root_partition_device" -
      printf "%s" "$password" | cryptsetup open "$root_partition_device" omarchy_root -
      root_mapper="/dev/mapper/omarchy_root"
    else
      wipefs -af "$root_partition_device"
      root_mapper="$root_partition_device"
    fi

    step "Creating Btrfs filesystem and subvolumes"
    mkfs.btrfs -f -L OMARCHY "$root_mapper"
    for i in 1 2 3 4 5; do
      [[ -b "$root_mapper" ]] && break
      sleep 1
    done

    mkdir -p /mnt/btrfs-root
    mount "$root_mapper" /mnt/btrfs-root
    btrfs subvolume create /mnt/btrfs-root/@
    btrfs subvolume create /mnt/btrfs-root/@home
    btrfs subvolume create /mnt/btrfs-root/@log
    btrfs subvolume create /mnt/btrfs-root/@pkg
    umount /mnt/btrfs-root
    rmdir /mnt/btrfs-root

    mount -o noatime,compress=zstd,subvol=@ "$root_mapper" /mnt
    mkdir -p /mnt/home /mnt/var/log /mnt/var/cache/pacman/pkg /mnt"$esp_mount_in_target"
    mount -o noatime,compress=zstd,subvol=@home "$root_mapper" /mnt/home
    mount -o noatime,compress=zstd,subvol=@log  "$root_mapper" /mnt/var/log
    mount -o noatime,compress=zstd,subvol=@pkg  "$root_mapper" /mnt/var/cache/pacman/pkg

    if [[ -z "$shared_esp" ]]; then
      mkfs.fat -F32 -n OMARCHY_EFI "$efi_dev"
    fi
    mount "$efi_dev" /mnt"$esp_mount_in_target"
  fi

  clear

  echo "$full_name" >user_full_name.txt
  echo "$email_address" >user_email_address.txt
  printf '%s\n' "$encrypt_installation" >user_encrypt_installation.txt

  local password_escaped password_hash_escaped username_escaped
  password_escaped=$(echo -n "$password" | jq -Rsa)
  password_hash_escaped=$(echo -n "$password_hash" | jq -Rsa)
  username_escaped=$(echo -n "$username" | jq -Rsa)

  local credentials_encryption_line=""
  if [[ "$encrypt_installation" == "true" ]]; then
    credentials_encryption_line="    \"encryption_password\": $password_escaped,"
  fi

  cat <<-_EOF_ >user_credentials.json
{
$credentials_encryption_line
    "root_enc_password": $password_hash_escaped,
    "users": [
        {
            "enc_password": $password_hash_escaped,
            "groups": [],
            "sudo": true,
            "username": $username_escaped
        }
    ]
}
_EOF_

  local luks_uuid_json="null" luks_uuid
  if [[ "$encrypt_installation" == "true" ]]; then
    if $dry; then
      luks_uuid="00000000-0000-0000-0000-000000000000"
    else
      luks_uuid=$(blkid -s UUID -o value "$root_partition_device")
    fi
    luks_uuid_json="\"$luks_uuid\""
  fi

  cat <<-_EOF_ >user_configuration.json
{
    "app_config": null,
    "archinstall-language": "English",
    "auth_config": {},
    "audio_config": { "audio": "pipewire" },
    "bootloader_config": { "bootloader": "Limine", "uki": false, "removable": false },
    "custom_commands": [],
    "omarchy_install": {
        "mode": "protected",
        "target_mount": "/mnt",
        "boot": {
            "esp_mount": "$esp_mount_in_target",
            "esp_path": "/EFI/limine",
            "efi_binary": "limine_x64.efi",
            "enable_fallback": false
        },
        "storage": {
            "esp_device": "$efi_dev",
            "root_device": "$root_partition_device",
            "root_mapper": "$root_mapper",
            "luks_uuid": $luks_uuid_json,
            "kernel": "$kernel_choice"
        }
    },
    "disk_config": {
        "config_type": "pre_mounted_config",
        "mountpoint": "/mnt"
    },
    "hostname": "$hostname",
    "kernels": [ "$kernel_choice" ],
    "network_config": { "type": "iso" },
    "ntp": true,
    "parallel_downloads": 8,
    "script": null,
    "services": [],
    "swap": true,
    "timezone": "$timezone",
    "locale_config": {
        "kb_layout": "$keyboard",
        "sys_enc": "UTF-8",
        "sys_lang": "en_US.UTF-8"
    },
    "mirror_config": {
        "custom_repositories": [],
        "custom_servers": [
            {"url": "https://mirror.omarchy.org/\$repo/os/\$arch"},
            {"url": "https://mirror.rackspace.com/archlinux/\$repo/os/\$arch"},
            {"url": "https://geo.mirror.pkgbuild.com/\$repo/os/\$arch"}
        ],
        "mirror_regions": {},
        "optional_repositories": []
    },
    "packages": [
        "base-devel",
        "git",
        "omarchy-keyring",
        "omarchy"
    ],
    "profile_config": {
        "gfx_driver": null,
        "greeter": null,
        "profile": {}
    },
    "version": "3.0.9"
}
_EOF_

  if $dry; then
    echo -e "User Configuration:"
    cat user_configuration.json
    echo -e "\n\nUser Credentials:"
    cat user_credentials.json
    echo -e "\n\nUser Full Name:"
    cat user_full_name.txt
    echo -e "\nUser Email Address:"
    cat user_email_address.txt
    echo -e "\nDisk Encryption:"
    cat user_encrypt_installation.txt
    rm -f user_configuration.json user_credentials.json user_full_name.txt user_email_address.txt user_encrypt_installation.txt
  fi
}

disk_form() {
  step "Let's select where to install Omarchy..."

  # Don't offer the install media as an option (Arch ISO mounts it here)
  local boot_source
  local exclude_disk
  boot_source=$(findmnt -no SOURCE /run/archiso/bootmnt 2>/dev/null || true)
  exclude_disk=$(get_root_disk "$boot_source")

  # List all installable disks, excluding the boot device if present
  local available_disks
  available_disks=$(
    lsblk -dpno NAME,TYPE |
      awk '$2=="disk"{print $1}' |
      grep -E '/dev/(sd|hd|vd|nvme|mmcblk|xv)' |
      { if [[ -n "$exclude_disk" ]]; then grep -Fvx "$exclude_disk"; else cat; fi; }
  )

  # Get available disks and format them with info
  local disk_options=""
  while IFS= read -r device; do
    if [[ -n "$device" ]]; then
      disk_info=$(get_disk_info "$device")
      disk_options="$disk_options$disk_info"$'\n'
    fi
  done <<<"$available_disks"

  selected_display=$(echo "$disk_options" | gum choose --header "Select install disk") || abort
  disk=$(echo "$selected_display" | awk '{print $1}')
}

disk_form

# STEP 4: INSTALL MODE

open_partition_tool() {
  step "Partition tool for $disk"
  gum style "Create unallocated free space for Omarchy, then write changes and quit."
  gum style --foreground 8 "Do not create an Omarchy partition here — leave the target area as Free space."
  gum style --foreground 8 "Tip: free-space install needs at least 32GB unallocated, plus 2GB more if no ESP already exists."
  echo
  gum confirm --affirmative "Open cfdisk" --negative "Back" "Open partition tool for $disk?" || return 0

  clear
  cfdisk "$disk" || true
  partprobe "$disk" 2>/dev/null || true
  sync
  sleep 1
}

install_mode_form() {
  local choices=("Full disk install")

  # The free-space/protected path registers an EFI boot entry and is UEFI-only.
  # Full-disk installs still support BIOS through the orchestrator's BIOS branch.
  if [[ -d /sys/firmware/efi ]]; then
    choices+=("Free space install (alongside existing data)")
  fi
  choices+=("Choose a different disk")

  step "Installation mode"
  install_mode=$(gum choose --header "How should Omarchy be installed on $disk?" "${choices[@]}") || abort
}

while true; do
  install_mode_form

  case "$install_mode" in
    "Full disk install")
      break
      ;;
    "Free space install (alongside existing data)")
      if run_partition_install "$@"; then
        exit 0
      fi
      ;;
    "Choose a different disk")
      disk_form
      ;;
  esac
done

confirm_disk_overwrite() {
  local mode="encrypted"
  local affirmative confirm_status

  while true; do
    clear_logo
    echo
    gum style "Everything will be overwritten. There is no recovery possible."

    if [[ $mode == "encrypted" ]]; then
      gum style --foreground 8 "Press Ctrl+C for unencrypted install."
      affirmative="Yes, install"
    else
      affirmative="Yes, install without encryption"
    fi

    echo
    gum confirm --affirmative "$affirmative" --negative "No, change it" "Confirm overwriting ${disk}"
    confirm_status=$?

    case $confirm_status in
      0)
        [[ $mode == "encrypted" ]] && encrypt_installation=true || encrypt_installation=false
        return 0
        ;;
      1)
        return 1
        ;;
      130)
        [[ $mode == "encrypted" ]] && mode="unencrypted" || mode="encrypted"
        ;;
      *)
        abort
        ;;
    esac
  done
}

while ! confirm_disk_overwrite; do
  disk_form
done

clear

# Save user full name and email address, so the autostart can
# set them as OMARCHY_USER_NAME and OMARCHY_USER_EMAIL
echo "$full_name" >user_full_name.txt
echo "$email_address" >user_email_address.txt

# Escape username and password for safe inclusion
password_escaped=$(echo -n "$password" | jq -Rsa)
password_hash_escaped=$(echo -n "$password_hash" | jq -Rsa)
username_escaped=$(echo -n "$username" | jq -Rsa)

if [[ $encrypt_installation == true ]]; then
  credentials_encryption_line="    \"encryption_password\": $password_escaped,"
else
  credentials_encryption_line=""
fi

cat <<-_EOF_ >user_credentials.json
{
$credentials_encryption_line
    "root_enc_password": $password_hash_escaped,
    "users": [
        {
            "enc_password": $password_hash_escaped,
            "groups": [],
            "sudo": true,
            "username": $username_escaped
        }
    ]
}
_EOF_

printf '%s\n' "$encrypt_installation" >user_encrypt_installation.txt

# Setup partition layout
disk_size=$(lsblk -bdno SIZE "$disk")
mib=$((1024 * 1024))
gib=$((mib * 1024))
disk_size_in_mib=$((disk_size / mib * mib)) # Rounds to nearest MiB

gpt_backup_reserve=$((mib))
boot_partition_start=$((mib))
boot_partition_size=$((2 * gib))

main_partition_start=$((boot_partition_size + boot_partition_start))
main_partition_size=$((disk_size_in_mib - main_partition_start - gpt_backup_reserve))

# Detect T2 Mac and set appropriate kernel
if lspci -nn 2>/dev/null | grep -q "106b:180[12]"; then
  kernel_choice="linux-t2"
else
  kernel_choice="linux"
fi

if [[ $encrypt_installation == true ]]; then
  disk_encryption_config=$(cat <<-_EOF_
,
        "disk_encryption": {
            "encryption_type": "luks",
            "lvm_volumes": [],
            "iter_time": 2000,
            "partitions": [ "8c2c2b92-1070-455d-b76a-56263bab24aa" ],
            "encryption_password": $password_escaped
        }
_EOF_
)
else
  disk_encryption_config=""
fi

cat <<-_EOF_ >user_configuration.json
{
    "app_config": null,
    "archinstall-language": "English",
    "auth_config": {},
    "audio_config": { "audio": "pipewire" },
    "bootloader_config": { "bootloader": "Limine", "uki": false, "removable": false },
    "custom_commands": [],
    "omarchy_install": {
        "mode": "full_disk",
        "target_mount": "/mnt",
        "boot": {
            "esp_mount": "/boot",
            "esp_path": "/EFI/limine",
            "efi_binary": "limine_x64.efi",
            "enable_fallback": true
        },
        "storage": {
            "kernel": "$kernel_choice"
        }
    },
    "disk_config": {
        "btrfs_options": { 
          "snapshot_config": {
            "type": "Snapper"
          }
        },
        "config_type": "default_layout",
        "device_modifications": [
            {
                "device": "$disk",
                "partitions": [
                    {
                        "btrfs": [],
                        "dev_path": null,
                        "flags": [ "boot", "esp" ],
                        "fs_type": "fat32",
                        "mount_options": [],
                        "mountpoint": "/boot",
                        "obj_id": "ea21d3f2-82bb-49cc-ab5d-6f81ae94e18d",
                        "size": {
                            "sector_size": { "unit": "B", "value": 512 },
                            "unit": "B",
                            "value": $boot_partition_size
                        },
                        "start": {
                            "sector_size": { "unit": "B", "value": 512 },
                            "unit": "B",
                            "value": $boot_partition_start
                        },
                        "status": "create",
                        "type": "primary"
                    },
                    {
                        "btrfs": [
                            { "mountpoint": "/", "name": "@" },
                            { "mountpoint": "/home", "name": "@home" },
                            { "mountpoint": "/var/log", "name": "@log" },
                            { "mountpoint": "/var/cache/pacman/pkg", "name": "@pkg" }
                        ],
                        "dev_path": null,
                        "flags": [],
                        "fs_type": "btrfs",
                        "mount_options": [ "compress=zstd" ],
                        "mountpoint": null,
                        "obj_id": "8c2c2b92-1070-455d-b76a-56263bab24aa",
                        "size": {
                            "sector_size": { "unit": "B", "value": 512 },
                            "unit": "B",
                            "value": $main_partition_size
                        },
                        "start": {
                            "sector_size": { "unit": "B", "value": 512 },
                            "unit": "B",
                            "value": $main_partition_start
                        },
                        "status": "create",
                        "type": "primary"
                    }
                ],
                "wipe": true
            }
        ]$disk_encryption_config
    },
    "hostname": "$hostname",
    "kernels": [ "$kernel_choice" ],
    "network_config": { "type": "iso" },
    "ntp": true,
    "parallel_downloads": 8,
    "script": null,
    "services": [],
    "swap": true,
    "timezone": "$timezone",
    "locale_config": {
        "kb_layout": "$keyboard",
        "sys_enc": "UTF-8",
        "sys_lang": "en_US.UTF-8"
    },
    "mirror_config": {
        "custom_repositories": [],
        "custom_servers": [
            {"url": "https://mirror.omarchy.org/\$repo/os/\$arch"},
            {"url": "https://mirror.rackspace.com/archlinux/\$repo/os/\$arch"},
            {"url": "https://geo.mirror.pkgbuild.com/\$repo/os/\$arch"}
        ],
        "mirror_regions": {},
        "optional_repositories": []
    },
    "packages": [
        "base-devel",
        "git",
        "omarchy-keyring",
        "omarchy"
    ],
    "profile_config": {
        "gfx_driver": null,
        "greeter": null,
        "profile": {}
    },
    "version": "3.0.9"
}
_EOF_

if [[ $1 == "dry" ]]; then
  echo -e "User Configuration:"
  cat user_configuration.json

  echo -e "\n\nUser Credentials:"
  cat user_credentials.json

  echo -e "\n\nUser Full Name:"
  cat user_full_name.txt

  echo -e "\nUser Email Address:"
  cat user_email_address.txt

  echo -e "\nDisk Encryption:"
  cat user_encrypt_installation.txt

  rm user_configuration.json user_credentials.json user_full_name.txt user_email_address.txt user_encrypt_installation.txt
fi
