#!/bin/bash
# LiDAR archaeology pipeline
# Usage: ./run.sh [options]
#
# Options:
#   -r RESOLUTION   Resolution in m/px, or several separated by commas
#                   (run.sh default: 0.5, e.g. 0.5,0.2; the CLI alone defaults to 0.2)
#   -w WORKERS      Number of parallel workers (run.sh default: 1)
#   -g [LIST]       Select GPU(s): -g (auto), -g 0, -g 0,2, -g all
#   -v              Verbose mode
#   --debug         Debug mode
#   -f / --force    Regenerate all files
#   --keep-tif      Keep the TIFF files
#   --force-classification
#   --ground-classification {auto,ign,smrf,csf}   (default: ign)
#   --ign-classes CLASSES
#                   LAS classes for the IGN DTM (e.g. sol,unclassified; default: sol)
#   --quality N     Image quality 1-100 (default: 60)
#   --lossless      Lossless compression
#   --format FMT    Output format: avif (default) or webp
#   --only VIZ...   Generate only these visualizations
#   --skip VIZ...   Exclude these visualizations
#   --file NAME...  Process one or more LAZ files
#   --fetch-tiles COL,ROW...
#                   Download these tiles from IGN, then process them
#   --rebuild-index Rebuild only the tile catalogue (no reprocessing)
#   --no-index      Do not build the catalogue
#   --test          Run the unit tests
#   --serve [PORT]  Map + XYZ tiles + generation (full image, default: 8973)
#   --serve-maps [PORT]
#                   Map container only (lightweight lidar-maps image): tiles
#                   served from the local cache, generation delegated via
#                   LIDAR_GENERATION_URL, tiles fetched via LIDAR_SOURCE_URL
#   -h              Show the full help

set -e

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
INPUT_DIR="${SCRIPT_DIR}/input"
OUTPUT_DIR="${SCRIPT_DIR}/output"
IMAGE_NAME="lidar-lidar"

show_help() {
    echo "LiDAR archaeology pipeline"
    echo ""
    echo "Usage: $0 [options]"
    echo ""
    echo "Options:"
    echo "  -r RESOLUTION   Resolution in m/px, or several separated by commas"
    echo "                  (default: 0.5, e.g. 0.5,0.2; the CLI alone defaults to 0.2)"
    echo "  -w WORKERS      Number of parallel CPU workers (default: 1)"
    echo "  -g [LIST]       Select GPU(s): -g (auto), -g 0, -g 0,2, -g all"
    echo "  -v              Verbose mode (timestamps + levels)"
    echo "  --debug         Debug mode (internal file:line details)"
    echo "  -f / --force    Regenerate all files even if the AVIF images exist"
    echo "  --force-classification"
    echo "                  Reclassify the ground even if the .las file exists"
    echo "  --keep-tif      Keep the TIFFs to regenerate the AVIF images"
    echo "  --ground-classification {auto,ign,smrf,csf}"
    echo "                  Ground classification method (default: ign = IGN pre-classification;"
    echo "                  auto = IGN if present, otherwise SMRF/CSF detection)"
    echo "  --ign-classes CLASSES"
    echo "                  LAS classes for the IGN DTM: sol,unclassified,eau,virtuel,pont,sursol"
    echo "                  or codes (e.g. 2,1). Pure mode, no retouching. (default: sol)"
    echo "  --quality N     Image quality 1-100 (default: 60, 100 = lossless)"
    echo "  --lossless      Lossless compression (same as --quality 100)"
    echo "  --format FMT    Output format: avif (default) or webp"
    echo "  --only VIZ...   Generate only these visualizations"
    echo "  --skip VIZ...   Exclude these visualizations"
    echo "  --file NAME...  Process one or more LAZ files (full name without .laz/.las)"
    echo "  --fetch-tiles COL,ROW..."
    echo "                  Download these tiles from IGN, then process them"
    echo "  --rebuild-index Rebuild only the tile catalogue (no reprocessing)"
    echo "  --no-index      Do not build the catalogue (thumbnails + inventory)"
    echo "  --serve [PORT]  Map + XYZ tiles + generation API (full image,"
    echo "                  default 8973) - see docker-compose.yml (serve service)"
    echo "  --serve-maps [PORT]"
    echo "                  XYZ tile map (lightweight lidar-maps image, default 8975):"
    echo "                  OpenStreetMap scheme, reusable in JOSM/QGIS/iD;"
    echo "                  generation delegated via LIDAR_GENERATION_URL"
    echo "                  (see docs/MAPS.md)"
    echo "  --test          Run the unit tests"
    echo "  -h              Show this help"
    echo ""
    echo "Available visualizations (default without --only/--skip: relief_oriente densite_sol):"
    echo "  relief_oriente densite_sol hillshade slope aspect mslrm sailore"
    echo "  pos_open neg_open svf roughness wavelet flow_acc solar anomaly"
    echo "  ortho topo"
    echo ""
    echo "Examples:"
    echo "  $0 -g                                    # Auto GPU (first available)"
    echo "  $0 -g 0                                  # GPU index 0"
    echo "  $0 -g 0,2                                # GPUs 0 and 2"
    echo "  $0 -g all                                # All GPUs"
    echo "  $0 -g -w 4                               # GPU + 4 workers"
    echo "  $0 -g -v                                 # GPU + verbose"
    echo "  $0 -g -r 0.2                             # High resolution"
    echo "  $0 -g -r 0.5,0.2                         # Multi-resolution (0.5 m + 0.2 m)"
    echo "  $0 -g --force                            # Regenerate AVIF (DTM kept with --keep-tif)"
    echo "  $0 -g --force-classification             # Reclassify the ground only"
    echo "  $0 -g --only hillshade svf               # Only 2 visualizations"
    echo "  $0 -g --lossless                         # Lossless AVIF"
    echo "  $0 -g --quality 90                       # AVIF quality 90"
    echo "  $0 -g --ground-classification csf        # Force CSF (complex terrain)"
    echo "  $0 -g --file LHD_...IGN69.copc           # A single file"
}

