#!/usr/bin/env bash
# ocr-scan.sh — recursively OCR un-OCR'd PDFs
# Deps: ocrmypdf, poppler-utils (pdftotext), tesseract-ocr
# Optional: python3-langdetect (for -a auto-detect)

set -euo pipefail

LANG_CODE="eng"
AUTO_DETECT=false
MIN_CHARS=50        # chars per page threshold — below = needs OCR
DIR="."

usage() {
    cat <<EOF
Usage: $(basename "$0") [OPTIONS] [DIRECTORY]

Recursively find un-OCR'd PDFs and produce OCR'd copies with _OCR suffix.

OPTIONS:
  -l LANG     Tesseract language code(s)  (default: eng)
              Single:       -l pol
              Multiple:     -l pol+eng
              Examples:     eng pol tha deu fra rus chi_sim jpn ara
  -a          Auto-detect language via langdetect (requires: pip install langdetect)
              Note: falls back to eng for fully un-OCR'd scans (no text layer)
  -h          Show this help and exit

ARGUMENTS:
  DIRECTORY   Root directory to search recursively (default: current directory)

OUTPUT:
  Each processed file is saved alongside the original:
    original.pdf → original_OCR.pdf
  Files already having an _OCR copy or sufficient text are skipped.

DEPENDENCIES:
  Required:   ocrmypdf, poppler-utils (pdftotext, pdfinfo), tesseract-ocr
  Per-lang:   apt install tesseract-ocr-pol tesseract-ocr-tha ...
  Optional:   pip install langdetect  (for -a flag)

EXAMPLES:
  $(basename "$0")                        # current dir, English
  $(basename "$0") ~/docs                 # ~/docs recursively, English
  $(basename "$0") -l pol ~/docs          # Polish
  $(basename "$0") -l pol+eng ~/docs      # Polish + English mixed
  $(basename "$0") -a ~/docs             # auto-detect language
EOF
    exit 0
}

while getopts ":l:ah" opt; do
    case $opt in
        l) LANG_CODE="$OPTARG" ;;
        a) AUTO_DETECT=true ;;
        h) usage ;;
        :) echo "Option -$OPTARG requires an argument."; usage ;;
        \?) echo "Unknown option: -$OPTARG"; usage ;;
    esac
done
shift $((OPTIND - 1))
[[ $# -gt 0 ]] && DIR="$1"

# Map ISO 639-1 → Tesseract lang codes
map_lang() {
    case "$1" in
        pl|pol) echo "pol" ;;
        en|eng) echo "eng" ;;
        th|tha) echo "tha" ;;
        de|deu) echo "deu" ;;
        fr|fra) echo "fra" ;;
        ru|rus) echo "rus" ;;
        zh|chi) echo "chi_sim" ;;
        ja|jpn) echo "jpn" ;;
        ar|ara) echo "ara" ;;
        *)       echo "eng" ;;
    esac
}

detect_lang() {
    local pdf="$1"
    local sample
    sample=$(pdftotext "$pdf" - 2>/dev/null | head -c 500)
    if [[ -z "$sample" ]]; then
        echo "eng"; return
    fi
    python3 - <<EOF 2>/dev/null || echo "eng"
from langdetect import detect
try:
    print("$(echo "$sample" | tr -d "'" | head -c 300)")
except:
    print("eng")
EOF
    # simpler: call python directly
    python3 -c "
from langdetect import detect
import sys
try:
    lang = detect('''$sample''')
    print(lang)
except:
    print('eng')
" 2>/dev/null || echo "eng"
}

needs_ocr() {
    local pdf="$1"
    local pages
    pages=$(pdfinfo "$pdf" 2>/dev/null | awk '/^Pages:/{print $2}')
    [[ -z "$pages" || "$pages" -eq 0 ]] && pages=1
    local chars
    chars=$(pdftotext "$pdf" - 2>/dev/null | wc -c)
    local threshold=$(( pages * MIN_CHARS ))
    [[ "$chars" -lt "$threshold" ]]
}

find "$DIR" -type f -iname "*.pdf" ! -iname "*_OCR.pdf" | while read -r pdf; do
    out="${pdf%.pdf}_OCR.pdf"

    if [[ -f "$out" ]]; then
        echo "[SKIP] Already has OCR copy: $out"
        continue
    fi

    if ! needs_ocr "$pdf"; then
        echo "[SKIP] Already has text: $pdf"
        continue
    fi

    lang="$LANG_CODE"
    if $AUTO_DETECT; then
        detected=$(detect_lang "$pdf")
        lang=$(map_lang "$detected")
        echo "[DETECT] $pdf → $detected → $lang"
    fi

    echo "[OCR] $pdf (lang: $lang) → $out"
    ocrmypdf \
        --language "$lang" \
        --output-type pdfa \
        --skip-text \
        --rotate-pages \
        --deskew \
        "$pdf" "$out" \
        && echo "[DONE] $out" \
        || echo "[FAIL] $pdf"
done