#!/bin/bash

# Test script for ZIP batch upload functionality
# This script exercises all batch endpoints with a running deployment
# use: task compose:up (or compose:refresh if you have made service changes)
# connect to the db like this
# jdbc:postgresql://localhost:5432/query_orchestration
#
# then to look at the assets
# List all buckets:
#  aws --endpoint-url=http://localhost:4566 s3 ls
#
#  List contents of a specific bucket:
#  aws --endpoint-url=http://localhost:4566 s3 ls s3://documentin/
#
#  List recursively to see all files in subdirectories:
#  aws --endpoint-url=http://localhost:4566 s3 ls s3://documentin/ --recursive
#
# REST API operations tested:
# - GET /metrics                                          - Check service health
# - POST /client                                          - Create new test client
# - POST /client/{CLIENT_ID}/document/batch               - Upload ZIP file with multiple PDFs
# - GET /client/{CLIENT_ID}/document/batch/{BATCH_ID}     - Get batch status and details
# - GET /client/{CLIENT_ID}/document/batch                - List all batches for client with pagination
# - POST /client/{CLIENT_ID}/document/batch (invalid)     - Test error handling with non-ZIP file
# - DELETE /client/{CLIENT_ID}/document/batch/{BATCH_ID}  - Cancel/delete batch processing

set -e  # Exit on any error

# Configuration
BASE_URL="http://localhost:8080"
CLIENT_ID="test_client_$(date +%s)"
ZIP_FILE="test_batch.zip"
TEST_PDF_COUNT=3

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

# Logging functions
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Check if service is running
check_service() {
    log_info "Checking if queryAPI service is running..."
    if ! curl -s "$BASE_URL/metrics" > /dev/null 2>&1; then
        log_error "queryAPI service is not running at $BASE_URL"
        log_info "Please start the stack with: task compose:up"
        exit 1
    fi
    log_info "Service is running"
}

# Create test ZIP file with PDFs
create_test_zip() {
    log_info "Creating test ZIP file with $TEST_PDF_COUNT PDFs..."
    
    # Remove existing file if it exists
    rm -f "$ZIP_FILE"
    
    # Create temporary directory for PDF files
    temp_dir=$(mktemp -d)
    
    # Create test PDF files
    for i in $(seq 1 $TEST_PDF_COUNT); do
        pdf_file="$temp_dir/test${i}.pdf"
        cat > "$pdf_file" << EOF
%PDF-1.4
%Test content for test${i}.pdf
1 0 obj
<<
/Type /Catalog
/Pages 2 0 R
>>
endobj
2 0 obj
<<
/Type /Pages
/Kids [3 0 R]
/Count 1
>>
endobj
3 0 obj
<<
/Type /Page
/Parent 2 0 R
/MediaBox [0 0 612 792]
>>
endobj
xref
0 4
0000000000 65535 f 
0000000010 00000 n 
0000000053 00000 n 
0000000125 00000 n 
trailer
<<
/Size 4
/Root 1 0 R
>>
startxref
284
%%EOF
EOF
    done
    
    # Create ZIP file using absolute path
    current_dir=$(pwd)
    cd "$temp_dir"
    zip -q "$current_dir/$ZIP_FILE" *.pdf
    cd "$current_dir"
    
    # Cleanup temp directory
    rm -rf "$temp_dir"
    
    file_size=$(stat -f%z "$ZIP_FILE" 2>/dev/null || stat -c%s "$ZIP_FILE" 2>/dev/null || echo "unknown")
    log_info "Created $ZIP_FILE with $TEST_PDF_COUNT PDFs ($file_size bytes)"
}

# Create client (idempotent)
create_client() {
    log_info "Creating/checking client: $CLIENT_ID"
    
    # Try to create client, ignore if already exists
    response=$(curl -s -w "\n%{http_code}" -X POST "$BASE_URL/client" \
        -H "Content-Type: application/json" \
        -d "{\"id\":\"$CLIENT_ID\",\"name\":\"Test Client $CLIENT_ID\"}" 2>/dev/null || echo "000")
    
    http_code=$(echo "$response" | tail -n1)
    
    if [ "$http_code" = "201" ] || [ "$http_code" = "200" ]; then
        log_info "Client created/exists: $CLIENT_ID"
    elif [ "$http_code" = "409" ] || [ "$http_code" = "400" ]; then
        log_info "Client already exists or duplicate name: $CLIENT_ID"
    else
        log_warn "Unexpected response when creating client (HTTP $http_code), continuing anyway..."
    fi
}

