#!/bin/bash

# omarchy:summary=Apply an Omarchy theme
# omarchy:args=<theme-name>
# omarchy:examples=omarchy theme list | omarchy theme set "Tokyo Night"

if [[ -z $1 ]]; then
  echo "Usage: omarchy-theme-set <theme-name>"
  exit 1
fi

CURRENT_THEME_PATH="$HOME/.local/state/omarchy/current/theme"
CURRENT_THEME_NAME_PATH="$HOME/.local/state/omarchy/current/theme.name"
NEXT_THEME_PATH="$HOME/.local/state/omarchy/current/next-theme"
CURRENT_BACKGROUND_LINK="$HOME/.local/state/omarchy/current/background"
THEME_BACKGROUND_STATE_PATH="$HOME/.local/state/omarchy/theme-backgrounds"
BACKGROUND_TRANSITION_CACHE="$HOME/.cache/omarchy/background-transitions"
THEME_SET_LOCK="${XDG_RUNTIME_DIR:-/tmp}/omarchy-theme-set.lock"
USER_THEMES_PATH="$HOME/.config/omarchy/themes"
OMARCHY_THEMES_PATH="$OMARCHY_PATH/themes"

# What a theme installed from a git repo may not ship, because these run code.
# Hyprland requires a theme's hyprland.lua and gum_env.lua at login and Neovim
# loads its neovim.lua at startup, so no .lua from such a theme is staged at all.
# Each terminal config names the program the terminal launches, and vscode.json
# names an extension omarchy-theme-set-vscode installs, a VS Code extension being
# arbitrary JavaScript. Everything else a theme ships is colour and is kept.
#
# Adding a template for another terminal, or for another editor that loads Lua,
# means adding it here. test/shell.d/theme-staging-test.sh fails on a generated
# theme file that is neither denied here nor recorded there as colour-only.
INSTALLED_THEME_DENIED=(alacritty.toml foot.ini ghostty.conf kitty.conf vscode.json)
IGNORED_THEME_FILES=()

run_parallel() {
  local pid
  local pids=()

  for command in "$@"; do
    bash -lc "$command" &
    pids+=("$!")
  done

  for pid in "${pids[@]}"; do
    wait "$pid"
  done
}

shell_ipc() {
  timeout 2 omarchy-shell "$@" >/dev/null 2>&1
}

is_video_path() {
  [[ ${1,,} =~ \.(mp4|m4v|mov|webm|mkv|avi)$ ]]
}