# Show the help when called without arguments
if [ $# -eq 0 ]; then
    show_help
    exit 0
fi

RESOLUTION=0.5
WORKERS=1
GPU_FLAG=""
GPU_ARG=""
VERBOSE_FLAG=""
FORCE_FLAG=""
FILE_ARGS=""
GROUND_METHOD=""
IGN_CLASSES_FLAG=""
FORCE_CLASSIFY_FLAG=""
KEEP_TIF_FLAG=""
QUALITY=""
FORMAT_FLAG=""
ONLY_FLAG=""
SKIP_FLAG=""
TEST_FLAG=0
REBUILD_INDEX_FLAG=""
NO_INDEX_FLAG=""
FETCH_TILES_ARGS=""
SERVE_FLAG=0
SERVE_PORT=8973
MAPS_FLAG=0
MAPS_SERVE_PORT=8975
MAPS_IMAGE="lidar-maps"

# Parse arguments manually (more robust than getopts for mixed short/long options)
while [ $# -gt 0 ]; do
    case "$1" in
        -r) RESOLUTION="$2"; shift 2 ;;
        -w) WORKERS="$2"; shift 2 ;;
        -g)
            GPU_FLAG="--gpus all"
            # Check if next arg is a GPU selector (not another flag)
            if [ $# -gt 1 ] && [[ ! "$2" =~ ^- ]]; then
                GPU_ARG="$2"
                GPU_FLAG="--gpus all"
                shift 2
            else
                shift
            fi
            ;;
        -v) VERBOSE_FLAG="-v"; shift ;;
        -f) FORCE_FLAG="--force"; shift ;;
        --debug) VERBOSE_FLAG="--debug"; shift ;;
        --force) FORCE_FLAG="--force"; shift ;;
        --force-classification) FORCE_CLASSIFY_FLAG="--force-classification"; shift ;;
        --keep-tif) KEEP_TIF_FLAG="--keep-tif"; shift ;;
        --ground-classification) GROUND_METHOD="$2"; shift 2 ;;
        --ground-classification=*) GROUND_METHOD="${1#--ground-classification=}"; shift ;;
        --ign-classes) IGN_CLASSES_FLAG="--ign-classes $2"; shift 2 ;;
        --ign-classes=*) IGN_CLASSES_FLAG="--ign-classes=${1#--ign-classes=}"; shift ;;
        --quality) QUALITY="--quality $2"; shift 2 ;;
        --lossless) QUALITY="--lossless"; shift ;;
        --format) FORMAT_FLAG="--format $2"; shift 2 ;;
        --only) shift; ONLY_FLAG="--only"; while [ $# -gt 0 ] && [[ ! "$1" =~ ^- ]]; do ONLY_FLAG="$ONLY_FLAG $1"; shift; done ;;
        --skip) shift; SKIP_FLAG="--skip"; while [ $# -gt 0 ] && [[ ! "$1" =~ ^- ]]; do SKIP_FLAG="$SKIP_FLAG $1"; shift; done ;;
        --file) shift; while [ $# -gt 0 ] && [[ ! "$1" =~ ^- ]]; do FILE_ARGS="$FILE_ARGS $1"; shift; done ;;
        --rebuild-index) REBUILD_INDEX_FLAG="--rebuild-index"; shift ;;
        --no-index) NO_INDEX_FLAG="--no-index"; shift ;;
        --fetch-tiles) shift; while [ $# -gt 0 ] && [[ ! "$1" =~ ^- ]]; do FETCH_TILES_ARGS="$FETCH_TILES_ARGS $1"; shift; done ;;
        --test) TEST_FLAG=1; shift ;;
        --serve)
            SERVE_FLAG=1; shift
            if [ $# -gt 0 ] && [[ ! "$1" =~ ^- ]]; then SERVE_PORT="$1"; shift; fi
            ;;
        --serve-maps)
            MAPS_FLAG=1; shift
            if [ $# -gt 0 ] && [[ ! "$1" =~ ^- ]]; then MAPS_SERVE_PORT="$1"; shift; fi
            ;;
        -h|--help|-help)
            show_help
            exit 0
            ;;
        *) echo "Invalid option: $1" >&2; exit 1 ;;
    esac
