#!/bin/bash
# organize-images -- Sort images into AI/, Photos/, Screenshots/, Unsorted/
# Default: sub-categorizes AI by model + LoRA combo (uses extract-ai-meta).
# Use --flat for one flat AI/ folder.

TARGET=""
DRY_RUN=true
VERBOSE=false
INSTALL=false
FLAT=false
LIMIT=0

usage() {
  echo "Usage: $(basename "$0") [path] [-e|--execute] [-v|--verbose] [-i|--install] [--flat]"
  echo "  path       Directory to organize (default: ~/Pictures)"
  echo "  -e         Execute moves (default: dry-run only)"
  echo "  -v         Verbose: show per-file classification"
  echo "  -i         Install this script to ~/.local/bin/"
  echo "  --flat     Don't sub-categorize by model/LoRA (flat AI/ folder)"
  echo "  --limit N  Only process N files (for batching)"
  exit 0
}

ARGS=()
LIMIT_NEXT=false
for arg in "$@"; do
  if [ "$LIMIT_NEXT" = true ]; then
    LIMIT="$arg"
    LIMIT_NEXT=false
    continue
  fi
  case "$arg" in
    -e|--execute) DRY_RUN=false ;;
    -v|--verbose) VERBOSE=true ;;
    -i|--install) INSTALL=true ;;
    --flat) FLAT=true ;;
    --limit) LIMIT_NEXT=true ;;
    --limit=*) LIMIT="${arg#*=}" ;;
    -h|--help) usage ;;
    *) ARGS+=("$arg") ;;
  esac
done

if [ "$INSTALL" = true ]; then
  SRC="$(cd "$(dirname "$0")" && pwd -P)/$(basename "$0")"
  DEST="$HOME/.local/bin/$(basename "$0")"
  mkdir -p "$HOME/.local/bin"
  if [ "$SRC" = "$DEST" ]; then
    echo "Already installed at $DEST"
  else
    cp "$SRC" "$DEST"
    chmod +x "$DEST"
    echo "Installed to $DEST"
  fi
  exit 0
fi

TARGET="${ARGS[0]:-$HOME/Pictures}"
TARGET="${TARGET%/}"

if [ ! -d "$TARGET" ]; then
  echo "Error: '$TARGET' is not a directory" >&2
  exit 1
fi

echo "organize-images — Classifying images in: $TARGET"
[ "$DRY_RUN" = true ] && echo "  DRY RUN  (use -e to execute)" || echo "  EXECUTING"
[ "$LIMIT" -gt 0 ] && echo "  Limit: first $LIMIT files"

# --- Pre-scan: extract AI metadata for model/LoRA routing ---
AI_META=""
if [ "$FLAT" = false ]; then
  META_SCRIPT="$(cd "$(dirname "$0")" && pwd -P)/extract-ai-meta"
  if [ -x "$META_SCRIPT" ]; then
    echo "  Scanning AI metadata for model/LoRA routing..."
    AI_META=$(mktemp)
    if ! "$META_SCRIPT" "$TARGET" > "$AI_META" 2>/dev/null; then
      echo "  Warning: AI metadata scan failed, falling back to flat AI/"
      FLAT=true
      rm -f "$AI_META"
    elif [ ! -s "$AI_META" ]; then
      FLAT=true
      rm -f "$AI_META"
    fi
  else
    echo "  Warning: extract-ai-meta not found, falling back to flat AI/"
    FLAT=true
  fi
fi

[ "$FLAT" = false ] && echo "  AI sub-categorization: by model/LoRA (use --flat to disable)"
echo

classify() {
  local f="$1"

  local name
  name=$(basename "$f")
  case "$name" in
    CleanShot*|Screen_Shot*|Screenshot*) echo "Screenshots"; return ;;
  esac

  local comment
  comment=$(mdls -name kMDItemComment "$f" 2>/dev/null)
  if echo "$comment" | grep -q '= "Screenshot"'; then
    echo "Screenshots"; return
  fi

  case "$name" in
    ComfyUI_*|*_DEFI_*|R.*.png|schmeeve-AI-*) echo "AI"; return ;;
  esac

  if mdls -name kMDItemAcquisitionMake "$f" 2>/dev/null | grep -q '= "'; then
    echo "Photos"; return
  fi

  local ext="${f##*.}"
  ext=$(echo "$ext" | tr '[:upper:]' '[:lower:]')
  case "$ext" in
    dng|cr2|nef|arw|orf|rw2|raf|3fr)
      echo "Photos"; return ;;
  esac

  if [ "$ext" = "png" ]; then
    local dims w h area ratio
    dims=$(sips --getProperty pixelWidth --getProperty pixelHeight "$f" 2>/dev/null | grep -E 'pixel(Width|Height)' | awk '{print $NF}')
    w=$(echo "$dims" | head -1)
    h=$(echo "$dims" | tail -1)
    if [ -n "$w" ] && [ -n "$h" ] && [ "$w" -gt 0 ] && [ "$h" -gt 0 ]; then
      area=$((w * h))
      ratio=$(echo "scale=4; $w / $h" | bc 2>/dev/null || echo 1)
      if [ "$area" -gt 1048576 ] && [ "$(echo "$ratio > 0.75" | bc -l)" = 1 ] && [ "$(echo "$ratio < 1.33" | bc -l)" = 1 ]; then
        echo "AI"; return
      fi
    fi
  fi

  echo "Unsorted"
}