snapshot_background_path() {
  local background="$1"
  local name="$2"
  local snapshot extension

  [[ -f $background ]] || return
  is_video_path "$background" && return

  mkdir -p "$BACKGROUND_TRANSITION_CACHE"
  extension=${background##*.}
  snapshot="$BACKGROUND_TRANSITION_CACHE/$name-$$.$extension"
  ln "$background" "$snapshot" 2>/dev/null || cp "$background" "$snapshot"
  echo "$snapshot"
}

snapshot_current_background() {
  local current_background

  current_background=$(readlink -f "$CURRENT_BACKGROUND_LINK" 2>/dev/null || true)
  snapshot_background_path "$current_background" "previous"
}

theme_background_state_file() {
  local theme="$1"

  [[ -n $theme && ${theme##*/} == $theme && $theme != "." && $theme != ".." ]] || return 1
  printf '%s/%s\n' "$THEME_BACKGROUND_STATE_PATH" "$theme"
}

remember_current_theme_background() {
  local current_background state_file

  state_file=$(theme_background_state_file "$PREVIOUS_THEME_NAME") || return
  current_background=$(readlink -f "$CURRENT_BACKGROUND_LINK" 2>/dev/null || true)
  [[ -f $current_background ]] || return

  mkdir -p "$THEME_BACKGROUND_STATE_PATH"
  printf '%s\n' "$current_background" >"$state_file"
}

background_transition_uses_snapshots() {
  local next_background="$1"
  local current_background

  current_background=$(readlink -f "$CURRENT_BACKGROUND_LINK" 2>/dev/null || true)
  ! is_video_path "$current_background" && ! is_video_path "$next_background"
}

choose_theme_background() {
  local theme_path="${1:-$CURRENT_THEME_PATH}"
  local current_theme_backgrounds="$CURRENT_THEME_PATH/backgrounds"
  local backgrounds=()
  local current_background remembered_background state_file index next_index i

  CHOSEN_THEME_BACKGROUND=""
  mapfile -d '' -t backgrounds < <(
    find -L "$HOME/.config/omarchy/backgrounds/$THEME_NAME/" "$theme_path/backgrounds/" -maxdepth 1 -type f \
      \( -iname '*.jpg' -o -iname '*.jpeg' -o -iname '*.png' -o -iname '*.gif' -o -iname '*.bmp' -o -iname '*.webp' \
         -o -iname '*.mp4' -o -iname '*.m4v' -o -iname '*.mov' -o -iname '*.webm' -o -iname '*.mkv' -o -iname '*.avi' \) \
      -print0 2>/dev/null | sort -z
  )

  (( ${#backgrounds[@]} > 0 )) || return 1

  if [[ $PREVIOUS_THEME_NAME == "$THEME_NAME" ]]; then
    current_background=$(readlink "$CURRENT_BACKGROUND_LINK" 2>/dev/null || true)
    if [[ $theme_path != $CURRENT_THEME_PATH && ${current_background%/*} == $current_theme_backgrounds ]]; then
      current_background="$theme_path/backgrounds/${current_background##*/}"
    fi

    index=-1
    for i in "${!backgrounds[@]}"; do
      if [[ ${backgrounds[$i]} == $current_background ]]; then
        index=$i
        break
      fi
    done

    if (( index != -1 )); then
      next_index=$(((index + 1) % ${#backgrounds[@]}))
      CHOSEN_THEME_BACKGROUND="${backgrounds[$next_index]}"
    else
      CHOSEN_THEME_BACKGROUND="${backgrounds[0]}"
    fi
    return
  fi

  state_file=$(theme_background_state_file "$THEME_NAME" 2>/dev/null || true)
  remembered_background=$(cat "$state_file" 2>/dev/null || true)
  if [[ $theme_path != $CURRENT_THEME_PATH && ${remembered_background%/*} == $current_theme_backgrounds ]]; then
    remembered_background="$theme_path/backgrounds/${remembered_background##*/}"
  fi
  if [[ -f $remembered_background ]]; then
    CHOSEN_THEME_BACKGROUND="$remembered_background"
    return
  fi

  CHOSEN_THEME_BACKGROUND="${backgrounds[0]}"
}

choose_staged_theme_background() {
  local next_theme_backgrounds="$NEXT_THEME_PATH/backgrounds"

  choose_theme_background "$NEXT_THEME_PATH" || return 1

  if [[ ${CHOSEN_THEME_BACKGROUND%/*} == $next_theme_backgrounds ]]; then
    CHOSEN_THEME_BACKGROUND="$CURRENT_THEME_PATH/backgrounds/${CHOSEN_THEME_BACKGROUND##*/}"
  fi
}

set_theme_background_link() {
  choose_theme_background || return 1
  ln -nsf "$CHOSEN_THEME_BACKGROUND" "$CURRENT_BACKGROUND_LINK"
}

set_theme_background() {
  local new_background new_background_snapshot

  if [[ -z $CHOSEN_THEME_BACKGROUND || ! -f $CHOSEN_THEME_BACKGROUND ]]; then
    if ! choose_theme_background; then
      omarchy-notification-send "No background was found for theme" -t 2000
      shell_ipc shell applyTheme "$colors_payload" "$shell_payload" || true
      return
    fi
  fi

  new_background="$CHOSEN_THEME_BACKGROUND"
  new_background_snapshot=""
  if [[ $BACKGROUND_TRANSITION_SNAPSHOTS == "true" ]]; then
    new_background_snapshot=$(snapshot_background_path "$new_background" "next")
  fi

  if [[ -f $OLD_BACKGROUND_SNAPSHOT && -f $new_background_snapshot ]]; then
    shell_ipc background themeTransition "$OLD_BACKGROUND_SNAPSHOT" "$new_background_snapshot" "$new_background" "$colors_payload" "$shell_payload" || \
      shell_ipc shell applyTheme "$colors_payload" "$shell_payload" || true
    (sleep 3; rm -f "$OLD_BACKGROUND_SNAPSHOT" "$new_background_snapshot") &
  elif [[ -f $new_background_snapshot ]]; then
    shell_ipc background themeTransition "" "$new_background_snapshot" "$new_background" "$colors_payload" "$shell_payload" || \
      shell_ipc shell applyTheme "$colors_payload" "$shell_payload" || true
    (sleep 3; rm -f "$new_background_snapshot") &
  else
    shell_ipc background themeTransition "$OLD_BACKGROUND_SNAPSHOT" "$new_background" "$new_background" "$colors_payload" "$shell_payload" || \
      shell_ipc shell applyTheme "$colors_payload" "$shell_payload" || true
    if [[ -f $OLD_BACKGROUND_SNAPSHOT ]]; then
      (sleep 3; rm -f "$OLD_BACKGROUND_SNAPSHOT") &
    fi
  fi

  ln -nsf "$new_background" "$CURRENT_BACKGROUND_LINK"
}

is_denied_installed_file() {
  local name="$1"
  local denied

  [[ $name == *.lua ]] && return 0

  for denied in "${INSTALLED_THEME_DENIED[@]}"; do
    [[ $name == "$denied" ]] && return 0
  done

  return 1
}

# Themes older than colors.toml still get their palette, but their
# alacritty.toml never reaches the staged theme: an Alacritty config names the
# program the terminal launches.
stage_installed_colors_from_alacritty() {
  local source="$1"
  local scratch

  if [[ -f $NEXT_THEME_PATH/colors.toml ]]; then
    return
  fi

  if [[ ! -f $source/alacritty.toml || -L $source/alacritty.toml ]]; then
    return
  fi

  scratch=$(mktemp -d)
  cp "$source/alacritty.toml" "$scratch/alacritty.toml"
  omarchy-theme-colors-from-alacritty "$scratch"

  if [[ -f $scratch/colors.toml ]]; then
    cp "$scratch/colors.toml" "$NEXT_THEME_PATH/colors.toml"
  fi

  rm -rf "$scratch"
}

# Copies a directory without ever following a symlink: in an installed theme one
# points wherever the theme author chose, which is how an unlock.png becomes a
# copy of any file the session can read.
stage_installed_dir() {
  local source="$1"
  local dest="$2"
  local entry name

  mkdir -p "$dest"

  for entry in "$source"/*; do
    [[ -e $entry && ! -L $entry ]] || continue
    name=${entry##*/}

    if [[ -d $entry ]]; then
      stage_installed_dir "$entry" "$dest/$name"
    else
      cp "$entry" "$dest/$name"
    fi
  done
}

# `omarchy theme install` clones into ~/.config/omarchy/themes, so a .git
# directory there means the contents came from a stranger and are held to the
# list above. A directory the user wrote themselves, and a symlink to their own
# working copy, are theirs to fill however they like -- the same distinction
# omarchy-theme-extras draws when it decides which themes it may pull.
theme_came_from_a_repo() {
  local source="$1"

  [[ ! -L $source && -d $source/.git ]]
}

stage_installed_theme() {
  local source="$1"
  local entry name

  [[ -d $source ]] || return 0

  for entry in "$source"/*; do
    [[ -e $entry ]] || continue
    name=${entry##*/}

    if [[ -L $entry ]] || is_denied_installed_file "$name"; then
      case "${name,,}" in
        readme* | license* | changelog* | *.md | *.txt) ;;
        *) IGNORED_THEME_FILES+=("$name") ;;
      esac
    elif [[ -d $entry ]]; then
      stage_installed_dir "$entry" "$NEXT_THEME_PATH/$name"
    else
      cp "$entry" "$NEXT_THEME_PATH/$name"
    fi
  done

  stage_installed_colors_from_alacritty "$source"
}

report_ignored_theme_files() {
  (( ${#IGNORED_THEME_FILES[@]} > 0 )) || return 0

  echo "Ignored in $USER_THEMES_PATH/$THEME_NAME: ${IGNORED_THEME_FILES[*]}" >&2
  echo "A theme installed from a git repo cannot supply Lua, a terminal config, or vscode.json." >&2
}

THEME_NAME=$(echo "$1" | sed -E 's/<[^>]+>//g' | tr '[:upper:]' '[:lower:]' | tr ' ' '-')
THEME_HEADLESS=0
if [[ ${OMARCHY_THEME_HEADLESS:-} == "1" || ${OMARCHY_THEME_OFFLINE:-} == "1" ]]; then
  THEME_HEADLESS=1
fi

if [[ -z $THEME_NAME || $THEME_NAME == .* || $THEME_NAME == */* ]]; then
  echo "Invalid theme name: $1"
  exit 1
fi

if [[ ! -d $OMARCHY_THEMES_PATH/$THEME_NAME ]] && [[ ! -d $USER_THEMES_PATH/$THEME_NAME ]]; then
  echo "Theme '$THEME_NAME' does not exist"
  exit 1
fi

# Serialize theme changes. Theme switching rebuilds a shared next-theme staging
# directory and updates current theme/background symlinks; concurrent calls can
# otherwise race and make one selection appear to be ignored.
exec 9>"$THEME_SET_LOCK"
flock 9

PREVIOUS_THEME_NAME=$(cat "$CURRENT_THEME_NAME_PATH" 2>/dev/null || true)
remember_current_theme_background

# Setup clean next theme directory (for atomic theme config swapping)
rm -rf "$NEXT_THEME_PATH"
mkdir -p "$NEXT_THEME_PATH"

# Copy official theme first, then overlay the user's theme on top
cp -r "$OMARCHY_THEMES_PATH/$THEME_NAME/"* "$NEXT_THEME_PATH/" 2>/dev/null

if theme_came_from_a_repo "$USER_THEMES_PATH/$THEME_NAME"; then
  stage_installed_theme "$USER_THEMES_PATH/$THEME_NAME"
  report_ignored_theme_files
else
  cp -r "$USER_THEMES_PATH/$THEME_NAME/"* "$NEXT_THEME_PATH/" 2>/dev/null
fi

# Generate colors.toml from alacritty.toml if theme is missing colors.toml
if [[ ! -f $NEXT_THEME_PATH/colors.toml && -f $NEXT_THEME_PATH/alacritty.toml ]]; then
  omarchy-theme-colors-from-alacritty "$NEXT_THEME_PATH"
fi

# Generate dynamic configs
omarchy-theme-set-templates

CHOSEN_THEME_BACKGROUND=""
OLD_BACKGROUND_SNAPSHOT=""
BACKGROUND_TRANSITION_SNAPSHOTS=true
if [[ $THEME_HEADLESS != "1" && $OMARCHY_THEME_SKIP_BACKGROUND != "1" ]]; then
  # Resolve the staged choice while the old theme still exists. Video changes
  # switch directly to the durable path in QML, so neither side needs a copy.
  if choose_staged_theme_background; then
    if background_transition_uses_snapshots "$CHOSEN_THEME_BACKGROUND"; then
      OLD_BACKGROUND_SNAPSHOT=$(snapshot_current_background)
    else
      BACKGROUND_TRANSITION_SNAPSHOTS=false
    fi
  fi
fi

# Swap next theme in as current
rm -rf "$CURRENT_THEME_PATH"
mv "$NEXT_THEME_PATH" "$CURRENT_THEME_PATH"

# Store theme name for reference
echo "$THEME_NAME" >"$CURRENT_THEME_NAME_PATH"

# Make the running shell pick up the new palette immediately while the rest of
# the theme hooks run.
colors_payload=$([[ -f $CURRENT_THEME_PATH/colors.toml ]] && base64 -w 0 "$CURRENT_THEME_PATH/colors.toml")
shell_payload=$([[ -f $CURRENT_THEME_PATH/shell.toml ]] && base64 -w 0 "$CURRENT_THEME_PATH/shell.toml")
if [[ $THEME_HEADLESS == "1" ]]; then
  # No shell/session bus exists during ISO chroot finalization, but the first
  # real login still needs a current background symlink for omarchy-shell to
  # render.
  [[ $OMARCHY_THEME_SKIP_BACKGROUND == "1" ]] || set_theme_background_link || true
elif [[ $OMARCHY_THEME_SKIP_BACKGROUND == "1" ]]; then
  shell_ipc shell applyTheme "$colors_payload" "$shell_payload" || true
else
  set_theme_background
fi

# The shared staging/current symlinks are updated and the shell has accepted
# the transition. Let another theme selection queue only behind that critical
# section, not behind slower app-retint hooks and selector cache warmups.
flock -u 9

post_theme_commands=(
  omarchy-restart-terminal
  omarchy-restart-hyprctl
  omarchy-restart-btop
  omarchy-restart-opencode
  omarchy-restart-helix
  omarchy-theme-set-foot
  omarchy-theme-set-tmux
  omarchy-theme-set-gnome
  omarchy-theme-set-pi
  omarchy-theme-set-claude
  omarchy-theme-set-hermes
  omarchy-theme-set-t3code
  omarchy-theme-set-browser
  omarchy-theme-set-vscode
  omarchy-theme-set-obsidian
  omarchy-theme-set-keyboard
)

if [[ $THEME_HEADLESS != "1" ]]; then
  run_parallel "${post_theme_commands[@]}"

  # Call hook on theme set
  omarchy-hook theme-set "$THEME_NAME" >/dev/null

  # Warm selector caches after the theme is applied. The shell hot-reloads theme
  # colors/backgrounds, so keep the running instance alive and preload the picker
  # rows/selection to avoid first-open carousel settling after a theme change.
  omarchy-theme-switcher --preload >/dev/null 2>&1
  omarchy-theme-bg-cache >/dev/null 2>&1 &
fi