done

# Trim FILE_ARGS whitespace
FILE_ARGS=$(echo "$FILE_ARGS" | xargs)
FETCH_TILES_ARGS=$(echo "$FETCH_TILES_ARGS" | xargs)

# input/ is read-only unless --fetch-tiles (IGN download into input/)
INPUT_RO=":ro"
if [ -n "$FETCH_TILES_ARGS" ]; then
    INPUT_RO=""
fi

# Always build the image before any docker run. The Docker cache makes this
# almost instant when the code has not changed; on failure the full build log
# is printed. This avoids running tests/server/pipeline on a stale image (the
# code is baked in at build time).
build_image() {
    local log="/tmp/lidar_build_$$.log"
    echo "Updating the Docker image (cached)..."
    if ! docker build -t "$IMAGE_NAME" "$SCRIPT_DIR" >"$log" 2>&1; then
        cat "$log" >&2
        echo "Image build failed (log: $log)" >&2
        exit 1
    fi
    rm -f "$log"
}

# XYZ tile map image (see Dockerfile.maps)
build_maps_image() {
    local log="/tmp/lidar_maps_build_$$.log"
    echo "Updating the map Docker image (cached)..."
    if ! docker build -f Dockerfile.maps -t "$MAPS_IMAGE" "$SCRIPT_DIR" >"$log" 2>&1; then
        cat "$log" >&2
        echo "Map image build failed (log: $log)" >&2
        exit 1
    fi
    rm -f "$log"
}

# Check for --test flag first
if [ "$TEST_FLAG" -eq 1 ]; then
    build_image
    echo "============================================"
    echo " LiDAR pipeline unit tests"
    echo "============================================"
    docker run --rm --init $GPU_FLAG \
        "$IMAGE_NAME" \
        python3 -m pytest -v --pyargs lidar_pipeline.tests
    exit $?
fi

build_image

# Create the directories if they do not exist
mkdir -p "$INPUT_DIR" "$OUTPUT_DIR"

# --serve mode: XYZ tile map + generation API (full image, input/ writable
# for the IGN download). Compose equivalent:
# docker compose up -d --build serve (port 8973).
if [ "$SERVE_FLAG" -eq 1 ]; then
    echo "============================================"
    echo " LiDAR map - XYZ tiles + generation"
    echo "============================================"
    echo "  http://127.0.0.1:${SERVE_PORT}/"
    echo "  GPU        : $([ -n "$GPU_FLAG" ] && echo 'yes' || echo 'no')"
    echo "  \"+ Area\" button: draw a rectangle -> IGN download + rendering"
    echo "============================================"
    GPU_ENV_FLAG=""
    [ -n "$GPU_FLAG" ] && GPU_ENV_FLAG="-e LIDAR_GPU=1"
    if [ -n "$GPU_ARG" ] && [ "$GPU_ARG" != "all" ]; then
        GPU_ENV_FLAG="-e LIDAR_GPU=1 -e CUDA_VISIBLE_DEVICES=${GPU_ARG} --gpus '""device=${GPU_ARG}""'"
    fi
    # Network allowed to start generations (/api/generate): taken from the
    # host environment variable LIDAR_REGEN_CIDR; without it, run.sh passes
    # the server default (localhost + RFC1918). An EMPTY string lifts the
    # restriction entirely.
    REGEN_ENV_FLAG="-e LIDAR_REGEN_CIDR=${LIDAR_REGEN_CIDR-127.0.0.0/8,::1,10.0.0.0/8,172.16.0.0/12,192.168.0.0/16}"
    exec docker run --rm --init $GPU_FLAG \
        --user 1000:1000 \
        -p "${SERVE_PORT}:8973" \
        -v "${INPUT_DIR}:/data/input" \
        -v "${OUTPUT_DIR}:/data/output" \
        -e LIDAR_INPUT_DIR=/data/input -e LIDAR_OUTPUT_DIR=/data/output \
        -e LIDAR_PORT=8973 \
        $GPU_ENV_FLAG $REGEN_ENV_FLAG \
        "$IMAGE_NAME" \
        python3 -m uvicorn lidar_pipeline.mapserve:app --host 0.0.0.0 --port 8973
