#!/usr/bin/env zsh
# 並列バッチOCR（zsh + xargs）
# 依存: img2pdf, ocrmypdf, Tesseract
# 使い方:
#   ./batch_ocr_parallel.zsh ROOT_DIR [OUT_DIR] [-P N] [--recursive]
# 例:
#   ./batch_ocr_parallel.zsh ./ROOT ./out_pdfs -P 4 --recursive

emulate -L zsh
set -euo pipefail

# ---- 引数処理 ----
ROOT="${1:-}"
[[ -z "${ROOT}" ]] && { print -u2 -- "Usage: $0 ROOT_DIR [OUT_DIR] [-P N] [--recursive]"; exit 2; }
shift

OUT="out_pdfs"
PROCS=""
RECURSIVE=0

if (( $# )); then
  # 2番目がディレクトリ名っぽければ OUT に採用
  if [[ -n "${1:-}" && "${1:-}" != "-P" && "${1:-}" != "--recursive" && "${1:-}" != "-r" ]]; then
    OUT="$1"; shift
  fi
  # 残りのオプション
  while (( $# )); do
    case "$1" in
      -P) PROCS="${2:-}"; shift 2 ;;
      --recursive|-r) RECURSIVE=1; shift ;;
      *) print -u2 -- "Unknown arg: $1"; exit 2 ;;
    esac
  done
fi

# 並列度のデフォルト（未指定ならCPUコア数）
if [[ -z "${PROCS}" ]]; then
  if command -v nproc >/dev/null 2>&1; then
    PROCS="$(nproc)"
  elif command -v sysctl >/dev/null 2>&1; then
    PROCS="$(sysctl -n hw.ncpu)"
  else
    PROCS=2
  fi
fi

# 出力先
mkdir -p -- "$OUT"

# ---- ワーカー（サブディレクトリ1つを処理）を一時ファイルに作成 ----
WORKER="$(mktemp -t ocrmypdf_worker.XXXXXX).zsh"
cat > "$WORKER" <<'ZSH'
#!/usr/bin/env zsh
emulate -L zsh
set -euo pipefail
setopt null_glob

ROOT="$1"
OUT="$2"
DIR="$3"

name="${${DIR#${ROOT%/}/}#/}"       # ROOT/ 以降を相対名に（深い場合は "A/B"）
base="${DIR:t}"                      # 末尾名（単純名）

# 画像拡張子
typeset -a exts; exts=( jpg jpeg png tif tiff JPG JPEG PNG TIF TIFF )

# 画像列挙
typeset -a imgs; imgs=()
for ext in $exts; do
  imgs+=("$DIR"/*.${ext})
done

if (( ${#imgs} == 0 )); then
  print -u2 -- "[skip] ${name} (no images)"
  exit 0
fi

# 一時PDF
tmp_pdf="$(mktemp -t "ocr_${base}.XXXXXX").pdf"

# 1) 画像 → 無劣化PDF結合（-o を先に、-- の後に入力）
img2pdf --auto-orient -o "$tmp_pdf" -- "${imgs[@]}"

# 出力ファイル名（ネストを __ でつなぐ）
out_pdf="$OUT/${name//\//__}.pdf"
mkdir -p -- "${out_pdf:h}"

# 2) OCR（タイプ打ち英語向けチューニング）
ocrmypdf \
  -l eng \
  --tesseract-oem 1 \
  --tesseract-pagesegmode 6 \
  --optimize 1 \
  --output-type pdf \
  "$tmp_pdf" "$out_pdf"

rm -f -- "$tmp_pdf"
print -r -- "✅ Wrote: $out_pdf"
ZSH
chmod +x "$WORKER"

# ---- 対象ディレクトリ列挙 → xargs 並列実行 ----
if (( RECURSIVE )); then
  # ROOT 配下の全サブディレクトリ（ROOT 自身は除外）
  find "$ROOT" -mindepth 1 -type d -print0 \
    | xargs -0 -I {} -n 1 -P "$PROCS" "$WORKER" "$ROOT" "$OUT" {}
else
  # 直下のサブディレクトリのみ
  find "$ROOT" -mindepth 1 -maxdepth 1 -type d -print0 \
    | xargs -0 -I {} -n 1 -P "$PROCS" "$WORKER" "$ROOT" "$OUT" {}
fi

print -r -- "Batch complete. (parallel: $PROCS)"