# Upload ZIP batch
upload_batch() {
    log_info "Uploading ZIP batch..."
    
    response=$(curl -s -w "\n%{http_code}" -X POST "$BASE_URL/client/$CLIENT_ID/document/batch" \
        -F "archive=@$ZIP_FILE" 2>/dev/null)
    
    http_code=$(echo "$response" | tail -n1)
    response_body=$(echo "$response" | head -n -1)
    
    if [ "$http_code" = "202" ]; then
        BATCH_ID=$(echo "$response_body" | grep -o '"batch_id":"[^"]*"' | cut -d'"' -f4)
        STATUS_URL=$(echo "$response_body" | grep -o '"status_url":"[^"]*"' | cut -d'"' -f4)
        log_info "Batch uploaded successfully"
        log_info "Batch ID: $BATCH_ID"
        log_info "Status URL: $STATUS_URL"
        return 0
    else
        log_error "Failed to upload batch (HTTP $http_code)"
        echo "$response_body"
        return 1
    fi
}

# Get batch status
get_batch_status() {
    log_info "Getting batch status..."
    
    response=$(curl -s -w "\n%{http_code}" "$BASE_URL/client/$CLIENT_ID/document/batch/$BATCH_ID" 2>/dev/null)
    
    http_code=$(echo "$response" | tail -n1)
    response_body=$(echo "$response" | head -n -1)
    
    if [ "$http_code" = "200" ]; then
        status=$(echo "$response_body" | grep -o '"status":"[^"]*"' | cut -d'"' -f4)
        filename=$(echo "$response_body" | grep -o '"originalFilename":"[^"]*"' | cut -d'"' -f4)
        log_info "Batch status: $status"
        log_info "Original filename: $filename"
        echo "$response_body" | python3 -m json.tool 2>/dev/null || echo "$response_body"
        return 0
    else
        log_error "Failed to get batch status (HTTP $http_code)"
        echo "$response_body"
        return 1
    fi
}

# List batches for client
list_batches() {
    log_info "Listing batches for client..."
    
    response=$(curl -s -w "\n%{http_code}" "$BASE_URL/client/$CLIENT_ID/document/batch?limit=10&offset=0" 2>/dev/null)
    
    http_code=$(echo "$response" | tail -n1)
    response_body=$(echo "$response" | head -n -1)
    
    if [ "$http_code" = "200" ]; then
        batch_count=$(echo "$response_body" | grep -o '"totalCount":[0-9]*' | cut -d':' -f2)
        log_info "Found $batch_count batches"
        echo "$response_body" | python3 -m json.tool 2>/dev/null || echo "$response_body"
        return 0
    else
        log_error "Failed to list batches (HTTP $http_code)"
        echo "$response_body"
        return 1
    fi
}

# Test invalid file upload
test_invalid_file() {
    log_info "Testing invalid file upload (should fail)..."
    
    # Create invalid file
    echo "This is not a ZIP file" > invalid.txt
    
    response=$(curl -s -w "\n%{http_code}" -X POST "$BASE_URL/client/$CLIENT_ID/document/batch" \
        -F "archive=@invalid.txt" 2>/dev/null)
    
    http_code=$(echo "$response" | tail -n1)
    response_body=$(echo "$response" | head -n -1)
    
    # Cleanup
    rm -f invalid.txt
    
    if [ "$http_code" = "400" ]; then
        log_info "Invalid file correctly rejected (HTTP $http_code)"
        return 0
    else
        log_warn "Expected HTTP 400 for invalid file, got HTTP $http_code"
        echo "$response_body"
        return 1
    fi
}

# Cancel batch
cancel_batch() {
    log_info "Canceling batch..."
    
    response=$(curl -s -w "\n%{http_code}" -X DELETE "$BASE_URL/client/$CLIENT_ID/document/batch/$BATCH_ID" 2>/dev/null)
    
    http_code=$(echo "$response" | tail -n1)
    
    if [ "$http_code" = "204" ]; then
        log_info "Batch canceled successfully"
        return 0
    elif [ "$http_code" = "404" ]; then
        log_warn "Batch not found or already completed (HTTP $http_code)"
        return 0
    else
        log_error "Failed to cancel batch (HTTP $http_code)"
        return 1
    fi
}

# Cleanup function
cleanup() {
    log_info "Cleaning up..."
    rm -f "$ZIP_FILE" invalid.txt
    log_info "Cleanup complete"
}

# Main test flow
main() {
    log_info "Starting ZIP batch upload test..."
    log_info "Using client ID: $CLIENT_ID"
    
    # Setup trap for cleanup
    trap cleanup EXIT
    
    # Run tests
    check_service
    create_test_zip
    create_client
    
    if upload_batch; then
        sleep 1  # Brief pause to let the system process
        get_batch_status
        list_batches
        test_invalid_file
        cancel_batch
        
        # Verify cancellation
        sleep 1
        log_info "Verifying batch was canceled..."
        get_batch_status
        
        log_info "✅ All tests completed successfully!"
    else
        log_error "❌ Batch upload failed, skipping remaining tests"
        exit 1
    fi
}

# Run main function
main "$@"