fi

# --serve-maps mode: XYZ tile map (lightweight lidar-maps image).
# Tiles follow the OpenStreetMap scheme ({z}/{x}/{y}.png): the map can be
# browsed here and the same URLs work as imagery background in JOSM, iD,
# QGIS or uMap. Optional variables from the host environment:
#   LIDAR_MAPS_URL          upstream tile server (remote map)
#   LIDAR_SOURCE_URL        upstream source-tile server (/api/tiles inventory)
#   LIDAR_GENERATION_URL    processing machine for delegated generation
#   LIDAR_REMOTE_TOKEN      token for calls to the processing machine
#   LIDAR_REGEN_CIDR        network allowed to start generations
#   LIDAR_ATTRIBUTION       attribution text served to clients
if [ "$MAPS_FLAG" -eq 1 ]; then
    build_maps_image
    mkdir -p "$OUTPUT_DIR"
    echo "============================================"
    echo " LiDAR map - XYZ tiles (lightweight image)"
    echo "============================================"
    echo "  Map            : http://127.0.0.1:${MAPS_SERVE_PORT}/"
    echo "  XYZ tiles      : http://127.0.0.1:${MAPS_SERVE_PORT}/tiles/{layer}/{z}/{x}/{y}.png"
    echo "  JOSM / QGIS    : /tiles/josm.imagery.xml, /tiles/wmts.xml"
    echo "  Upstream tiles : ${LIDAR_MAPS_URL:-none (local rendering)}"
    echo "  Upstream data  : ${LIDAR_SOURCE_URL:-none (local cache only)}"
    echo "  Generation     : ${LIDAR_GENERATION_URL:-local if the pipeline is present}"
    echo "============================================"
    MAPS_ENV=()
    [ -n "${LIDAR_MAPS_URL:-}" ] && MAPS_ENV+=(-e LIDAR_MAPS_URL="$LIDAR_MAPS_URL")
    [ -n "${LIDAR_SOURCE_URL:-}" ] && MAPS_ENV+=(-e LIDAR_SOURCE_URL="$LIDAR_SOURCE_URL")
    [ -n "${LIDAR_GENERATION_URL:-}" ] && MAPS_ENV+=(-e LIDAR_GENERATION_URL="$LIDAR_GENERATION_URL")
    [ -n "${LIDAR_REMOTE_TOKEN:-}" ] && MAPS_ENV+=(-e LIDAR_REMOTE_TOKEN="$LIDAR_REMOTE_TOKEN")
    [ -n "${LIDAR_REGEN_CIDR:-}" ] && MAPS_ENV+=(-e LIDAR_REGEN_CIDR="$LIDAR_REGEN_CIDR")
    [ -n "${LIDAR_ATTRIBUTION:-}" ] && MAPS_ENV+=(-e LIDAR_ATTRIBUTION="$LIDAR_ATTRIBUTION")
    exec docker run --rm --init \
        --user 1000:1000 \
        -p "${MAPS_SERVE_PORT}:8975" \
        -v "${OUTPUT_DIR}:/data/output" \
        -e LIDAR_OUTPUT_DIR=/data/output \
        "${MAPS_ENV[@]}" \
        "$MAPS_IMAGE"
fi

