#!/bin/bash
# sort-images — classify all images in a directory and group by category
#
# Usage:
#   sort-images                   # classify images in current directory
#   sort-images <directory>       # classify images in specified directory
#   sort-images -j <directory>    # JSON output
#   sort-images -n 3 <directory>  # show top N categories per image
#
# Examples:
#   sort-images ~/Photos
#   sort-images . -n 1
#   sort-images ~/Desktop -j | jq
#   sort-images /tmp/screenshots
#
# Requires: auge

set -euo pipefail

dir="."
json=false
topn=1

while [[ $# -gt 0 ]]; do
  case "$1" in
    -j|--json) json=true; shift ;;
    -n|--top)  topn="${2:-1}"; shift 2 ;;
    -h|--help)
      sed -n '2,/^$/{ s/^# //; s/^#//; p; }' "$0"
      exit 0 ;;
    *)
      dir="$1"; shift ;;
  esac
done

[[ -d "$dir" ]] || { echo "error: not a directory: $dir" >&2; exit 1; }

command -v auge >/dev/null || { echo "error: auge not found" >&2; exit 1; }

# Find image files
images=()
while IFS= read -r -d '' f; do
  images+=("$f")
done < <(find "$dir" -maxdepth 1 -type f \( -iname '*.png' -o -iname '*.jpg' -o -iname '*.jpeg' -o -iname '*.heic' -o -iname '*.tiff' -o -iname '*.bmp' -o -iname '*.gif' \) -print0 2>/dev/null | sort -z)

if [[ ${#images[@]} -eq 0 ]]; then
  echo "No images found in $dir" >&2
  exit 0
fi

echo "Classifying ${#images[@]} images..." >&2

if $json; then
  echo "["
  first=true
  for img in "${images[@]}"; do
    $first && first=false || echo ","
    name=$(basename "$img")
    labels=$(auge --classify "$img" --top "$topn" -q 2>/dev/null | head -"$topn")
    top=$(echo "$labels" | head -1 | sed 's/: [0-9]*%//')
    printf '  {"file": "%s", "category": "%s", "labels": "%s"}' "$name" "$top" "$(echo "$labels" | tr '\n' '; ' | sed 's/; $//')"
  done
  echo ""
  echo "]"
else
  # Collect category:file pairs, then group (bash 3.2 compatible)
  tmpcat=$(mktemp /tmp/sort-images-XXXXXX.txt)
  trap 'rm -f "$tmpcat"' EXIT

  for img in "${images[@]}"; do
    name=$(basename "$img")
    top=$(auge --classify "$img" --top 1 -q 2>/dev/null | head -1 | sed 's/: [0-9]*%//')
    [[ -z "$top" ]] && top="unknown"
    echo "$top	$name" >> "$tmpcat"
    printf '  %s -> %s\n' "$name" "$top" >&2
  done

  echo ""
  ncats=0
  sort "$tmpcat" | while IFS=$'\t' read -r cat name; do
    if [[ "$cat" != "${prev_cat:-}" ]]; then
      printf '\033[1m%s\033[0m\n' "$cat"
      prev_cat="$cat"
    fi
    echo "  $name"
  done

  ncats=$(cut -f1 "$tmpcat" | sort -u | wc -l | tr -d ' ')
  echo ""
  echo "${#images[@]} images in $ncats categories."
fi
