#!/bin/bash

# omarchy:summary=Open a generic image selector menu
# omarchy:args=[--selected <image>] [--print-name] [--show-labels] [--filterable] [--lazy-thumbnails] [--preload] [--cache-only] <image-dir>...

selected_image=""
print_name=false
show_labels=false
filterable=false
lazy_thumbnails=false
prepare_only=false
preload=false
cache_only=false
image_dirs=()

usage() {
  echo "Usage: omarchy-menu-images [--selected <image>] [--print-name] [--show-labels] [--filterable] [--lazy-thumbnails] [--preload] [--cache-only] <image-dir>..."
}

while (( $# > 0 )); do
  case "$1" in
    --selected)
      if (( $# < 2 )); then
        usage >&2
        exit 1
      fi

      selected_image="$2"
      shift 2
      ;;
    --print-name)
      print_name=true
      shift
      ;;
    --show-labels)
      show_labels=true
      shift
      ;;
    --filterable)
      filterable=true
      shift
      ;;
    --lazy-thumbnails)
      lazy_thumbnails=true
      shift
      ;;
    --prepare-only)
      prepare_only=true
      shift
      ;;
    --preload)
      preload=true
      shift
      ;;
    --cache-only)
      cache_only=true
      shift
      ;;
    --help|-h)
      usage
      exit 0
      ;;
    *)
      image_dirs+=("$1")
      shift
      ;;
  esac
done

