#!/bin/bash
# Validate CNN v2: Export checkpoint → Build → Test → Display results

set -e

# Default paths
BUILD_DIR="build"
WORKSPACE="main"
TEST_IMAGES_DIR="training/validation"
OUTPUT_DIR="validation_results"
PYTHON="python3"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

print_usage() {
    cat << EOF
Usage: $0 CHECKPOINT [OPTIONS]

End-to-end CNN v2 validation: export shaders, rebuild, test images, show results.

Arguments:
  CHECKPOINT              Path to .pth checkpoint file (required)

Options:
  -b, --build-dir DIR     Build directory (default: build)
  -w, --workspace NAME    Workspace name (default: main)
  -i, --images DIR        Test images directory (default: training/validation)
  -o, --output DIR        Output directory (default: validation_results)
  --python CMD            Python command (default: python3)
  --skip-build            Skip cnn_test rebuild
  --skip-export           Skip shader export (use existing .wgsl)
  -h, --help              Show this help

Example:
  $0 checkpoints/checkpoint_epoch_5000.pth
  $0 checkpoint.pth -i my_test_images/ -o results/
  $0 checkpoint.pth --skip-build  # Use existing cnn_test binary

EOF
}

log() { echo -e "${GREEN}[validate]${NC} $*"; }
warn() { echo -e "${YELLOW}[validate]${NC} $*"; }
error() { echo -e "${RED}[validate]${NC} $*" >&2; exit 1; }

# Parse arguments
CHECKPOINT=""
SKIP_BUILD=false
SKIP_EXPORT=false

while [[ $# -gt 0 ]]; do
    case $1 in
        -h|--help)
            print_usage
            exit 0
            ;;
        -b|--build-dir)
            BUILD_DIR="$2"
            shift 2
            ;;
        -w|--workspace)
            WORKSPACE="$2"
            shift 2
            ;;
        -i|--images)
            TEST_IMAGES_DIR="$2"
            shift 2
            ;;
        -o|--output)
            OUTPUT_DIR="$2"
            shift 2
            ;;
        --python)
            PYTHON="$2"
            shift 2
            ;;
        --skip-build)
            SKIP_BUILD=true
            shift
            ;;
        --skip-export)
            SKIP_EXPORT=true
            shift
            ;;
        -*)
            error "Unknown option: $1"
            ;;
        *)
            if [[ -z "$CHECKPOINT" ]]; then
                CHECKPOINT="$1"
            else
                error "Unexpected argument: $1"
            fi
            shift
            ;;
    esac
done

# Validate inputs
[[ -z "$CHECKPOINT" ]] && error "Checkpoint file required. Use -h for help."
[[ ! -f "$CHECKPOINT" ]] && error "Checkpoint not found: $CHECKPOINT"
[[ ! -d "$TEST_IMAGES_DIR" ]] && error "Test images directory not found: $TEST_IMAGES_DIR"

SHADER_DIR="workspaces/$WORKSPACE/shaders"
CNN_TEST="$BUILD_DIR/cnn_test"

log "Configuration:"
log "  Checkpoint:  $CHECKPOINT"
log "  Build dir:   $BUILD_DIR"
log "  Workspace:   $WORKSPACE"
log "  Shader dir:  $SHADER_DIR"
log "  Test images: $TEST_IMAGES_DIR"
log "  Output dir:  $OUTPUT_DIR"
echo

# Step 1: Export shaders
if [[ "$SKIP_EXPORT" = false ]]; then
    log "Step 1/4: Exporting shaders from checkpoint..."
    [[ ! -d "$SHADER_DIR" ]] && error "Shader directory not found: $SHADER_DIR"

    if [[ ! -f "training/export_cnn_v2_shader.py" ]]; then
        error "Export script not found: training/export_cnn_v2_shader.py"
    fi

    $PYTHON training/export_cnn_v2_shader.py "$CHECKPOINT" --output-dir "$SHADER_DIR" \
        || error "Shader export failed"

    log "✓ Shaders exported to $SHADER_DIR"
else
    warn "Skipping shader export (using existing .wgsl files)"
fi

# Step 2: Rebuild cnn_test
if [[ "$SKIP_BUILD" = false ]]; then
    log "Step 2/4: Rebuilding cnn_test..."

    cmake --build "$BUILD_DIR" -j4 --target cnn_test \
        || error "Build failed"

    log "✓ Built $CNN_TEST"
else
    warn "Skipping build (using existing binary)"
fi

[[ ! -x "$CNN_TEST" ]] && error "cnn_test not found or not executable: $CNN_TEST"

# Step 3: Process test images
log "Step 3/4: Processing test images..."
mkdir -p "$OUTPUT_DIR"

# Find PNG images
mapfile -t IMAGES < <(find "$TEST_IMAGES_DIR" -maxdepth 1 -name "*.png" | sort)
[[ ${#IMAGES[@]} -eq 0 ]] && error "No PNG images found in $TEST_IMAGES_DIR"

log "Found ${#IMAGES[@]} test image(s)"

for img in "${IMAGES[@]}"; do
    basename=$(basename "$img" .png)
    output="$OUTPUT_DIR/${basename}_output.png"

    log "  Processing $basename.png..."
    "$CNN_TEST" "$img" "$output" --cnn-version v2 \
        || warn "  Failed: $basename.png"
done

log "✓ Processed ${#IMAGES[@]} image(s)"

# Step 4: Display results
log "Step 4/4: Opening results..."

case "$(uname -s)" in
    Darwin*)
        open "$OUTPUT_DIR"
        ;;
    Linux*)
        if command -v xdg-open &> /dev/null; then
            xdg-open "$OUTPUT_DIR"
        else
            log "Results saved to: $OUTPUT_DIR"
        fi
        ;;
    MINGW*|MSYS*|CYGWIN*)
        explorer "$OUTPUT_DIR"
        ;;
    *)
        log "Results saved to: $OUTPUT_DIR"
        ;;
esac

log "✓ Validation complete!"
log ""
log "Results:"
log "  Input:  $TEST_IMAGES_DIR/*.png"
log "  Output: $OUTPUT_DIR/*_output.png"