#!/bin/bash

set -euo pipefail

SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
TEST_DIR="${SCRIPT_DIR}/test_tmp"
TEST_TOKEN="test-token-$(date +%s)"
CONTAINER_NAME="sos-test-$(date +%s)"
TEST_PORT=$((7700 + RANDOM % 100))
CLEANUP_NEEDED=false

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

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

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

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

function cleanup() {
    if [ "$CLEANUP_NEEDED" = true ]; then
        log_info "Cleaning up test environment..."
        
        # Stop and remove container
        if docker ps -a --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
            docker stop "${CONTAINER_NAME}" >/dev/null 2>&1 || true
            docker rm "${CONTAINER_NAME}" >/dev/null 2>&1 || true
            log_info "Removed test container"
        fi
        
        # Clean up test directory
        if [ -d "${TEST_DIR}" ]; then
            rm -rf "${TEST_DIR}"
            log_info "Removed test directory"
        fi
    fi
}

trap cleanup EXIT

function die() {
    log_error "$@"
    exit 1
}

function check_dependencies() {
    log_info "Checking dependencies..."
    
    # Check for Docker
    if ! command -v docker &> /dev/null; then
        die "Docker is required but not installed"
    fi
    
    # Check for jq
    if ! command -v jq &> /dev/null; then
        die "jq is required but not installed"
    fi
    
    # Check for curl
    if ! command -v curl &> /dev/null; then
        die "curl is required but not installed"
    fi
    
    # Check Docker daemon is running
    if ! docker info >/dev/null 2>&1; then
        die "Docker daemon is not running"
    fi
    
    log_info "All dependencies satisfied"
}

function setup_test_environment() {
    log_info "Setting up test environment..."
    CLEANUP_NEEDED=true
    
    # Create test directory
    rm -rf "${TEST_DIR}"
    mkdir -p "${TEST_DIR}/storage"
    mkdir -p "${TEST_DIR}/test_files"
    
    # Generate hashed token
    log_info "Generating authentication token..."
    HASHED_TOKEN=$(docker run --rm gitea.jde.nz/public/simple-object-server /sos/hash_token "${TEST_TOKEN}" 2>/dev/null | grep '^\$2[aby]\$' | head -1)
    
    if [ -z "${HASHED_TOKEN}" ]; then
        die "Failed to generate hashed token"
    fi
    
    # Create test configuration
    cat > "${TEST_DIR}/sos_config.json" <<EOF
{
  "write_tokens": ["${HASHED_TOKEN}"]
}
EOF
    
    log_info "Created test configuration"
    
    # Create test files
    echo "This is test file 1" > "${TEST_DIR}/test_files/test1.txt"
    echo "This is test file 2 with more content" > "${TEST_DIR}/test_files/test2.txt"
    dd if=/dev/urandom of="${TEST_DIR}/test_files/binary_test.bin" bs=1024 count=10 2>/dev/null
    
    # Create the sos binary for testing
    cp "${SCRIPT_DIR}/sos" "${TEST_DIR}/test_files/sos_binary"
    
    log_info "Created test files"
}

function start_test_server() {
    log_info "Starting SOS test server on port ${TEST_PORT}..."
    
    # Start the container (server runs on port 80 inside container)
    local container_id=$(docker run -d \
        --name "${CONTAINER_NAME}" \
        -p "${TEST_PORT}:80" \
        -v "${TEST_DIR}/sos_config.json:/data/sos_config.json:ro" \
        -v "${TEST_DIR}/storage:/data/storage" \
        gitea.jde.nz/public/simple-object-server 2>&1)
    
    if [ $? -ne 0 ]; then
        log_error "Failed to start container: ${container_id}"
        die "Failed to start test server"
    fi
    
    log_info "Container started with ID: ${container_id:0:12}"
    
    # Give container a moment to initialize
    sleep 2
    
    # Check if container is still running
    if ! docker ps --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
        log_error "Container stopped unexpectedly. Logs:"
        docker logs "${CONTAINER_NAME}" 2>&1
        die "Container failed to stay running"
    fi
    
    # Wait for server to be ready
    log_info "Waiting for server to be ready..."
    local max_attempts=30
    local attempt=0
    
    while [ $attempt -lt $max_attempts ]; do
        # Try multiple endpoints
        if curl -s "http://localhost:${TEST_PORT}/health" >/dev/null 2>&1 || \
           curl -s "http://localhost:${TEST_PORT}/" >/dev/null 2>&1; then
            log_info "Server is ready!"
            return 0
        fi
        
        # Check container is still running every 5 attempts
        if [ $((attempt % 5)) -eq 0 ]; then
            if ! docker ps --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
                log_error "Container stopped while waiting. Logs:"
                docker logs "${CONTAINER_NAME}" 2>&1 | tail -20
                die "Container stopped unexpectedly"
            fi
        fi
        
        sleep 1
        attempt=$((attempt + 1))
    done
    
    log_error "Server did not become ready. Container logs:"
    docker logs "${CONTAINER_NAME}" 2>&1 | tail -20
    die "Server failed to start after ${max_attempts} seconds"
}

