#!/bin/bash
# organize-images -- Sort images into AI/, Photos/, Screenshots/, Unsorted/
# AI files route into AI/{model}/{lora}/ via extract-ai-meta (canonical scheme).
#
# Modes:
#   (default)  Sort top-level files of [path] into the 4 buckets.
#   --resort   Recursively re-sort [path] (the Images root containing AI/),
#              moving any file with ComfyUI metadata to its canonical
#              AI/{model}/{lora}/ location. Filenames are never changed.
#
# Generic ComfyUI_* filenames are renamed first via rename-ai-snaps
# (--only-generic) before sorting, when that script is available.

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

usage() {
  echo "Usage: $(basename "$0") [path] [-e|--execute] [-v|--verbose] [-i|--install] [--flat] [--resort] [--limit N]"
  echo "  path       Directory to organize (default: ~/Pictures)"
  echo "  -e         Execute moves (default: dry-run only)"
  echo "  -v         Verbose: show per-file classification/moves"
  echo "  -i         Install this script to ~/.local/bin/"
  echo "  --flat     Don't sub-categorize by model/LoRA (flat AI/ folder)"
  echo "  --resort   Recursively re-sort [path] into canonical AI/{model}/{lora}/"
  echo "             ([path] must be the Images root containing AI/)"
  echo "  --limit N  Only process N files (for batching; default mode only)"
  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 ;;
    --resort) RESORT=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

# --- Resolve helper scripts (same dir as this script, else PATH) ---
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd -P)"
find_tool() {
  local n="$1"
  if [ -x "$SCRIPT_DIR/$n" ]; then
    echo "$SCRIPT_DIR/$n"
  elif command -v "$n" >/dev/null 2>&1; then
    command -v "$n"
  else
    return 1
  fi
}

META_SCRIPT="$(find_tool extract-ai-meta)" || {
  echo "Error: extract-ai-meta not found (looked next to $0 and in PATH)" >&2
  exit 1
}
RENAME_SCRIPT="$(find_tool rename-ai-snaps || true)"

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

# --- Resort mode: recursive canonical re-sort of an existing tree ---
if [ "$RESORT" = true ]; then
  if [ ! -d "$TARGET/AI" ]; then
    echo "Error: --resort expects the Images root (a directory containing AI/)" >&2
    exit 1
  fi

  # Generic-name hook: rename ComfyUI_* files in place before sorting
  GENERIC_COUNT=$(find "$TARGET" -type f -name 'ComfyUI*.png' 2>/dev/null | wc -l | tr -d ' ')
  if [ "$GENERIC_COUNT" -gt 0 ]; then
    if [ -n "$RENAME_SCRIPT" ]; then
      if [ "$DRY_RUN" = true ]; then
        echo "  $GENERIC_COUNT generic ComfyUI_*.png file(s) would be renamed first (rename-ai-snaps --only-generic)"
      else
        echo "  Renaming $GENERIC_COUNT generic ComfyUI_*.png file(s) first..."
        "$RENAME_SCRIPT" "$TARGET" -r -n --only-generic
      fi
    else
      echo "  Note: $GENERIC_COUNT generic ComfyUI_*.png file(s) found; rename-ai-snaps not available, keeping names"
    fi
  fi

  # Metadata manifest (recursive)
  echo "  Scanning AI metadata (recursive)..."
  AI_META=$(mktemp)
  if ! "$META_SCRIPT" "$TARGET" -r > "$AI_META"; then
    echo "Error: AI metadata scan failed" >&2
    rm -f "$AI_META"
    exit 1
  fi

  python3 - "$TARGET" "$AI_META" "$DRY_RUN" "$VERBOSE" <<'PYEOF'
import json, os, shutil, sys

root = os.path.abspath(sys.argv[1])
meta = json.load(open(sys.argv[2]))
dry = sys.argv[3] == "true"
verbose = sys.argv[4] == "true"

IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp"}
LEGACY_DIRS = {"None", "unknown"}

moved = ok = nometa = collisions = errors = 0
samples = []
legacy_no_meta = []
planned = set()

for dirpath, dirnames, filenames in os.walk(root):
    dirnames.sort()
    for fn in sorted(filenames):
        if os.path.splitext(fn)[1].lower() not in IMAGE_EXTS:
            continue
        src = os.path.join(dirpath, fn)
        rel = os.path.relpath(src, root)
        entry = meta.get(rel)
        if not entry:
            nometa += 1
            top = rel.split(os.sep)[0]
            if top in LEGACY_DIRS:
                legacy_no_meta.append(rel)
            continue

        dest_dir = os.path.join(root, "AI", entry["dest"])
        if os.path.abspath(dest_dir) == os.path.abspath(dirpath):
            ok += 1
            continue

        dest = os.path.join(dest_dir, fn)
        taken = planned if dry else None
        if os.path.exists(dest) or (taken is not None and dest in taken):
            stem, ext = os.path.splitext(fn)
            c = 1
            while os.path.exists(os.path.join(dest_dir, f"{stem}_{c}{ext}")) or \
                  (taken is not None and os.path.join(dest_dir, f"{stem}_{c}{ext}") in taken):
                c += 1
            dest = os.path.join(dest_dir, f"{stem}_{c}{ext}")
            collisions += 1
        if taken is not None:
            taken.add(dest)

        dest_rel = os.path.relpath(dest, root)
        if dry:
            moved += 1
            if len(samples) < 20:
                samples.append((rel, dest_rel))
            if verbose:
                print(f"  WOULD MOVE  {rel}  →  {dest_rel}")
        else:
            try:
                os.makedirs(dest_dir, exist_ok=True)
                shutil.move(src, dest)
                moved += 1
                if verbose:
                    print(f"  MOVED  {rel}  →  {dest_rel}")
                elif moved % 200 == 0:
                    print(f"\r  …moved {moved}", end="", flush=True)
            except OSError as e:
                errors += 1
                print(f"  SKIPPED ({e}): {rel}")

