#!/bin/bash

# Source project-specific API keys
source "$(dirname "$0")/source_keys.sh"

# Help text function
show_help() {
    echo "Usage: $0 [OPTIONS]"
    echo "Run ETR problems evaluation with OpenRouter using specified parameters."
    echo
    echo "Options:"
    echo "  -m, --model MODEL      OpenRouter model to use (e.g., google/gemini-2-5-pro)"
    echo "  -p, --path PATH        Path to lm-evaluation-harness directory"
    echo "  -i, --include PATH     Path to include for task definitions"
    echo "  -d, --dataset PATH     Path to dataset JSONL file to evaluate"
    echo "  -v, --verbosity LEVEL  Verbosity level (default: WARNING)"
    echo "  -o, --output NAME      Output directory name (default: results)"
    echo "  -h, --help            Show this help message"
    echo
    echo "Example:"
    echo "  $0 -m google/gemini-2-5-pro -d /path/to/dataset.jsonl -o custom_results"
}

# Default values
MODEL="openai/chatgpt-4o-latest"
EVAL_PATH="${EVAL_PATH:-/home/keenan/Dev/lm-evaluation-harness/}"
INCLUDE_PATH="${INCLUDE_PATH:-/home/keenan/Dev/etr_case_generator/}"
DATASET=""
TASK="etr_problems"
VERBOSITY="WARNING"
OUTPUT_DIR="results"

# Check for required OpenRouter API key
check_api_key() {
    if [ -z "$OPENROUTER_API_KEY" ]; then
        echo "Error: OPENROUTER_API_KEY is not set"
        echo "Please ensure it is correctly set in keys.env"
        exit 1
    fi
}

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -m|--model)
            MODEL="$2"
            shift 2
            ;;
        -p|--path)
            EVAL_PATH="$2"
            shift 2
            ;;
        -i|--include)
            INCLUDE_PATH="$2"
            shift 2
            ;;
        -d|--dataset)
            DATASET="$2"
            shift 2
            ;;
        -v|--verbosity)
            VERBOSITY="$2"
            shift 2
            ;;
        -o|--output)
            OUTPUT_DIR="$2"
            shift 2
            ;;
        -h|--help)
            show_help
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            show_help
            exit 1
            ;;
    esac
done

# Handle dataset routing if specified
if [ -n "$DATASET" ]; then
    if [[ $DATASET == *"yes_no"* ]]; then
        TASK="etr_problems"
    elif [[ $DATASET == *"open_ended"* ]]; then
        TASK="etr_problems_open_ended"
    elif [[ $DATASET == *"multiple_choice"* ]]; then
        echo "Error: multiple-choice questions are not yet implemented"
        exit 1
    else
        echo "Warning: Dataset filename does not contain 'yes_no', 'open_ended', or 'multiple_choice'"
        echo "Defaulting to yes/no questions task configuration"
        TASK="etr_problems"
    fi

    # Copy dataset to the root datasets directory
    mkdir -p "datasets"
    cp "$DATASET" "datasets/etr_for_lm_eval.jsonl"
    echo "Copied $DATASET to datasets/etr_for_lm_eval.jsonl"
fi

# Assert that EVAL_PATH is a real directory
if [ ! -d "$EVAL_PATH" ]; then
    echo "Error: $EVAL_PATH is not a valid directory, please rerun with the '--path' option. Sorry for hardcoding my own."
    exit 1
fi

# Assert that INCLUDE_PATH is a real directory
if [ ! -d "$INCLUDE_PATH" ]; then
    echo "Error: $INCLUDE_PATH is not a valid directory, please rerun with the '--include' option. Sorry for hardcoding my own."
    exit 1
fi

# Export PYTHONPATH if needed
if [ -n "$EVAL_PATH" ]; then
    export PYTHONPATH="${EVAL_PATH}:${PYTHONPATH:+:$PYTHONPATH}"
fi

echo "Configuration:"
echo "  OpenRouter Model: $MODEL"
echo "  Evaluation harness path: $EVAL_PATH"
echo "  Include path: $INCLUDE_PATH"
echo "  Task: $TASK"
echo "  Output directory: lm_eval/tasks/etr_problems/$OUTPUT_DIR"
echo ""

# Check API key requirements
check_api_key

# Store original OPENAI_API_KEY if it exists and set up restoration trap
if [ -n "$OPENAI_API_KEY" ]; then
    export ORIGINAL_OPENAI_KEY=$OPENAI_API_KEY
    # This will run on any exit (normal, interrupt, error)
    trap 'export OPENAI_API_KEY=$ORIGINAL_OPENAI_KEY; echo "Restored original OpenAI API key"' EXIT SIGINT SIGTERM
fi

# Set OPENAI_API_KEY to OPENROUTER_API_KEY
# We need to do this because the evaluation harness looks directly at OPENAI_API_KEY
export OPENAI_API_KEY=$OPENROUTER_API_KEY

# NOTE! You can use reasoning.exclued=true, like this:
#    --model_args "base_url=https://openrouter.ai/api/v1/chat/completions,model=${MODEL},max_tokens=3000,num_concurrent=1,reasoning.exclude=true" \

# Run evaluation with OpenRouter
lm_eval --model openai-chat-completions \
    --model_args "base_url=https://openrouter.ai/api/v1/chat/completions,model=${MODEL},max_tokens=3000,num_concurrent=5" \
    --include_path "$INCLUDE_PATH" \
    --tasks $TASK \
    --num_fewshot 0 \
    --batch_size 1 \
    --output_path "lm_eval/tasks/etr_problems/${OUTPUT_DIR}" \
    --apply_chat_template \
    --log_samples \
    --write_out \
    ${VERBOSITY:+--verbosity "$VERBOSITY"}
