#!/bin/bash

# Test script for ZIP batch upload functionality
#!/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
# - GET /document/{DOCUMENT_ID}/download                  - Download a processed document
# - GET /document/{FAKE_ID}/download                      - Verify 404 for non-existent document

set -e  # Exit on any error

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

# 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 real PDFs from test.pdf.batch directory, each in its own subfolder
create_test_zip() {
    log_info "Creating test ZIP file with real PDFs from ./test.pdf.batch, each in its own subfolder..."

    # Remove existing file if it exists
    rm -f "$ZIP_FILE"

    # Check if test.pdf.batch directory exists
    pdf_dir="./test.pdf.batch"
    if [ ! -d "$pdf_dir" ]; then
        log_error "Directory $pdf_dir does not exist"
        exit 1
    fi

    # Check if there are PDF files in the directory
    pdf_count=$(find "$pdf_dir" -name "*.pdf" -type f | wc -l)
    if [ "$pdf_count" -eq 0 ]; then
        log_error "No PDF files found in $pdf_dir"
        exit 1
    fi

    # Create temporary directory structure
    temp_dir="temp_zip_structure"
    rm -rf "$temp_dir"
    mkdir -p "$temp_dir"

    # Copy each PDF to its own subfolder
    folder_num=1
    for pdf_file in "$pdf_dir"/*.pdf; do
        if [ -f "$pdf_file" ]; then
            folder_name="folder$folder_num"
            mkdir -p "$temp_dir/$folder_name"
            cp "$pdf_file" "$temp_dir/$folder_name/"
            log_info "Placed $(basename "$pdf_file") in $folder_name/"
            folder_num=$((folder_num + 1))
        fi
    done

    # Create ZIP file with the folder structure
    current_dir=$(pwd)
    cd "$temp_dir"
    zip -q -r "$current_dir/$ZIP_FILE" .
    cd "$current_dir"

    # Cleanup temporary 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 $pdf_count real PDFs in separate subfolders ($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
}

# Enable CanSync flag for client to allow document processing
enable_client_sync() {
    log_info "Enabling CanSync flag for client: $CLIENT_ID"

    # Check if entry already exists
    existing=$(docker exec deployments-db-1 psql -U postgres -d query_orchestration -t -c \
        "SELECT COUNT(*) FROM clientcansync WHERE clientid='$CLIENT_ID';" 2>/dev/null | tr -d ' ')

    if [ "$existing" = "0" ]; then
        # Insert new record
        docker exec deployments-db-1 psql -U postgres -d query_orchestration -c \
            "INSERT INTO clientcansync (clientid, cansync) VALUES ('$CLIENT_ID', true);" >/dev/null 2>&1
    else
        # Update existing record
        docker exec deployments-db-1 psql -U postgres -d query_orchestration -c \
            "UPDATE clientcansync SET cansync = true WHERE clientid='$CLIENT_ID';" >/dev/null 2>&1
    fi

    if [ $? -eq 0 ]; then
        log_info "CanSync flag enabled - documents can now be processed"
        return 0
    else
        log_error "Failed to enable CanSync flag"
        return 1
    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
}

# Wait for at least one document in the batch to have a document_id assigned.
# Polls the batch status endpoint until a document_id appears or timeout is reached.
wait_for_document_id() {
    log_info "Waiting for documents to be processed (up to 60s)..."

    local max_wait=60
    local elapsed=0
    local interval=3

    while [ $elapsed -lt $max_wait ]; do
        response=$(curl -s "$BASE_URL/client/$CLIENT_ID/document/batch/$BATCH_ID" 2>/dev/null)

        # Extract first non-null document_id from document_outcomes array
        DOCUMENT_ID=$(echo "$response" | python3 -c "
import sys, json
data = json.load(sys.stdin)
for outcome in data.get('document_outcomes', []):
    doc_id = outcome.get('document_id')
    if doc_id:
        print(doc_id)
        sys.exit(0)
sys.exit(1)
" 2>/dev/null)

        if [ $? -eq 0 ] && [ -n "$DOCUMENT_ID" ]; then
            log_info "Document ready for download: $DOCUMENT_ID"
            return 0
        fi

        sleep $interval
        elapsed=$((elapsed + interval))
        log_info "Still waiting... ($elapsed/${max_wait}s)"
    done

    log_error "Timed out waiting for documents to be processed"
    return 1
}

# Test downloading a document via GET /document/{id}/download.
# Verifies HTTP 200, Content-Disposition header, Content-Type, non-empty body, and PDF magic bytes.
test_document_download() {
    log_info "Testing document download: GET /document/$DOCUMENT_ID/download"

    local output_file="downloaded_document.pdf"

    # Download with headers captured
    curl -s -D headers.txt -o "$output_file" \
        "$BASE_URL/document/$DOCUMENT_ID/download" 2>/dev/null

    http_code=$(grep -m1 "^HTTP/" headers.txt | awk '{print $2}')

    if [ "$http_code" != "200" ]; then
        log_error "Download failed (HTTP $http_code)"
        cat headers.txt 2>/dev/null
        rm -f "$output_file" headers.txt
        return 1
    fi

    log_info "Download returned HTTP 200"

    # Verify Content-Disposition header is present
    if grep -qi "Content-Disposition" headers.txt; then
        disposition=$(grep -i "Content-Disposition" headers.txt | tr -d '\r')
        log_info "$disposition"
    else
        log_error "Missing Content-Disposition header"
        rm -f "$output_file" headers.txt
        return 1
    fi

    # Log Content-Type
    content_type=$(grep -i "Content-Type" headers.txt | head -1 | tr -d '\r')
    log_info "$content_type"

    # Verify file is non-empty
    file_size=$(stat -f%z "$output_file" 2>/dev/null || stat -c%s "$output_file" 2>/dev/null || echo "0")
    if [ "$file_size" -eq 0 ]; then
        log_error "Downloaded file is empty"
        rm -f "$output_file" headers.txt
        return 1
    fi
    log_info "Downloaded file size: $file_size bytes"

    # Verify file starts with PDF magic bytes (%PDF)
    file_header=$(head -c 4 "$output_file")
    if [ "$file_header" = "%PDF" ]; then
        log_info "File is a valid PDF (starts with %%PDF header)"
    else
        log_warn "File does not start with %%PDF header (got: $file_header)"
    fi

    # Cleanup
    rm -f "$output_file" headers.txt

    log_info "Document download test PASSED"
    return 0
}

# Test that downloading a non-existent document returns 404.
test_download_not_found() {
    log_info "Testing download of non-existent document (should return 404)..."

    local fake_id="00000000-0000-0000-0000-000000000000"

    response=$(curl -s -w "\n%{http_code}" \
        "$BASE_URL/document/$fake_id/download" 2>/dev/null)

    http_code=$(echo "$response" | tail -n1)

    if [ "$http_code" = "404" ]; then
        log_info "Non-existent document correctly returned 404"
        return 0
    else
        log_warn "Expected HTTP 404 for non-existent document, got HTTP $http_code"
        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 downloaded_document.pdf headers.txt
    rm -rf temp_zip_structure
    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
    enable_client_sync
    
    if upload_batch; then
        sleep 1  # Brief pause to let the system process
        get_batch_status
        list_batches

        # Document download tests -- must run before cancel_batch
        if wait_for_document_id; then
            test_document_download
            test_download_not_found
        else
            log_warn "Skipping download tests - no documents were processed in time"
        fi

        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 "$@"