# --- Pre-scan: collect all image files ---
FILES=()
for f in "$TARGET"/*; do
  [ -f "$f" ] || continue
  ext="${f##*.}"; ext=$(echo "$ext" | tr '[:upper:]' '[:lower:]')
  case "$ext" in
    jpg|jpeg|png|webp|gif|bmp|tiff|tif|heic|heif|dng|cr2|nef|arw|orf|rw2|raf|3fr)
      FILES+=("$f") ;;
  esac
done
TOTAL=${#FILES[@]}
if [ "$LIMIT" -gt 0 ] && [ "$LIMIT" -lt "$TOTAL" ]; then
  FILES=("${FILES[@]:0:$LIMIT}")
  TOTAL=$LIMIT
fi
if [ "$TOTAL" -eq 0 ]; then
  echo "  No image files found."
  exit 0
fi

mkdir -p "$TARGET/AI" "$TARGET/Photos" "$TARGET/Screenshots" "$TARGET/Unsorted"

cAI=0; cPhotos=0; cScreenshots=0; cUnsorted=0

progress_bar() {
  local current=$1 total=$2 w=40
  local pct=0 filled=0 empty=0
  [ "$total" -gt 0 ] && pct=$((current * 100 / total))
  [ "$total" -gt 0 ] && filled=$((current * w / total))
  empty=$((w - filled))
  printf "\r  ["
  for ((i=0; i<filled; i++)); do printf "#"; done
  for ((i=0; i<empty; i++)); do printf "."; done
  printf "] %3d%% (%d/%d)" "$pct" "$current" "$total"
}

ai_dest() {
  python3 -c "
import json, sys
m = json.load(open('$AI_META'))
entry = m.get(sys.argv[1])
if not entry or not isinstance(entry, dict):
    print('AI/unknown')
    sys.exit(0)
def clean(s):
    if not s:
        return 'unknown'
    return s.replace('/', '-').replace('\\\\', '-')

model = clean(entry.get('model'))
loras = entry.get('loras') or []
if loras:
    combo = '+'.join(clean(l.strip()) for l in loras if l)
    combo = (combo[:200]).rstrip('+')
    print(f'AI/{model}/{combo}')
else:
    print(f'AI/{model}/no-lora')
" "$1"
}

idx=0
for f in "${FILES[@]}"; do
  idx=$((idx + 1))

  cat=$(classify "$f")

  if [ "$cat" = "AI" ] && [ "$FLAT" = false ] && [ -n "$AI_META" ]; then
    sub=$(ai_dest "$(basename "$f")" 2>/dev/null) || sub=""
    if [ -z "$sub" ]; then
      dest="$TARGET/AI/unknown/"
    else
      dest="$TARGET/$sub"
    fi
  else
    dest="$TARGET/$cat/"
  fi

  if [ "$VERBOSE" = true ]; then
    short="${dest#$TARGET/}"
    printf "  (%d/%d) %s  →  %s\n" "$idx" "$TOTAL" "$(basename "$f")" "$short"
  else
    progress_bar "$idx" "$TOTAL"
  fi

  if [ "$DRY_RUN" = false ]; then
    mkdir -p "$dest"
    mv -f "$f" "$dest"
  fi

  case "$cat" in
    AI) cAI=$((cAI + 1)) ;;
    Photos) cPhotos=$((cPhotos + 1)) ;;
    Screenshots) cScreenshots=$((cScreenshots + 1)) ;;
    Unsorted) cUnsorted=$((cUnsorted + 1)) ;;
  esac
done

if [ "$VERBOSE" = false ]; then
  printf "\r  %-60s\n" "Done."
fi

echo
echo "  Results:"
if [ "$FLAT" = false ] && [ -n "$AI_META" ]; then
  echo "    AI (by model/LoRA): $cAI"
else
  echo "    AI: $cAI"
fi
echo "    Photos: $cPhotos"
echo "    Screenshots: $cScreenshots"
echo "    Unsorted: $cUnsorted"
echo

if [ "$DRY_RUN" = true ]; then
  echo "  Dry run complete. Re-run with -e to execute."
fi