if not verbose and not dry and moved:
    print(f"\r  Moved {moved} file(s).          ")

# Prune dirs left empty by the moves (keep top-level buckets)
pruned = 0
if not dry and moved:
    keep = {"AI", "Photos", "Screenshots", "Unsorted"}
    for dirpath, dirnames, filenames in os.walk(root, topdown=False):
        if dirpath == root:
            continue
        if os.path.relpath(dirpath, root).split(os.sep)[0] in keep and \
           os.path.dirname(dirpath) == root:
            continue
        try:
            os.rmdir(dirpath)  # only succeeds when empty
            pruned += 1
        except OSError:
            pass

print()
print(f"  {'DRY RUN — would move' if dry else 'Moved'}: {moved}")
print(f"  Already correctly placed: {ok}")
print(f"  Collisions (renamed with _N suffix): {collisions}")
if errors:
    print(f"  Errors: {errors}")
print(f"  No AI metadata (left in place): {nometa}")
if pruned:
    print(f"  Empty folders pruned: {pruned}")

if samples:
    print()
    print("  Sample moves:")
    for s, d in samples:
        print(f"    {s}")
        print(f"      → {d}")

if legacy_no_meta:
    print()
    print(f"  Note: {len(legacy_no_meta)} file(s) in None//unknown/ have no AI metadata")
    print("  and were left in place for manual triage (not all are AI images).")
    for rel in legacy_no_meta[:10]:
        print(f"    {rel}")
    if len(legacy_no_meta) > 10:
        print(f"    … and {len(legacy_no_meta) - 10} more")
print()
PYEOF
  STATUS=$?
  rm -f "$AI_META"
  exit $STATUS
fi

# --- Default mode: classify top-level files into AI/Photos/Screenshots/Unsorted ---

# Generic-name hook (non-recursive)
GENERIC_COUNT=$(find "$TARGET" -maxdepth 1 -type f -name 'ComfyUI*.png' 2>/dev/null | wc -l | tr -d ' ')
if [ "$GENERIC_COUNT" -gt 0 ] && [ -n "$RENAME_SCRIPT" ]; then
  if [ "$DRY_RUN" = true ]; then
    echo "  $GENERIC_COUNT generic ComfyUI_*.png file(s) would be renamed first (rename-ai-snaps --only-generic)"
  else
    echo "  Renaming $GENERIC_COUNT generic ComfyUI_*.png file(s) first..."
    "$RENAME_SCRIPT" "$TARGET" -n --only-generic
  fi
fi

# Pre-scan: extract AI metadata for model/LoRA routing
AI_META=""
if [ "$FLAT" = false ]; then
  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"
    AI_META=""
  elif [ ! -s "$AI_META" ]; then
    FLAT=true
    rm -f "$AI_META"
    AI_META=""
  fi
fi

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

# Basenames of files with AI metadata (manifest presence = is an AI image)
AI_KEYS=""
if [ -n "$AI_META" ]; then
  AI_KEYS=$(mktemp)
  python3 -c "
import json, sys
for k in json.load(open('$AI_META')):
    print(k)
" > "$AI_KEYS" 2>/dev/null
fi

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

  # Files with ComfyUI metadata are AI regardless of name
  if [ -n "$AI_KEYS" ] && grep -Fxq "$name" "$AI_KEYS" 2>/dev/null; then
    echo "AI"; return
  fi

  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

  echo "Unsorted"
}

# --- Collect top-level 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."
  [ -n "$AI_META" ] && rm -f "$AI_META"
  [ -n "$AI_KEYS" ] && rm -f "$AI_KEYS"
  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"
}

# Destination for an AI file: canonical AI/{model}/{lora}/ from the manifest
ai_dest() {
  AI_META_PATH="$AI_META" python3 -c "
import json, os, sys
m = json.load(open(os.environ['AI_META_PATH']))
entry = m.get(sys.argv[1])
if not entry or not isinstance(entry, dict) or not entry.get('dest'):
    print('AI/unknown')
else:
    print('AI/' + entry['dest'])
" "$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"
    base=$(basename "$f")
    final="$dest/$base"
    if [ -e "$final" ]; then
      stem="${base%.*}"; fext="${base##*.}"
      n=1
      while [ -e "$dest/${stem}_${n}.${fext}" ]; do n=$((n + 1)); done
      final="$dest/${stem}_${n}.${fext}"
    fi
    mv "$f" "$final"
  fi

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

[ -n "$AI_META" ] && rm -f "$AI_META"
[ -n "$AI_KEYS" ] && rm -f "$AI_KEYS"

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

echo
echo "  Results:"
if [ "$FLAT" = false ] && [ "$cAI" -gt 0 ]; 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