#!/bin/bash
# Test script for the sim-search API using curl commands

# Configuration
API_URL="http://localhost:8000"
API_V1="${API_URL}/api/v1"
TOKEN=""
EMAIL="test@example.com"
PASSWORD="password123"
FULL_NAME="Test User"

# Colors for output
GREEN='\033[0;32m'
RED='\033[0;31m'
YELLOW='\033[0;33m'
NC='\033[0m' # No Color

# Function to print section header
print_header() {
    echo -e "\n${YELLOW}=== $1 ===${NC}"
}

# Function to print success message
print_success() {
    echo -e "${GREEN}✓ $1${NC}"
}

# Function to print error message
print_error() {
    echo -e "${RED}✗ $1${NC}"
}

# Function to check if the API is running
check_api() {
    print_header "Checking if API is running"
    
    response=$(curl -s -o /dev/null -w "%{http_code}" ${API_URL})
    
    if [ "$response" == "200" ]; then
        print_success "API is running"
    else
        print_error "API is not running. Please start the API server first."
        exit 1
    fi
}

# Function to register a user
register_user() {
    print_header "Registering a user"
    
    response=$(curl -s -X POST \
        -H "Content-Type: application/json" \
        -d "{\"email\":\"${EMAIL}\",\"password\":\"${PASSWORD}\",\"full_name\":\"${FULL_NAME}\",\"is_active\":true,\"is_superuser\":false}" \
        ${API_V1}/auth/register)
    
    if echo "$response" | grep -q "email"; then
        print_success "User registered successfully"
    else
        # If user already exists, that's fine
        if echo "$response" | grep -q "already exists"; then
            print_success "User already exists, continuing with login"
        else
            print_error "Failed to register user: $response"
        fi
    fi
}

# Function to get an authentication token
get_token() {
    print_header "Getting authentication token"
    
    response=$(curl -s -X POST \
        -H "Content-Type: application/x-www-form-urlencoded" \
        -d "username=${EMAIL}&password=${PASSWORD}" \
        ${API_V1}/auth/token)
    
    if echo "$response" | grep -q "access_token"; then
        TOKEN=$(echo "$response" | grep -o '"access_token":"[^"]*' | sed 's/"access_token":"//')
        print_success "Got authentication token"
    else
        print_error "Failed to get authentication token: $response"
        exit 1
    fi
}