function test_upload() {
    local file="$1"
    local label="$2"
    shift 2
    local extra_labels="$@"
    
    log_info "Testing upload: ${file} with label ${label}"
    
    export SOS_WRITE_TOKEN="${TEST_TOKEN}"
    # Override to use HTTP for local testing
    export SOS_TEST_MODE=1
    "${SCRIPT_DIR}/sos" upload "localhost:${TEST_PORT}" "${file}" "${label}" ${extra_labels} 2>&1 | tee "${TEST_DIR}/upload_output.txt"
    
    if [ ${PIPESTATUS[0]} -ne 0 ]; then
        log_error "Upload failed for ${file}"
        return 1
    fi
    
    # Verify upload was successful
    if grep -q "Download URL:" "${TEST_DIR}/upload_output.txt"; then
        log_info "Upload successful for ${file}"
        return 0
    else
        log_error "Upload verification failed for ${file}"
        return 1
    fi
}

function test_retrieval() {
    local identifier="$1"
    local expected_content="$2"
    
    log_info "Testing retrieval by ${identifier}"
    
    # Try to retrieve the file
    local response=$(curl -s "http://localhost:${TEST_PORT}/${identifier}")
    
    if [ -z "${response}" ]; then
        log_error "Failed to retrieve by ${identifier}"
        return 1
    fi
    
    # For hash retrieval, test the actual hash endpoint
    if [[ "${identifier}" =~ ^[a-f0-9]{64}$ ]]; then
        local meta_response=$(curl -s "http://localhost:${TEST_PORT}/meta/${identifier}")
        if echo "${meta_response}" | jq -e '.metadata' >/dev/null 2>&1; then
            log_info "Successfully retrieved metadata for hash ${identifier}"
        else
            log_error "Failed to retrieve metadata for hash ${identifier}"
            return 1
        fi
    fi
    
    log_info "Retrieval successful for ${identifier}"
    return 0
}

function test_duplicate_upload() {
    log_info "Testing duplicate file upload (deduplication)..."
    
    # Create a duplicate file with different name
    cp "${TEST_DIR}/test_files/test1.txt" "${TEST_DIR}/test_files/test1_dup.txt"
    
    export SOS_WRITE_TOKEN="${TEST_TOKEN}"
    export SOS_TEST_MODE=1
    
    # Upload original
    "${SCRIPT_DIR}/sos" upload "localhost:${TEST_PORT}" "${TEST_DIR}/test_files/test1.txt" "dup:original" >/dev/null 2>&1
    
    # Upload duplicate (should detect existing file)
    local output=$("${SCRIPT_DIR}/sos" upload "localhost:${TEST_PORT}" "${TEST_DIR}/test_files/test1_dup.txt" "dup:copy" 2>&1)
    
    if echo "${output}" | grep -q "File already exists, skipping upload"; then
        log_info "Deduplication working correctly"
        return 0
    else
        log_error "Deduplication test failed"
        return 1
    fi
}

function test_metadata_update() {
    log_info "Testing metadata update..."
    
    export SOS_WRITE_TOKEN="${TEST_TOKEN}"
    export SOS_TEST_MODE=1
    
    # Upload a file
    "${SCRIPT_DIR}/sos" upload "localhost:${TEST_PORT}" "${TEST_DIR}/test_files/test2.txt" "meta:v1" >/dev/null 2>&1
    
    # Upload same file with different metadata
    local output=$("${SCRIPT_DIR}/sos" upload "localhost:${TEST_PORT}" "${TEST_DIR}/test_files/test2.txt" "meta:v2" "meta:updated" 2>&1)
    
    if echo "${output}" | grep -q "File already exists"; then
        log_info "Metadata update successful"
        return 0
    else
        log_error "Metadata update test failed"
        return 1
    fi
}

function test_exists_endpoint() {
    log_info "Testing exists endpoint..."
    
    # Get hash of an uploaded file
    local hash=$(curl -s "http://localhost:${TEST_PORT}/hash/test:file1" | jq -r '.hash')
    
    if [ -z "${hash}" ] || [ "${hash}" = "null" ]; then
        log_error "Failed to get hash for exists test"
        return 1
    fi
    
    # Test exists endpoint
    local exists=$(curl -s "http://localhost:${TEST_PORT}/exists/${hash}" | jq -r '.exists')
    
    if [ "${exists}" = "true" ]; then
        log_info "Exists endpoint working correctly"
        return 0
    else
        log_error "Exists endpoint test failed"
        return 1
    fi
}

