#!/bin/bash

# Define a timeout function that tries multiple approaches
multiplatform_timeout() {
    # Check if 'timeout' command is available (standard on Linux)
    if command -v timeout >/dev/null 2>&1; then
        command timeout -k 60 "$@"
    # Check if 'gtimeout' is available (macOS with coreutils)
    elif command -v gtimeout >/dev/null 2>&1; then
        gtimeout -k 60 "$@"
    else
        echo "WARNING: Neither 'timeout' nor 'gtimeout' command found. Falling back to Perl implementation." >&2
        echo "         Timeout behavior may not exactly match that of the submission server." >&2
        
        # Use perl with proper signal handling
        perl -e 'alarm shift; exec @ARGV' "$@"
    fi
}

# Initialize default values
SUBMISSION_FOLDER_PATH="submission"
RUN_ID=""
NO_TIMESTAMP=false
REMOVE_IMAGES=false

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -f|--folder)
            SUBMISSION_FOLDER_PATH="$2"
            shift 2
            ;;
        -i|--id)
            RUN_ID="$2"
            shift 2
            ;;
        -n|--no-timestamp)
            NO_TIMESTAMP=true
            shift
            ;;
        -r|--remove-images)
            REMOVE_IMAGES=true
            shift
            ;;
        -h|--help)
            echo "Usage: $0 [OPTIONS]"
            echo ""
            echo "Options:"
            echo "  -f, --folder PATH       Path to submission folder (default: submission)"
            echo "  -i, --id NAME           Add run id to output folder name"
            echo "  -n, --no-timestamp      Do not add timestamp to output folder name"
            echo "  -r, --remove-images     Remove Docker images after execution (default: false)"
            echo "  -h, --help          Show this help message"
            echo ""
            echo "Output folders are created as: outputs/<submission_folder_name>_<timestamp>_<run_id>"
            echo ""
            echo "Examples:"
            echo "  $0                                          # Output to outputs/submission_YYYYMMDD_HHMMSS"
            echo "  $0 --folder my_submission                   # Output to outputs/my_submission_YYYYMMDD_HHMMSS"
            echo "  $0 --id run1                                # Output to outputs/submission_YYYYMMDD_HHMMSS_run1"
            echo "  $0 --folder sub2 --id run2                  # Output to outputs/sub2_YYYYMMDD_HHMMSS_run2"
            echo "  $0 --folder sub2 --id run2 --no-timestamp   # Output to outputs/sub2_run2"
            exit 0
            ;;
        *)
            echo "Error: Unknown option $1"
            echo "Use $0 --help for usage information"
            exit 1
            ;;
    esac
done

# Check if submission folder exists
if [ ! -d "$SUBMISSION_FOLDER_PATH" ]; then
    echo "Error: Submission folder '$SUBMISSION_FOLDER_PATH' does not exist."
    exit 1
fi

# Extract folder name from path (basename)
SUBMISSION_FOLDER_NAME=$(basename "$SUBMISSION_FOLDER_PATH")

# # Determine output directory name
# if [ -z "$RUN_ID" ] && [ "$NO_TIMESTAMP" = true ]; then
#     echo "Cannot use --no-timestamp without --id. Exiting."
#     exit 1
# fi

# Create outputs directory and subdirectory for this submission
# Prefix with submission folder name
OUTPUT_DIR="outputs/${SUBMISSION_FOLDER_NAME}"

if [ "$NO_TIMESTAMP" != "true" ]; then
    OUTPUT_DIR="${OUTPUT_DIR}_$(date +"%Y%m%d_%H%M%S")"
fi

if [ -n "$RUN_ID" ]; then
    OUTPUT_DIR="${OUTPUT_DIR}_${RUN_ID}"
fi

if [ -d "$OUTPUT_DIR" ]; then
    rm -rf "$OUTPUT_DIR"
fi
mkdir -p "$OUTPUT_DIR"
LOG_FILE="$OUTPUT_DIR/run.log"

# Start logging - capture ALL output from this point forward
{
    echo "Running submission from: $SUBMISSION_FOLDER_PATH"
    echo "Submission folder name: $SUBMISSION_FOLDER_NAME"
    echo "Output directory: $OUTPUT_DIR"
    echo "Log file: $LOG_FILE"
    echo "Remove images after execution: $REMOVE_IMAGES"
    echo "=================================="

    # Export the submission folder path for docker-compose
    # Note: We need to handle both relative and absolute paths
    # Convert to absolute path for docker-compose
    export SUBMISSION_FOLDER=$(realpath "$SUBMISSION_FOLDER_PATH")

    # Run docker-compose with timeout (20 minutes = 1200 seconds)
    echo "Starting docker-compose with 20 minute timeout..."
    multiplatform_timeout 1200 docker compose up --build --abort-on-container-exit

    DOCKER_EXIT_CODE=$?
    
    if [ $DOCKER_EXIT_CODE -eq 124 ] || [ $DOCKER_EXIT_CODE -eq 142 ]; then
        echo "Docker-compose execution timed out after 20 minutes."
    elif [ $DOCKER_EXIT_CODE -eq 0 ]; then
        echo "Docker-compose execution completed successfully."
    else
        echo "Docker-compose execution failed with exit code: $DOCKER_EXIT_CODE"
    fi

    # Try to copy result.json from the agent container
    echo "Attempting to copy result.json from agent container..."
    if docker cp agent_controller:/app/result.json "$OUTPUT_DIR/result.json"; then
        echo "Successfully copied result.json from agent container"
    else
        echo "Failed to copy result.json from agent container, creating fallback result"
        
        if [ $DOCKER_EXIT_CODE -eq 124 ] || [ $DOCKER_EXIT_CODE -eq 142 ]; then
            # Timeout occurred
            cat > "$OUTPUT_DIR/result.json" << EOF
{
  "status": "TIMEOUT",
  "evaluation": 0,
  "reasoning": "Execution timed out after 20 minutes. Check the run.log file for details.",
  "error_type": "execution_timeout"
}
EOF
        else
            # Other failure
            cat > "$OUTPUT_DIR/result.json" << EOF
{
  "status": "BUILD_OR_RUN_FAILURE",
  "evaluation": 0,
  "reasoning": "Validation failed due to build or run failure. Check the run.log file in this directory for detailed error information.",
  "error_type": "container_failure"
}
EOF
        fi
    fi

    # Clean up - conditionally remove images based on flag
    if [ "$REMOVE_IMAGES" = true ]; then
        echo "Cleaning up containers, volumes, and images..."
        docker compose down -v --rmi all
    else
        echo "Cleaning up containers and volumes..."
        docker compose down -v
    fi

    echo "=================================="
    echo "Submission execution completed."
    echo "Results saved to: $OUTPUT_DIR"
} 2>&1 | tee "$LOG_FILE"
