#!/bin/bash

# Generic tool testing script for OpenRouter
# Usage: ./test_generic_model.sh YOUR_API_KEY [quick|full]
#
# CONFIGURATION: Edit the variables below to customize your test
# ================================================================

# MODEL CONFIGURATION
# Examples: "openai/gpt-4o", "anthropic/claude-3.5-sonnet", "qwen/qwen3-coder", "deepseek/deepseek-chat-v3.1"
MODEL="qwen/qwen3-coder"

# TEMPERATURE CONFIGURATION (0.0 to 2.0)
# Lower = more deterministic, Higher = more creative
# Recommended: 0.0-0.3 for tool calling tests
TEMPERATURE=0.6

# PROVIDERS TO TEST
# Leave empty to test without specific provider routing
# For a single provider test, use: providers=("openai")
# For multiple providers, list them all
# Common providers: openai, anthropic, google, deepseek, fireworks, together, deepinfra, etc.
# Some providers support precision specifiers: "fireworks/fp8", "deepinfra/fp4", etc.
providers=(
    # Uncomment and modify the providers you want to test:
    # "openai"
    # "anthropic"
    # "google"
    # "deepseek"
    # "fireworks"
    # "fireworks/fp8"
    # "together"
    # "together/fp8"
    # "deepinfra/fp4"
    # "baseten/fp8"
    # "chutes/fp8"
    # "parasail/fp8"
    # "novita/fp8"
    # "atlas-cloud/fp8"
    # "phala"
    # "gmicloud/fp8"
    # "targon/fp8"
    "alibaba/opensource"
    # "hyperbolic/fp8"
    # "cerebras/fp8"
)