if (( ${#image_dirs[@]} == 0 )); then
  usage >&2
  exit 1
fi

selection_file=$(mktemp)
done_file=$(mktemp)
pending_file=$(mktemp)
pending_video_file=$(mktemp)
rm -f "$done_file"
trap 'rm -f "$selection_file" "$done_file" "$pending_file" "$pending_video_file"' EXIT

image_dirs_env=""
for dir in "${image_dirs[@]}"; do
  if [[ -z $image_dirs_env ]]; then
    image_dirs_env="$dir"
  else
    image_dirs_env+=$'\n'"$dir"
  fi
done

current_image=$(readlink -f "$selected_image" 2>/dev/null)
selected_list_image=""

if [[ -n $current_image ]]; then
  for dir in "${image_dirs[@]}"; do
    if [[ -d $dir && -f $selected_image && ${selected_image%/*} == "$dir" ]]; then
      selected_list_image="$selected_image"
      break
    elif [[ -d $dir ]]; then
      selected_list_image=$(find -L "$dir" -maxdepth 1 -type f -samefile "$current_image" -print -quit 2>/dev/null)
      [[ -n $selected_list_image ]] && break
    fi
  done
fi

cache_dir=${XDG_CACHE_HOME:-$HOME/.cache}/omarchy/image-selector
index_file="$cache_dir/index.tsv"
rows=""
mkdir -p "$cache_dir"

cache_key=$(printf '%s' "$image_dirs_env" | md5sum | cut -d ' ' -f 1)
rows_cache_file="$cache_dir/$cache_key.rows"
rows_signature_file="$cache_dir/$cache_key.signature"
rows_fast_signature_file="$cache_dir/$cache_key.fast-signature"
rows_signature="v4"$'\n'
rows_fast_signature="v3"$'\n'
rows_cacheable=true
rows_cache_hit=false
image_files=()

for dir in "${image_dirs[@]}"; do
  [[ -d $dir ]] && rows_fast_signature+="$dir:$(stat -Lc '%Y' "$dir")"$'\n'
done

if [[ -f $rows_cache_file && -f $rows_fast_signature_file ]] && cmp -s "$rows_fast_signature_file" <(printf '%s' "$rows_fast_signature"); then
  rows=$(<"$rows_cache_file")
  rows_cache_hit=true
else
  for dir in "${image_dirs[@]}"; do
    if [[ -d $dir ]]; then
      rows_signature+="$dir:$(stat -Lc '%Y' "$dir")"$'\n'

      while IFS= read -r -d '' image; do
        image_files+=("$image")
        image_signature=$(stat -Lc '%s:%Y' "$image") || continue
        rows_signature+="$image:$image_signature"$'\n'
      done < <(find -L "$dir" -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)
    fi
  done
fi

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

generate_thumbnail() {
  local image="$1"
  local thumbnail="$2"
  local lock="$thumbnail.lock"
  local lock_fd
  local tmp="$thumbnail.$$.jpg"
  local thumbnail_command

  # Older releases used directories as locks, which could survive a killed
  # generator and block this thumbnail forever. Only reap aged ones, so a
  # legacy generator still running through an upgrade keeps its lock.
  if [[ -d $lock ]] && (( $(date +%s) - $(stat -c '%Y' "$lock" 2>/dev/null || date +%s) > 120 )); then
    rmdir "$lock" 2>/dev/null
  fi

  exec {lock_fd}>"$lock" || return
  flock -w 30 "$lock_fd" || return

  # A generator killed mid-write leaves its partial $thumbnail.<pid>.jpg
  # behind. Only the lock holder writes these, so any found now are stale.
  rm -f "$thumbnail".*.jpg

  [[ -f $thumbnail ]] && return

  # Callers fan out one generator per file, so keep each image conversion
  # single-threaded. ffmpegthumbnailer provides a still preview for videos.
  if is_video_path "$image"; then
    # Videos are generated before the picker opens, so one unreadable or
    # stalled file must not hold it shut. A failed run drops the row.
    thumbnail_command=(timeout -k 5 10 ffmpegthumbnailer -i "$image" -o "$tmp" -s 1536 -q 8)
  else
    thumbnail_command=(env VIPS_CONCURRENCY=1 vipsthumbnail "$image" --size 1536x864 --smartcrop=centre --path "$tmp[Q=82,strip]")
  fi

  # Close the lock fd for the converter: an orphaned or hung child must not
  # keep holding the lock after this shell is killed.
  if "${thumbnail_command[@]}" {lock_fd}>&-; then
    mv -f "$tmp" "$thumbnail"
  else
    status=$?
    rm -f "$tmp" "$thumbnail"
    # Remember a video the converter rejected, so it costs nothing on the next
    # open. The key covers size and mtime, so a repaired file starts clean. A
    # timeout is left to retry: the machine may only have been busy.
    if is_video_path "$image" && (( status != 124 && status != 137 )); then
      : >"$thumbnail.failed"
    fi
  fi
}

thumbnail_for() {
  local image="$1"
  local signature hash thumbnail

  signature=$(stat -Lc '%s:%Y' "$image") || return
  hash=$(awk -F '\t' -v path="$image" -v sig="$signature" '$1 == path && $2 == sig { print $3; exit }' "$index_file" 2>/dev/null)

  if [[ -z $hash ]]; then
    hash=$(printf '%s\t%s' "$image" "$signature" | md5sum | cut -d ' ' -f 1)
    printf '%s\t%s\t%s\n' "$image" "$signature" "$hash" >>"$index_file"
  fi

  thumbnail="$cache_dir/$hash.jpg"

  if [[ ! -f $thumbnail ]]; then
    # A video that already failed to convert has no row to offer. Hand the
    # marker back so the caller can keep the rows uncached over its absence.
    if is_video_path "$image" && [[ -f $thumbnail.failed ]]; then
      printf '%s' "$thumbnail.failed"
      return
    fi

    # A lazy row stands in with the media file itself, which the picker draws
    # with an Image -- fine for a picture, blank for a video. Videos take the
    # queue instead, which also keeps them under its narrower fan out.
    if [[ $lazy_thumbnails == true && $cache_only != true ]] && ! is_video_path "$image"; then
      rows_cacheable=false

      if [[ $prepare_only != true ]]; then
        generate_thumbnail "$image" "$thumbnail" >/dev/null 2>&1 &
      fi

      printf '%s' "$image"
      return
    fi

    if is_video_path "$image"; then
      printf '%s\0%s\0' "$image" "$thumbnail" >>"$pending_video_file"
    else
      printf '%s\0%s\0' "$image" "$thumbnail" >>"$pending_file"
    fi
  fi

  printf '%s' "$thumbnail"
}

# Each vips run is single-threaded, so still images can fill every core.
# ffmpegthumbnailer leaves FFmpeg's automatic threading on, so a full-width fan
# out of those would put a codec thread pool on every core at once.
drain_pending_thumbnails() {
  local video_jobs

  export -f generate_thumbnail is_video_path

  if [[ -s $pending_file ]]; then
    xargs -a "$pending_file" -0 -n 2 -P "$(nproc)" \
      bash -c 'generate_thumbnail "$1" "$2"' _ >/dev/null 2>&1 || true
  fi

  if [[ -s $pending_video_file ]]; then
    video_jobs=$(( $(nproc) / 4 ))
    (( video_jobs > 0 )) || video_jobs=1
    xargs -a "$pending_video_file" -0 -n 2 -P "$video_jobs" \
      bash -c 'generate_thumbnail "$1" "$2"' _ >/dev/null 2>&1 || true
  fi
}

if [[ $rows_cache_hit != true && -f $rows_cache_file && -f $rows_signature_file ]] && cmp -s "$rows_signature_file" <(printf '%s' "$rows_signature"); then
  rows=$(<"$rows_cache_file")
  printf '%s' "$rows_fast_signature" >"$rows_fast_signature_file"
elif [[ $rows_cache_hit != true ]]; then
  for image in "${image_files[@]}"; do
    thumbnail=$(thumbnail_for "$image")
    [[ -n $thumbnail ]] || continue
    # Cached rows are trusted on the directory's mtime alone, which a file
    # repaired in place never changes. Leave them uncached instead.
    if [[ $thumbnail == *.failed ]]; then
      rows_cacheable=false
      continue
    fi
    if [[ $lazy_thumbnails == true && $cache_only != true && $thumbnail == $image ]]; then
      rows_cacheable=false
    fi

    if [[ -z $rows ]]; then
      rows="$image"$'\t'"$thumbnail"
    else
      rows+=$'\n'"$image"$'\t'"$thumbnail"
    fi
  done

  drain_pending_thumbnails

  if [[ -s $pending_file || -s $pending_video_file ]]; then
    pruned=""
    while IFS=$'\t' read -r row_image row_thumbnail; do
      if [[ ! -e $row_thumbnail ]]; then
        rows_cacheable=false
        continue
      fi

      if [[ -z $pruned ]]; then
        pruned="$row_image"$'\t'"$row_thumbnail"
      else
        pruned+=$'\n'"$row_image"$'\t'"$row_thumbnail"
      fi
    done <<<"$rows"
    rows="$pruned"
  fi

  # Publish the cache under a lock and via renames: a picker killed mid-write,
  # or two pickers interleaving, must never leave truncated or mismatched rows
  # behind signatures that still validate. Rows go first so a kill between
  # renames leaves signatures that are either older (a harmless miss) or
  # describe the same directory state.
  if exec {rows_lock_fd}>"$rows_cache_file.lock" && flock -w 30 "$rows_lock_fd"; then
    if [[ $rows_cacheable == true ]]; then
      rm -f "$cache_dir/$cache_key".*.tmp
      printf '%s' "$rows" >"$rows_cache_file.$$.tmp" && mv -f "$rows_cache_file.$$.tmp" "$rows_cache_file"
      printf '%s' "$rows_signature" >"$rows_signature_file.$$.tmp" && mv -f "$rows_signature_file.$$.tmp" "$rows_signature_file"
      printf '%s' "$rows_fast_signature" >"$rows_fast_signature_file.$$.tmp" && mv -f "$rows_fast_signature_file.$$.tmp" "$rows_fast_signature_file"
    else
      rm -f "$rows_cache_file" "$rows_signature_file" "$rows_fast_signature_file"
    fi
    exec {rows_lock_fd}>&-
  fi
fi

if [[ $cache_only == true || $prepare_only == true ]]; then
  exit 0
fi

# Image rows can contain newlines and tabs, which don't survive positional
# shell IPC arguments. Base64-encode for transit; the ImagePicker plugin
# Qt.atob()s on the other side.
rows_b64=$(printf '%s' "$rows" | base64 -w 0)

if [[ $preload == true ]]; then
  omarchy-shell image-selector preload "$rows_b64" "$selected_list_image" "$show_labels" "$filterable" >/dev/null || true
  exit 0
fi

if ! open_result=$(omarchy-shell image-selector open \
     "" \
     "$rows_b64" \
     "$selected_list_image" \
     "$selection_file" \
     "$done_file" \
     "$show_labels" \
     "$filterable"); then
  echo "Image selector failed to accept request" >&2
  exit 1
fi

if [[ $open_result != "ok" ]]; then
  echo "Image selector failed to accept request" >&2
  exit 1
fi

while [[ ! -e $done_file ]]; do
  sleep 0.01
done

if [[ -s $selection_file ]]; then
  if [[ $print_name == true ]]; then
    selection=$(<"$selection_file")
    selection=${selection##*/}
    printf '%s\n' "${selection%.*}"
  else
    cat "$selection_file"
  fi
fi