function test_invalid_auth() {
    log_info "Testing invalid authentication..."
    
    # Create a unique test file for this test
    echo "auth test content" > "${TEST_DIR}/test_files/auth_test.txt"
    
    export SOS_WRITE_TOKEN="invalid-token"
    export SOS_TEST_MODE=1
    
    # Run upload with invalid token and capture output + exit code
    set +e  # Temporarily disable exit on error
    "${SCRIPT_DIR}/sos" upload "localhost:${TEST_PORT}" "${TEST_DIR}/test_files/auth_test.txt" "auth:test" 2>&1 | tee "${TEST_DIR}/auth_test_output.txt"
    local exit_code=${PIPESTATUS[0]}
    set -e  # Re-enable exit on error
    
    local output=$(cat "${TEST_DIR}/auth_test_output.txt")
    
    # Check if it failed with non-zero exit code and error message
    if [ ${exit_code} -ne 0 ] && (echo "${output}" | grep -q "Failed to upload" || echo "${output}" | grep -q "Invalid write token"); then
        log_info "Invalid auth correctly rejected"
        return 0
    else
        log_error "Invalid auth test failed - upload should have been rejected (exit code: ${exit_code})"
        echo "Output was: ${output}"
        return 1
    fi
}

function run_tests() {
    local total_tests=0
    local passed_tests=0
    
    # Disable exit on error for test functions
    set +e
    
    log_info "Starting test suite..."
    echo ""
    
    # Test 1: Basic upload
    ((total_tests++))
    test_upload "${TEST_DIR}/test_files/test1.txt" "test:file1"
    if [ $? -eq 0 ]; then
        ((passed_tests++))
    fi
    echo ""
    
    # Test 2: Upload with multiple labels
    ((total_tests++))
    test_upload "${TEST_DIR}/test_files/test2.txt" "test:file2" "version:1.0" "env:test"
    if [ $? -eq 0 ]; then
        ((passed_tests++))
    fi
    echo ""
    
    # Test 3: Binary file upload
    ((total_tests++))
    test_upload "${TEST_DIR}/test_files/binary_test.bin" "binary:test"
    if [ $? -eq 0 ]; then
        ((passed_tests++))
    fi
    echo ""
    
    # Test 4: Upload sos binary itself
    ((total_tests++))
    test_upload "${TEST_DIR}/test_files/sos_binary" "sos:test" "sos:latest"
    if [ $? -eq 0 ]; then
        ((passed_tests++))
    fi
    echo ""
    
    # Test 5: Retrieval by label
    ((total_tests++))
    test_retrieval "test:file1" "This is test file 1"
    if [ $? -eq 0 ]; then
        ((passed_tests++))
    fi
    echo ""
    
    # Test 6: Deduplication
    ((total_tests++))
    test_duplicate_upload
    if [ $? -eq 0 ]; then
        ((passed_tests++))
    fi
    echo ""
    
    # Test 7: Metadata update
    ((total_tests++))
    test_metadata_update
    if [ $? -eq 0 ]; then
        ((passed_tests++))
    fi
    echo ""
    
    # Test 8: Exists endpoint
    ((total_tests++))
    test_exists_endpoint
    if [ $? -eq 0 ]; then
        ((passed_tests++))
    fi
    echo ""
    
    # Test 9: Invalid authentication
    ((total_tests++))
    test_invalid_auth
    if [ $? -eq 0 ]; then
        ((passed_tests++))
    fi
    echo ""
    
    # Re-enable exit on error
    set -e
    
    # Summary
    echo "=================================="
    if [ ${passed_tests} -eq ${total_tests} ]; then
        echo -e "${GREEN}All tests passed!${NC} (${passed_tests}/${total_tests})"
        return 0
    else
        echo -e "${RED}Some tests failed!${NC} (${passed_tests}/${total_tests})"
        return 1
    fi
}

function main() {
    echo "SOS Test Suite"
    echo "=============="
    echo ""
    
    check_dependencies
    setup_test_environment
    start_test_server
    
    echo ""
    run_tests
    local test_result=$?
    
    echo ""
    log_info "Test suite completed"
    
    # Show container logs if tests failed
    if [ ${test_result} -ne 0 ]; then
        log_warning "Showing server logs for debugging:"
        docker logs "${CONTAINER_NAME}" 2>&1 | tail -20
    fi
    
    exit ${test_result}
}

main "$@"