#!/usr/bin/env bash
# ============================================================
# batch-clean.sh
# Batch manga text removal (inpainting) using bubbles.json
#
# Usage:
#   ./batch-clean.sh <folder>
#   ./batch-clean.sh <folder> --start 3 --end 7
#
# Output per page lands in:
#   <folder>/translated/<page_stem>/
#     └── <page_stem>_cleaned.png
# ============================================================

set -uo pipefail

# ─────────────────────────────────────────────────────────────
# CONFIGURATION
# ─────────────────────────────────────────────────────────────
START_PAGE=1
END_PAGE=999999
PYTHON_BIN="python"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CLEANER="${SCRIPT_DIR}/clean_bubbles.py"

# ─────────────────────────────────────────────────────────────
# COLOURS
# ─────────────────────────────────────────────────────────────
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
BOLD='\033[1m'
RESET='\033[0m'

# ─────────────────────────────────────────────────────────────
# HELPERS
# ─────────────────────────────────────────────────────────────
usage() {
    echo ""
    echo -e "${BOLD}Usage:${RESET}"
    echo "  $0 <folder> [options]"
    echo ""
    echo -e "${BOLD}Options:${RESET}"
    echo "  --start         First page number     (default: 1)"
    echo "  --end           Last  page number     (default: all)"
    echo "  --python        Python binary         (default: python)"
    echo "  --help,    -h   Show this help"
    echo ""
    echo -e "${BOLD}Examples:${RESET}"
    echo "  $0 pages-for-tests"
    echo "  $0 pages-for-tests --start 3 --end 7"
    echo ""
}

log_info()    { echo -e "${CYAN}ℹ️  $*${RESET}"; }
log_ok()      { echo -e "${GREEN}✅  $*${RESET}"; }
log_warn()    { echo -e "${YELLOW}⚠️  $*${RESET}"; }
log_error()   { echo -e "${RED}❌  $*${RESET}"; }
log_section() {
    echo -e "\n${BOLD}${CYAN}══════════════════════════════════════════${RESET}"
    echo -e "${BOLD}${CYAN}  🧹  $*${RESET}"
    echo -e "${BOLD}${CYAN}══════════════════════════════════════════${RESET}"
}

# ─────────────────────────────────────────────────────────────
# ARGUMENT PARSING
# ─────────────────────────────────────────────────────────────
if [[ $# -eq 0 ]]; then
    log_error "No folder specified."
    usage
    exit 1
fi

FOLDER="$1"
shift

while [[ $# -gt 0 ]]; do
    case "$1" in
        --start)      START_PAGE="$2";  shift 2 ;;
        --end)        END_PAGE="$2";    shift 2 ;;
        --python)     PYTHON_BIN="$2";  shift 2 ;;
        --help|-h)    usage; exit 0 ;;
        *)
            log_error "Unknown option: $1"
            usage
            exit 1
            ;;
    esac
done

# ─────────────────────────────────────────────────────────────
# VALIDATION
# ─────────────────────────────────────────────────────────────
if [[ ! -d "$FOLDER" ]]; then
    log_error "Folder not found: $FOLDER"
    exit 1
fi

if [[ ! -f "$CLEANER" ]]; then
    log_error "clean_bubbles.py not found at: $CLEANER"
    exit 1
fi

if ! command -v "$PYTHON_BIN" &>/dev/null; then
    log_error "Python binary not found: $PYTHON_BIN"
    log_error "Try --python python3"
    exit 1
fi

# ─────────────────────────────────────────────────────────────
# DISCOVER IMAGES
# ─────────────────────────────────────────────────────────────
ALL_IMAGES=()
while IFS= read -r -d '' img; do
    ALL_IMAGES+=("$img")
done < <(
    find "$FOLDER" -maxdepth 1 -type f \
        \( -iname "*.jpg" -o -iname "*.jpeg" \
           -o -iname "*.png" -o -iname "*.webp" \) \
        -print0 | sort -z
)