# Function to process a query
process_query() {
    print_header "Processing a query"
    
    response=$(curl -s -X POST \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer ${TOKEN}" \
        -d "{\"query\":\"What are the environmental impacts of electric vehicles?\"}" \
        ${API_V1}/query/process)
    
    if echo "$response" | grep -q "structured_query"; then
        print_success "Query processed successfully"
    else
        print_error "Failed to process query: $response"
    fi
}

# Function to classify a query
classify_query() {
    print_header "Classifying a query"
    
    response=$(curl -s -X POST \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer ${TOKEN}" \
        -d "{\"query\":\"What are the environmental impacts of electric vehicles?\"}" \
        ${API_V1}/query/classify)
    
    if echo "$response" | grep -q "structured_query"; then
        print_success "Query classified successfully"
    else
        print_error "Failed to classify query: $response"
    fi
}

# Function to get available search engines
get_search_engines() {
    print_header "Getting available search engines"
    
    response=$(curl -s -X GET \
        -H "Authorization: Bearer ${TOKEN}" \
        ${API_V1}/search/engines)
    
    if echo "$response" | grep -q "\["; then
        print_success "Got search engines successfully"
    else
        print_error "Failed to get search engines: $response"
    fi
}

# Function to execute a search
execute_search() {
    print_header "Executing a search"
    
    response=$(curl -s -X POST \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer ${TOKEN}" \
        -d "{\"structured_query\":{\"original_query\":\"What are the environmental impacts of electric vehicles?\",\"enhanced_query\":\"What are the environmental impacts of electric vehicles?\",\"type\":\"factual\",\"domain\":\"environmental\"},\"search_engines\":[\"google\",\"arxiv\"],\"num_results\":5,\"timeout\":30}" \
        ${API_V1}/search/execute)
    
    if echo "$response" | grep -q "search_id"; then
        SEARCH_ID=$(echo "$response" | grep -o '"search_id":"[^"]*' | sed 's/"search_id":"//')
        print_success "Search executed successfully with ID: $SEARCH_ID"
    else
        print_error "Failed to execute search: $response"
    fi
}

# Function to get search history
get_search_history() {
    print_header "Getting search history"
    
    response=$(curl -s -X GET \
        -H "Authorization: Bearer ${TOKEN}" \
        ${API_V1}/search/history)
    
    if echo "$response" | grep -q "searches"; then
        print_success "Got search history successfully"
    else
        print_error "Failed to get search history: $response"
    fi
}

# Function to get search results
get_search_results() {
    print_header "Getting search results"
    
    if [ -z "$SEARCH_ID" ]; then
        print_error "No search ID available. Please execute a search first."
        return
    fi
    
    response=$(curl -s -X GET \
        -H "Authorization: Bearer ${TOKEN}" \
        ${API_V1}/search/${SEARCH_ID})
    
    if echo "$response" | grep -q "search_id"; then
        print_success "Got search results successfully"
    else
        print_error "Failed to get search results: $response"
    fi
}

# Function to generate a report
generate_report() {
    print_header "Generating a report"
    
    if [ -z "$SEARCH_ID" ]; then
        print_error "No search ID available. Please execute a search first."
        return
    fi
    
    response=$(curl -s -X POST \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer ${TOKEN}" \
        -d "{\"search_id\":\"${SEARCH_ID}\",\"query\":\"What are the environmental impacts of electric vehicles?\",\"detail_level\":\"standard\",\"query_type\":\"factual\",\"model\":\"llama-3.1-8b-instant\"}" \
        ${API_V1}/report/generate)
    
    if echo "$response" | grep -q "id"; then
        REPORT_ID=$(echo "$response" | grep -o '"id":"[^"]*' | sed 's/"id":"//')
        print_success "Report generated successfully with ID: $REPORT_ID"
    else
        print_error "Failed to generate report: $response"
    fi
}

# Function to get report progress
get_report_progress() {
    print_header "Getting report progress"
    
    if [ -z "$REPORT_ID" ]; then
        print_error "No report ID available. Please generate a report first."
        return
    fi
    
    response=$(curl -s -X GET \
        -H "Authorization: Bearer ${TOKEN}" \
        ${API_V1}/report/${REPORT_ID}/progress)
    
    if echo "$response" | grep -q "progress"; then
        print_success "Got report progress successfully"
    else
        print_error "Failed to get report progress: $response"
    fi
}

# Function to get report list
get_report_list() {
    print_header "Getting report list"
    
    response=$(curl -s -X GET \
        -H "Authorization: Bearer ${TOKEN}" \
        ${API_V1}/report/list)
    
    if echo "$response" | grep -q "reports"; then
        print_success "Got report list successfully"
    else
        print_error "Failed to get report list: $response"
    fi
}

# Function to get a specific report
get_report() {
    print_header "Getting a specific report"
    
    if [ -z "$REPORT_ID" ]; then
        print_error "No report ID available. Please generate a report first."
        return
    fi
    
    response=$(curl -s -X GET \
        -H "Authorization: Bearer ${TOKEN}" \
        ${API_V1}/report/${REPORT_ID})
    
    if echo "$response" | grep -q "id"; then
        print_success "Got report successfully"
    else
        print_error "Failed to get report: $response"
    fi
}

# Function to download a report
download_report() {
    print_header "Downloading a report"
    
    if [ -z "$REPORT_ID" ]; then
        print_error "No report ID available. Please generate a report first."
        return
    fi
    
    response=$(curl -s -X GET \
        -H "Authorization: Bearer ${TOKEN}" \
        ${API_V1}/report/${REPORT_ID}/download?format=markdown)
    
    if [ -n "$response" ]; then
        print_success "Downloaded report successfully"
    else
        print_error "Failed to download report"
    fi
}

# Function to delete a report
delete_report() {
    print_header "Deleting a report"
    
    if [ -z "$REPORT_ID" ]; then
        print_error "No report ID available. Please generate a report first."
        return
    fi
    
    response=$(curl -s -o /dev/null -w "%{http_code}" -X DELETE \
        -H "Authorization: Bearer ${TOKEN}" \
        ${API_V1}/report/${REPORT_ID})
    
    if [ "$response" == "204" ]; then
        print_success "Report deleted successfully"
    else
        print_error "Failed to delete report: $response"
    fi
}

# Function to delete a search
delete_search() {
    print_header "Deleting a search"
    
    if [ -z "$SEARCH_ID" ]; then
        print_error "No search ID available. Please execute a search first."
        return
    fi
    
    response=$(curl -s -o /dev/null -w "%{http_code}" -X DELETE \
        -H "Authorization: Bearer ${TOKEN}" \
        ${API_V1}/search/${SEARCH_ID})
    
    if [ "$response" == "204" ]; then
        print_success "Search deleted successfully"
    else
        print_error "Failed to delete search: $response"
    fi
}

# Main function
main() {
    echo "Starting API tests..."
    
    # Check if the API is running
    check_api
    
    # Register a user
    register_user
    
    # Get an authentication token
    get_token
    
    # Process a query
    process_query
    
    # Classify a query
    classify_query
    
    # Get available search engines
    get_search_engines
    
    # Execute a search
    execute_search
    
    # Get search history
    get_search_history
    
    # Get search results
    get_search_results
    
    # Generate a report
    generate_report
    
    # Get report progress
    get_report_progress
    
    # Get report list
    get_report_list
    
    # Get a specific report
    get_report
    
    # Download a report
    download_report
    
    # Delete a report
    delete_report
    
    # Delete a search
    delete_search
    
    echo -e "\n${GREEN}All tests completed!${NC}"
}

# Run the main function
main