# If providers array is empty, run a single test without provider routing
if [ ${#providers[@]} -eq 0 ]; then
    providers=("none")
fi

# DELAY BETWEEN TESTS (in seconds)
# Increase if you encounter rate limiting
DELAY_SECONDS=3

# REASONING EFFORT (optional, for models that support it)
# Options: low, medium, high, or leave empty
REASONING_EFFORT=""

# ================================================================
# END OF CONFIGURATION - DO NOT EDIT BELOW THIS LINE
# ================================================================

if [ $# -lt 1 ]; then
    echo "Usage: $0 YOUR_OPENROUTER_API_KEY [quick|full]"
    echo ""
    echo "Current Configuration:"
    echo "  Model: $MODEL"
    echo "  Temperature: $TEMPERATURE"
    echo "  Providers: ${providers[*]}"
    echo ""
    echo "Examples:"
    echo "  $0 sk-or-xxx quick    # Run quick tests"
    echo "  $0 sk-or-xxx full     # Run full test suite"
    echo ""
    echo "Edit this script to change the model, temperature, and providers to test."
    exit 1
fi

API_KEY=$1
TEST_MODE=${2:-quick}  # Default to quick mode if not specified

# Create results directory with timestamp
TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
SAFE_MODEL_NAME=$(echo "$MODEL" | sed 's/[\/:]/_/g')
RESULTS_DIR="test_results_${SAFE_MODEL_NAME}_${TIMESTAMP}"
mkdir -p "$RESULTS_DIR"

echo "================================================================"
echo "TOOL TESTING CONFIGURATION"
echo "================================================================"
echo "  Model: $MODEL"
echo "  Temperature: $TEMPERATURE"
echo "  Test Mode: $TEST_MODE"
echo "  Providers to test: ${#providers[@]}"
echo "  Results Directory: $RESULTS_DIR"
if [ -n "$REASONING_EFFORT" ]; then
    echo "  Reasoning Effort: $REASONING_EFFORT"
fi
echo ""

# Function to run test and capture results
run_test() {
    local provider=$1
    local test_name=""
    local provider_param=""
    
    if [ "$provider" == "none" ]; then
        test_name="no_provider"
        echo "================================================"
        echo "Testing model: $MODEL (no provider routing)"
    else
        test_name="${provider//\//_}"
        provider_param="--provider \"$provider\""
        echo "================================================"
        echo "Testing model: $MODEL"
        echo "Provider: $provider"
    fi
    
    local output_file="$RESULTS_DIR/test_${test_name}.txt"
    echo "Output file: $output_file"
    echo "================================================"
    
    # Build the command
    cmd="python3 tool_tester_v2.py \
        --api-base https://openrouter.ai/api/v1 \
        --api-key $API_KEY \
        --model \"$MODEL\" \
        --temperature $TEMPERATURE"
    
    # Add provider if specified
    if [ -n "$provider_param" ]; then
        cmd="$cmd $provider_param"
    fi
    
    # Add reasoning effort if specified
    if [ -n "$REASONING_EFFORT" ]; then
        cmd="$cmd --reasoning-effort $REASONING_EFFORT"
    fi
    
    # Add test mode flag
    if [ "$TEST_MODE" == "quick" ]; then
        cmd="$cmd --quick"
    else
        cmd="$cmd --max-tools 40"
    fi
    
    # Add output file
    cmd="$cmd --output \"$output_file\""
    
    # Run the test and capture both stdout and the result
    echo "Running: $cmd"
    eval $cmd 2>&1 | tee "${output_file}.log"
    
    # Check if the test completed successfully
    if [ ${PIPESTATUS[0]} -eq 0 ]; then
        echo "PASS: Test completed successfully"
        
        # Extract the overall score from the output file if it exists
        if [ -f "$output_file" ]; then
            score=$(grep "OVERALL SCORE:" "$output_file" | tail -1)
            echo "  $score"
        fi
        return 0
    else
        echo "FAIL: Test failed"
        if [ "$provider" == "none" ]; then
            echo "FAILED: No provider routing" >> "$RESULTS_DIR/failed_tests.txt"
        else
            echo "FAILED: $provider" >> "$RESULTS_DIR/failed_tests.txt"
        fi
        return 1
    fi
    
    echo ""
}

# Test each provider/configuration
successful_tests=0
failed_tests=0
total_tests=${#providers[@]}
current_test=0

for provider in "${providers[@]}"; do
    ((current_test++))
    echo ""
    echo "Test $current_test of $total_tests"
    
    run_test "$provider"
    
    # Check if test completed (not if it got a perfect score)
    if [ $? -eq 0 ]; then
        ((successful_tests++))
    else
        ((failed_tests++))
    fi
    
    # Add delay between tests (except for the last one)
    if [ $current_test -lt $total_tests ]; then
        echo "Waiting ${DELAY_SECONDS}s before next test..."
        sleep $DELAY_SECONDS
    fi
done

# Generate summary report
summary_file="$RESULTS_DIR/summary.txt"
echo "" | tee "$summary_file"
echo "================================================================" | tee -a "$summary_file"
echo "TOOL TESTING SUMMARY REPORT" | tee -a "$summary_file"
echo "================================================================" | tee -a "$summary_file"
echo "Model: $MODEL" | tee -a "$summary_file"
echo "Temperature: $TEMPERATURE" | tee -a "$summary_file"
echo "Test Mode: $TEST_MODE" | tee -a "$summary_file"
echo "Timestamp: $TIMESTAMP" | tee -a "$summary_file"
echo "Total Tests Run: $total_tests" | tee -a "$summary_file"
echo "Tests Completed: $successful_tests" | tee -a "$summary_file"
echo "Tests Failed to Run: $failed_tests" | tee -a "$summary_file"
echo "" | tee -a "$summary_file"

# Extract scores from all successful tests
echo "Test Results:" | tee -a "$summary_file"
echo "-----------------------------------------" | tee -a "$summary_file"

for file in "$RESULTS_DIR"/test_*.txt; do
    if [ -f "$file" ]; then
        test_name=$(basename "$file" .txt | sed 's/test_//')
        
        # Convert filename back to provider name
        if [ "$test_name" == "no_provider" ]; then
            display_name="No provider routing"
        else
            display_name=$(echo "$test_name" | sed 's/_/\//g')
        fi
        
        score=$(grep "OVERALL SCORE:" "$file" | tail -1 | awk '{print $3, $4}')
        if [ -n "$score" ]; then
            printf "%-30s %s\n" "$display_name:" "$score" | tee -a "$summary_file"
        else
            printf "%-30s %s\n" "$display_name:" "Test failed" | tee -a "$summary_file"
        fi
    fi
done

echo "" | tee -a "$summary_file"

# Create a CSV summary for easy analysis
csv_file="$RESULTS_DIR/results.csv"
echo "Provider,Model,Temperature,OverallScore,Grade,ScenarioSuccessRate,ToolPrecision,ToolRecall,ToolF1,ParamAccuracy,ExecutionSuccess" > "$csv_file"

for file in "$RESULTS_DIR"/test_*.txt; do
    if [ -f "$file" ]; then
        test_name=$(basename "$file" .txt | sed 's/test_//')
        
        # Convert filename back to provider name
        if [ "$test_name" == "no_provider" ]; then
            provider_name="none"
        else
            provider_name=$(echo "$test_name" | sed 's/_/\//g')
        fi
        
        # Extract metrics using grep and awk
        overall_score=$(grep "OVERALL SCORE:" "$file" | tail -1 | awk '{print $3}' | sed 's/%//')
        grade=$(grep "OVERALL SCORE:" "$file" | tail -1 | awk '{print $4}' | sed 's/[()]//g')
        scenario_rate=$(grep "Scenario Success Rate:" "$file" | tail -1 | awk '{print $4}' | sed 's/%//')
        tool_precision=$(grep "Tool Precision (LLM only):" "$file" | tail -1 | awk '{print $5}' | sed 's/%//')
        tool_recall=$(grep "Tool Recall (LLM only):" "$file" | tail -1 | awk '{print $5}' | sed 's/%//')
        tool_f1=$(grep "Tool F1 (LLM only):" "$file" | tail -1 | awk '{print $5}' | sed 's/%//')
        param_accuracy=$(grep "Parameter Accuracy (structural):" "$file" | tail -1 | awk '{print $4}' | sed 's/%//')
        exec_success=$(grep "Execution Success Rate:" "$file" | tail -1 | awk '{print $4}' | sed 's/%//')
        
        if [ -n "$overall_score" ]; then
            echo "$provider_name,$MODEL,$TEMPERATURE,$overall_score,$grade,$scenario_rate,$tool_precision,$tool_recall,$tool_f1,$param_accuracy,$exec_success" >> "$csv_file"
        fi
    fi
done

echo "Full results saved in: $RESULTS_DIR" | tee -a "$summary_file"
echo "CSV results saved in: $csv_file" | tee -a "$summary_file"
echo "" | tee -a "$summary_file"

# Sort results by score if there are multiple tests
if [ $total_tests -gt 1 ]; then
    echo "Top Performers (sorted by score):" | tee -a "$summary_file"
    echo "-----------------------------------------" | tee -a "$summary_file"
    
    # Create temp file for sorting
    temp_scores="/tmp/scores_$$.txt"
    
    for file in "$RESULTS_DIR"/test_*.txt; do
        if [ -f "$file" ]; then
            test_name=$(basename "$file" .txt | sed 's/test_//')
            if [ "$test_name" == "no_provider" ]; then
                display_name="No provider routing"
            else
                display_name=$(echo "$test_name" | sed 's/_/\//g')
            fi
            
            score=$(grep "OVERALL SCORE:" "$file" | tail -1 | awk '{print $3}' | sed 's/%//')
            if [ -n "$score" ]; then
                echo "$score|$display_name" >> "$temp_scores"
            fi
        fi
    done
    
    if [ -f "$temp_scores" ]; then
        sort -t'|' -k1 -rn "$temp_scores" | while IFS='|' read -r score provider; do
            printf "%-30s %s%%\n" "$provider:" "$score" | tee -a "$summary_file"
        done
        rm -f "$temp_scores"
    fi
fi

echo ""
echo "================================================================"
echo "Testing Complete!"
echo "================================================================"
echo ""
echo "View detailed results:"
echo "  Summary: cat $RESULTS_DIR/summary.txt"
echo "  CSV data: cat $RESULTS_DIR/results.csv"
echo ""

# Keep terminal open if running in Windows Git Bash
if [[ "$OSTYPE" == "msys" ]] || [[ "$OSTYPE" == "cygwin" ]]; then
    echo ""
    read -p "Press Enter to exit..."
fi

exit 0