TOTAL=${#ALL_IMAGES[@]}

if [[ $TOTAL -eq 0 ]]; then
    log_error "No image files found in: $FOLDER"
    exit 1
fi

# ─────────────────────────────────────────────────────────────
# SLICE TO REQUESTED PAGE RANGE
# ─────────────────────────────────────────────────────────────
PAGES=()
for i in "${!ALL_IMAGES[@]}"; do
    PAGE_NUM=$(( i + 1 ))
    if [[ $PAGE_NUM -ge $START_PAGE && $PAGE_NUM -le $END_PAGE ]]; then
        PAGES+=("${ALL_IMAGES[$i]}")
    fi
done

if [[ ${#PAGES[@]} -eq 0 ]]; then
    log_error "No pages in range [${START_PAGE}, ${END_PAGE}] (total: ${TOTAL})"
    exit 1
fi

# ─────────────────────────────────────────────────────────────
# SUMMARY HEADER
# ─────────────────────────────────────────────────────────────
log_section "BATCH MANGA CLEANER"
log_info "📂  Folder   : $(realpath "$FOLDER")"
log_info "📄  Pages    : ${#PAGES[@]} of ${TOTAL} total"
log_info "🔢  Range    : ${START_PAGE} → ${END_PAGE}"
echo ""

# ─────────────────────────────────────────────────────────────
# PROCESS EACH PAGE
# ─────────────────────────────────────────────────────────────
PASS=0
FAIL=0
FAIL_LIST=()

for i in "${!PAGES[@]}"; do
    IMAGE="${PAGES[$i]}"
    PAGE_NUM=$(( START_PAGE + i ))
    STEM="$(basename "${IMAGE%.*}")"
    WORKDIR="${FOLDER}/translated/${STEM}"
    
    echo ""
    echo -e "${BOLD}──────────────────────────────────────────${RESET}"
    echo -e "${BOLD}  🖼️  [${PAGE_NUM}/${TOTAL}]  ${STEM}${RESET}"
    echo -e "${BOLD}──────────────────────────────────────────${RESET}"

    OUTPUT_JSON="${WORKDIR}/bubbles.json"
    OUTPUT_CLEANED="${WORKDIR}/${STEM}_cleaned.png"

    if [[ ! -f "$OUTPUT_JSON" ]]; then
        log_warn "Skipping: bubbles.json not found in ${WORKDIR}"
        FAIL=$(( FAIL + 1 ))
        FAIL_LIST+=("${STEM} (No JSON)")
        continue
    fi

    log_info "🗂️  Image : $(basename "$IMAGE")"
    log_info "🧹  Cleaning text..."

    # ── Run the cleaner ───────────────────────────────────────
    if "$PYTHON_BIN" "$CLEANER" \
            -i "$IMAGE" \
            -j "$OUTPUT_JSON" \
            -o "$OUTPUT_CLEANED"; then

        if [[ -f "$OUTPUT_CLEANED" ]]; then
            log_ok "Cleaned image saved → ${STEM}_cleaned.png"
            PASS=$(( PASS + 1 ))
        else
            log_error "Script ran but output image is missing."
            FAIL=$(( FAIL + 1 ))
            FAIL_LIST+=("${STEM} (Missing Output)")
        fi
    else
        log_error "Page ${PAGE_NUM} FAILED — check output above."
        FAIL=$(( FAIL + 1 ))
        FAIL_LIST+=("${STEM} (Script Error)")
    fi

done

# ─────────────────────────────────────────────────────────────
# FINAL SUMMARY
# ─────────────────────────────────────────────────────────────
log_section "BATCH CLEANING COMPLETE"
echo -e "  ✅  ${GREEN}Passed : ${PASS}${RESET}"
echo -e "  ❌  ${RED}Failed : ${FAIL}${RESET}"

if [[ ${#FAIL_LIST[@]} -gt 0 ]]; then
    echo ""
    log_warn "Failed pages:"
    for NAME in "${FAIL_LIST[@]}"; do
        echo -e "    ❌  ${RED}${NAME}${RESET}"
    done
fi

echo ""
log_info "📦  Output folder: $(realpath "${FOLDER}/translated")"
echo ""

[[ $FAIL -eq 0 ]] && exit 0 || exit 1