#!/bin/bash

# omarchy:summary=Set the Plymouth boot theme colors and logo
# omarchy:args=<background-hex> <text-hex> <path-to-logo.png>
# omarchy:examples=omarchy plymouth set '#1d2021' '#ebdbb2' ~/.local/state/omarchy/current/theme/plymouth/logo.png
# omarchy:requires-sudo=true

set -euo pipefail

# Build the authoritative theme in a root-owned directory, then publish each
# fixed destination atomically. The caller opens the selected logo before sudo,
# so the privileged process never resolves a user-controlled input path.

refresh_default=false
if (( $# == 1 )) && [[ $1 == "--refresh-default" ]]; then
  refresh_default=true
elif (( $# != 3 )); then
  echo "Usage: omarchy-plymouth-set <background-hex> <text-hex> <path-to-logo.png>" >&2
  exit 1
fi

logo_fd=
if $refresh_default; then
  mode=refresh
  bg_hex=
  text_hex=
else
  mode=set
  bg_hex="${1#\#}"
  text_hex="${2#\#}"
  logo_path="$3"

  if ! [[ $bg_hex =~ ^[0-9a-fA-F]{6}$ ]]; then
    echo "Invalid background color: $1 (expected #RRGGBB)" >&2
    exit 1
  fi

  if ! [[ $text_hex =~ ^[0-9a-fA-F]{6}$ ]]; then
    echo "Invalid text color: $2 (expected #RRGGBB)" >&2
    exit 1
  fi

  if [[ ! -f $logo_path ]]; then
    echo "Logo file not found: $logo_path" >&2
    exit 1
  fi

  if [[ -L $logo_path ]]; then
    echo "Logo file is a symlink, which is not accepted: $logo_path" >&2
    exit 1
  fi

  # Open the logo while still unprivileged. A replacement symlink to a root-only
  # file therefore fails here instead of being followed after sudo starts.
  if ! exec {logo_fd}<"$logo_path"; then
    echo "Unable to open logo file as the current user: $logo_path" >&2
    exit 1
  fi
  if [[ ! -f /proc/$$/fd/$logo_fd ]]; then
    echo "Logo input is no longer a regular file: $logo_path" >&2
    exit 1
  fi
fi

run_root_transaction() {
  sudo /bin/bash -c '
    set -eEuo pipefail
    PATH=/usr/bin:/bin
    export PATH

    # Every check below is a bare assertion that aborts under set -e. Name the
    # subject of each one so a refusal reaches the user instead of exiting mute.
    failure_context="the privileged Plymouth transaction"
    failure_reported=
    report_failure() {
      [[ -z $failure_reported ]] || return 0
      failure_reported=1
      printf "omarchy-plymouth-set: refusing to publish: %s failed validation\n" "$failure_context" >&2
    }
    trap report_failure ERR

    mode=$1
    source_root=$2
    bg_hex=$3
    text_hex=$4
    max_asset_size=$5

    failure_context="the arguments of the privileged transaction"
    [[ $mode == "set" || $mode == "refresh" ]]
    [[ $source_root == /* ]]
    [[ $max_asset_size =~ ^[0-9]+$ ]]
    (( max_asset_size > 0 ))

    failure_context="the Omarchy source tree $source_root"
    canonical_source_root=$(realpath -e -- "$source_root")
    [[ $canonical_source_root == "$source_root" ]]

    # The one refusal a healthy machine can hit: omarchy dev link points
    # OMARCHY_PATH at a checkout the desktop user can rewrite, and executable
    # boot assets cannot come from there. Say so rather than failing mute.
    source_root_uid=$(stat -c %u -- "$source_root")
    if (( source_root_uid != 0 )); then
      printf "omarchy-plymouth-set: %s is not root-owned, so it cannot supply executable boot assets.\n" "$source_root" >&2
      printf "omarchy-plymouth-set: this is expected on a development checkout; run omarchy dev unlink to publish from /usr/share/omarchy.\n" >&2
      exit 1
    fi

    if [[ $mode == "set" ]]; then
      [[ $bg_hex =~ ^[0-9a-fA-F]{6}$ ]]
      [[ $text_hex =~ ^[0-9a-fA-F]{6}$ ]]
    fi

    theme_dir=/usr/share/plymouth/themes/omarchy
    sddm_dir=/usr/share/sddm/themes/omarchy
    plymouth_theme_assets=(
      bullet.png
      entry.png
      lock.png
      logo.png
      omarchy.plymouth
      omarchy.script
      preview-unlock.png
      progress_bar.png
      progress_box.png
    )
    plymouth_default_assets=("${plymouth_theme_assets[@]}" logos/oma.png)
    sddm_theme_assets=(Main.qml bullet.png entry-failed.png entry.png lock-failed.png lock.png logo.png)

    validate_trusted_directory() {
      local directory=$1 canonical uid directory_mode

      failure_context="directory $directory"
      canonical=$(realpath -e -- "$directory")
      [[ $canonical == "$directory" && -d $directory && ! -L $directory ]]

      while :; do
        failure_context="directory $directory (must be root-owned and not group- or world-writable)"
        uid=$(stat -c %u -- "$directory")
        directory_mode=$(stat -c %a -- "$directory")
        (( uid == 0 ))
        (( (8#$directory_mode & 0022) == 0 ))
        [[ $directory == "/" ]] && break
        directory=${directory%/*}
        [[ -n $directory ]] || directory=/
      done
    }

    validate_trusted_file() {
      local source=$1 canonical uid file_mode size

      failure_context="packaged source file $source"
      [[ -f $source && ! -L $source ]]
      canonical=$(realpath -e -- "$source")
      [[ $canonical == "$source" ]]
      validate_trusted_directory "${source%/*}"
      uid=$(stat -c %u -- "$source")
      file_mode=$(stat -c %a -- "$source")
      size=$(stat -c %s -- "$source")
      (( uid == 0 ))
      (( (8#$file_mode & 0022) == 0 ))
      (( size > 0 && size <= max_asset_size ))
    }

    copy_trusted_file() {
      local source=$1 destination=$2

      validate_trusted_file "$source"
      cp --reflink=never -- "$source" "$destination"
      chown 0:0 -- "$destination"
      chmod 0600 -- "$destination"
    }

    staging_dir=$(mktemp -d /tmp/omarchy-plymouth.XXXXXXXX)
    temporary=
    cleanup() {
      [[ -z $temporary ]] || rm -f -- "$temporary"
      rm -rf -- "$staging_dir"
    }
    trap cleanup EXIT HUP INT TERM
    chown 0:0 -- "$staging_dir"
    chmod 0700 -- "$staging_dir"

    plymouth_stage=$staging_dir/plymouth
    sddm_stage=$staging_dir/sddm
    mkdir -m 0700 -p -- "$plymouth_stage/logos" "$sddm_stage"

    if [[ $mode == "refresh" ]]; then
      assets_to_stage=("${plymouth_default_assets[@]}")
    else
      assets_to_stage=("${plymouth_theme_assets[@]}")
    fi

    for asset in "${assets_to_stage[@]}"; do
      copy_trusted_file "$source_root/default/plymouth/$asset" "$plymouth_stage/$asset"
    done

    if [[ $mode == "set" ]]; then
      # stdin was opened by the unprivileged caller. Read no more than the
      # documented limit into the root-owned stage before doing other work.
      failure_context="the selected logo (expected 1 to $max_asset_size bytes)"
      head -c "$((max_asset_size + 1))" >"$plymouth_stage/logo.png"
      logo_size=$(stat -c %s -- "$plymouth_stage/logo.png")
      (( logo_size > 0 && logo_size <= max_asset_size ))
      chown 0:0 -- "$plymouth_stage/logo.png"
      chmod 0600 -- "$plymouth_stage/logo.png"
      cp --reflink=never -- "$plymouth_stage/logo.png" "$sddm_stage/logo.png"

      bg_r=$(awk -v n=$((16#${bg_hex:0:2})) "BEGIN{printf \"%.3f\", n/255}")
      bg_g=$(awk -v n=$((16#${bg_hex:2:2})) "BEGIN{printf \"%.3f\", n/255}")
      bg_b=$(awk -v n=$((16#${bg_hex:4:2})) "BEGIN{printf \"%.3f\", n/255}")

      sed -i \
        -e "s/^Window.SetBackgroundTopColor.*/Window.SetBackgroundTopColor($bg_r, $bg_g, $bg_b);/" \
        -e "s/^Window.SetBackgroundBottomColor.*/Window.SetBackgroundBottomColor($bg_r, $bg_g, $bg_b);/" \
        "$plymouth_stage/omarchy.script"

      for asset in bullet.png entry.png lock.png progress_bar.png; do
        magick "$plymouth_stage/$asset" -channel RGB +level-colors "#$text_hex","#$text_hex" "$plymouth_stage/$asset"
      done

      copy_trusted_file "$source_root/default/sddm/omarchy/Main.qml" "$sddm_stage/Main.qml"
      sed -i \
        -e "s/#1a1b26/#__OMARCHY_SDDM_BG__/g" \
        -e "s/#ffffff/#__OMARCHY_SDDM_TEXT__/g" \
        -e "s/#__OMARCHY_SDDM_BG__/#$bg_hex/g" \
        -e "s/#__OMARCHY_SDDM_TEXT__/#$text_hex/g" \
        "$sddm_stage/Main.qml"

      for asset in bullet.png entry.png lock.png; do
        cp --reflink=never -- "$plymouth_stage/$asset" "$sddm_stage/$asset"
      done
      for asset in entry lock; do
        magick "$plymouth_stage/$asset.png" -channel RGB +level-colors "#f7768e","#f7768e" "$sddm_stage/$asset-failed.png"
      done
      chown -R 0:0 -- "$staging_dir"
      find "$staging_dir" -type f -exec chmod 0600 -- {} +
    fi

    publish_asset() {
      local source=$1 destination=$2 parent filename source_size copied_size

      failure_context="destination $destination"
      [[ -f $source && ! -L $source ]]
      (( $(stat -c %u -- "$source") == 0 ))
      source_size=$(stat -c %s -- "$source")
      (( source_size > 0 && source_size <= max_asset_size ))

      [[ $destination == /* && $destination != */ && $destination != *"/../"* ]]
      parent=${destination%/*}
      filename=${destination##*/}
      [[ -n $parent && -n $filename && $filename != "." && $filename != ".." ]]
      validate_trusted_directory "$parent"

      temporary=$(mktemp --tmpdir="$parent" ".$filename.omarchy-new.XXXXXXXX")
      cp --reflink=never -- "$source" "$temporary"
      copied_size=$(stat -c %s -- "$temporary")
      (( copied_size == source_size ))
      cmp -s -- "$source" "$temporary"
      chown 0:0 -- "$temporary"
      chmod 0644 -- "$temporary"
      sync -f -- "$temporary"
      mv --no-copy -fT -- "$temporary" "$destination"
      temporary=
    }

    for asset in "${assets_to_stage[@]}"; do
      publish_asset "$plymouth_stage/$asset" "$theme_dir/$asset"
    done

    if [[ $mode == "set" ]]; then
      for asset in "${sddm_theme_assets[@]}"; do
        publish_asset "$sddm_stage/$asset" "$sddm_dir/$asset"
      done
      validate_trusted_directory "$sddm_dir"
      rm -f -- "$sddm_dir/logo.svg"
    fi
  ' bash "$mode" "$OMARCHY_PATH" "$bg_hex" "$text_hex" "$((64 * 1024 * 1024))"
}

if $refresh_default; then
  run_root_transaction </dev/null
else
  run_root_transaction <&"$logo_fd"
fi

sudo plymouth-set-default-theme omarchy

if omarchy-cmd-present limine-mkinitcpio; then
  sudo limine-mkinitcpio
else
  sudo mkinitcpio -P
fi