# Run the pipeline
echo "============================================"
echo " LiDAR archaeology pipeline"
echo "============================================"
echo "  Resolution : ${RESOLUTION} m/px"
echo "  Workers    : ${WORKERS}"
echo "  GPU        : $([ -n "$GPU_FLAG" ] && echo "OUI${GPU_ARG:+ ($GPU_ARG)}" || echo 'non')"
echo "  Verbose    : $([ -n "$VERBOSE_FLAG" ] && echo 'yes' || echo 'no')"
echo "  Force      : $([ -n "$FORCE_FLAG" ] && echo 'yes' || echo 'no')"
echo "  Force classif. : $([ -n "$FORCE_CLASSIFY_FLAG" ] && echo 'yes' || echo 'no')"
echo "  Keep TIFF    : $([ -n "$KEEP_TIF_FLAG" ] && echo 'yes' || echo 'no')"
echo "  Image quality : $([ -n "$QUALITY" ] && echo "$QUALITY" || echo '60')"
echo "  Format      : $([ -n "$FORMAT_FLAG" ] && echo "${FORMAT_FLAG#--format }" || echo 'avif')"
echo "  Ground classification : $([ -n "$GROUND_METHOD" ] && echo "$GROUND_METHOD" || echo 'ign')$([ -n "$IGN_CLASSES_FLAG" ] && echo " ${IGN_CLASSES_FLAG#--ign-classes }")"
if [ -n "$ONLY_FLAG" ]; then
    echo "  Visualizations: only${ONLY_FLAG#--only}"
elif [ -n "$SKIP_FLAG" ]; then
    echo "  Visualizations: all except${SKIP_FLAG#--skip}"
fi
if [ -n "$FILE_ARGS" ]; then
    echo "  Files      :${FILE_ARGS}"
fi
if [ -n "$FETCH_TILES_ARGS" ]; then
    echo "  IGN tiles  :${FETCH_TILES_ARGS} (input/ mounted writable)"
fi
echo "============================================"

CMD_ARGS="-o /data/output -r $RESOLUTION -w $WORKERS $VERBOSE_FLAG $FORCE_FLAG $FORCE_CLASSIFY_FLAG $KEEP_TIF_FLAG $QUALITY $FORMAT_FLAG"
if [ -n "$GPU_ARG" ]; then
    CMD_ARGS="$CMD_ARGS -g $GPU_ARG"
fi
if [ -z "$GPU_ARG" ] && [ -n "$GPU_FLAG" ]; then
    CMD_ARGS="$CMD_ARGS -g"
fi
if [ -n "$GROUND_METHOD" ]; then
    CMD_ARGS="$CMD_ARGS --ground-classification $GROUND_METHOD"
fi
if [ -n "$IGN_CLASSES_FLAG" ]; then
    CMD_ARGS="$CMD_ARGS $IGN_CLASSES_FLAG"
fi
if [ -n "$ONLY_FLAG" ]; then
    CMD_ARGS="$CMD_ARGS $ONLY_FLAG"
fi
if [ -n "$SKIP_FLAG" ]; then
    CMD_ARGS="$CMD_ARGS $SKIP_FLAG"
fi
if [ -n "$FILE_ARGS" ]; then
    CMD_ARGS="$CMD_ARGS --file $FILE_ARGS"
fi
if [ -n "$REBUILD_INDEX_FLAG" ]; then
    CMD_ARGS="$CMD_ARGS $REBUILD_INDEX_FLAG"
fi
if [ -n "$NO_INDEX_FLAG" ]; then
    CMD_ARGS="$CMD_ARGS $NO_INDEX_FLAG"
fi
if [ -n "$FETCH_TILES_ARGS" ]; then
    CMD_ARGS="$CMD_ARGS --fetch-tiles $FETCH_TILES_ARGS"
fi

# Build CUDA_VISIBLE_DEVICES env var from GPU_ARG
CUDA_ENV_FLAG=""
if [ -n "$GPU_ARG" ]; then
    # Convert GPU_ARG (e.g. "0,2" or "1" or "all") to CUDA_VISIBLE_DEVICES
    if [ "$GPU_ARG" = "all" ]; then
        CUDA_VISIBLE="0,1,2,3"
    else
        CUDA_VISIBLE="$GPU_ARG"
    fi
    CUDA_ENV_FLAG="-e CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE}"
fi

docker run --rm --init $GPU_FLAG $CUDA_ENV_FLAG \
    --user 1000:1000 \
    -v "${INPUT_DIR}:/data/input${INPUT_RO}" \
    -v "${OUTPUT_DIR}:/data/output" \
    "$IMAGE_NAME" \
    python3 -m lidar_pipeline /data/input \
    $CMD_ARGS