#!/bin/bash

# SigNoz MCP Server Test Runner
# This script runs the complete test suite for the MCP server

set -e

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

# Get the directory of this script
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"

cd "$PROJECT_ROOT"

echo -e "${BLUE}🧪 SigNoz MCP Server Test Suite${NC}"
echo "========================================"

# Function to print section headers
print_section() {
    echo
    echo -e "${BLUE}$1${NC}"
    echo "----------------------------------------"
}

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

# Function to print warning messages
print_warning() {
    echo -e "${YELLOW}⚠ $1${NC}"
}

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

# Check if Go is installed
if ! command -v go &> /dev/null; then
    print_error "Go is not installed or not in PATH"
    exit 1
fi

print_success "Go found: $(go version)"

# Set up test environment (use defaults if not provided)
export SIGNOZ_MCP_SIGNOZ_URL="${SIGNOZ_MCP_SIGNOZ_URL:-http://localhost:3301}"
export SIGNOZ_MCP_SIGNOZ_API_KEY="${SIGNOZ_MCP_SIGNOZ_API_KEY:-test-api-key}"
export SIGNOZ_MCP_LOG_LEVEL="${SIGNOZ_MCP_LOG_LEVEL:-error}"  # Reduce log noise during tests

# Parse command line arguments
VERBOSE=false
COVERAGE=false
BENCHMARKS=false
INTEGRATION=false
UNIT_ONLY=false
QUICK=false
PACKAGE=""

while [[ $# -gt 0 ]]; do
    case $1 in
        -v|--verbose)
            VERBOSE=true
            shift
            ;;
        -c|--coverage)
            COVERAGE=true
            shift
            ;;
        -b|--benchmarks)
            BENCHMARKS=true
            shift
            ;;
        -i|--integration)
            INTEGRATION=true
            shift
            ;;
        -u|--unit-only)
            UNIT_ONLY=true
            shift
            ;;
        -q|--quick)
            QUICK=true
            shift
            ;;
        -p|--package)
            PACKAGE="$2"
            shift 2
            ;;
        -h|--help)
            echo "Usage: $0 [OPTIONS]"
            echo "Options:"
            echo "  -v, --verbose      Enable verbose output"
            echo "  -c, --coverage     Generate coverage report"
            echo "  -b, --benchmarks   Run benchmark tests"
            echo "  -i, --integration  Run integration tests only"
            echo "  -u, --unit-only    Run unit tests only"
            echo "  -q, --quick        Quick mode (skip slow validations, for pre-commit)"
            echo "  -p, --package PKG  Run tests for specific package"
            echo "  -h, --help         Show this help message"
            exit 0
            ;;
        *)
            print_error "Unknown option: $1"
            exit 1
            ;;
    esac
done

# Build test flags
TEST_FLAGS=""
if [ "$VERBOSE" = true ]; then
    TEST_FLAGS="$TEST_FLAGS -v"
fi

if [ "$COVERAGE" = true ]; then
    TEST_FLAGS="$TEST_FLAGS -coverprofile=coverage.out -covermode=atomic"
fi

# Determine which tests to run
if [ -n "$PACKAGE" ]; then
    TEST_PACKAGES="./$PACKAGE"
elif [ "$INTEGRATION" = true ]; then
    TEST_PACKAGES="./cmd/server"
elif [ "$UNIT_ONLY" = true ]; then
    TEST_PACKAGES="./pkg/... ./internal/..."
else
    TEST_PACKAGES="./..."
fi

print_section "🏗️  Building Project"

# Build the project first
if go build -o bin/signoz-mcp-server cmd/server/main.go; then
    print_success "Project built successfully"
else
    print_error "Build failed"
    exit 1
fi

print_section "🔍 Running Go Vet"

# Run go vet
if go vet $TEST_PACKAGES; then
    print_success "Go vet passed"
else
    print_error "Go vet failed"
    exit 1
fi

print_section "📋 Running Tests"

# Run tests
echo "Running tests for: $TEST_PACKAGES"
if go test $TEST_FLAGS $TEST_PACKAGES; then
    print_success "All tests passed"
else
    print_error "Some tests failed"
    exit 1
fi

# Run benchmarks if requested
if [ "$BENCHMARKS" = true ]; then
    print_section "🏃 Running Benchmarks"
    
    if go test -bench=. -benchmem $TEST_PACKAGES; then
        print_success "Benchmarks completed"
    else
        print_warning "Some benchmarks failed"
    fi
fi

# Generate coverage report if requested
if [ "$COVERAGE" = true ] && [ -f coverage.out ]; then
    print_section "📊 Coverage Report"
    
    # Generate HTML coverage report
    go tool cover -html=coverage.out -o coverage.html
    
    # Show coverage summary
    COVERAGE_PERCENT=$(go tool cover -func=coverage.out | grep total | awk '{print $3}')
    echo "Total coverage: $COVERAGE_PERCENT"
    
    # Check coverage threshold
    COVERAGE_NUM=$(echo $COVERAGE_PERCENT | sed 's/%//')
    if (( $(echo "$COVERAGE_NUM >= 70" | bc -l) )); then
        print_success "Coverage is above 70% threshold"
    else
        print_warning "Coverage is below 70% threshold"
    fi
    
    print_success "Coverage report generated: coverage.html"
fi

# Initialize validation status tracking
MCP_VALIDATION_STATUS="skipped"
TOOL_VALIDATION_STATUS="skipped"

# Only run MCP protocol validation for full server tests, not for individual packages, and not in quick mode
if [ "$QUICK" = false ] && ([ "$TEST_PACKAGES" = "./..." ] || [[ "$TEST_PACKAGES" == *"./cmd/server"* ]]); then
    print_section "🧪 MCP Protocol Validation"

    # Test MCP protocol compliance
    if [[ "$SIGNOZ_MCP_SIGNOZ_URL" == *"signoz.cloud"* ]]; then
        echo "Testing MCP protocol compliance using real SigNoz instance..."
        echo "URL: $SIGNOZ_MCP_SIGNOZ_URL"
    else
        echo "Testing MCP protocol compliance using test credentials..."
        echo "(Use real credentials via environment variables for full validation)"
    fi

    # Create test input for MCP server
    TEST_INPUT=$(cat << 'EOF'
{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test-client","version":"1.0.0"}}}
{"jsonrpc":"2.0","id":2,"method":"tools/list"}
{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"signoz_test_connection","arguments":{}}}
EOF
)

    # Test the server with sample MCP messages  
    if echo "$TEST_INPUT" | SIGNOZ_MCP_SIGNOZ_URL="$SIGNOZ_MCP_SIGNOZ_URL" SIGNOZ_MCP_SIGNOZ_API_KEY="$SIGNOZ_MCP_SIGNOZ_API_KEY" ./bin/signoz-mcp-server stdio 2>/dev/null | grep -q '"jsonrpc":"2.0"'; then
        print_success "MCP protocol compliance validated"
        MCP_VALIDATION_STATUS="passed"
    else
        if [[ "$SIGNOZ_MCP_SIGNOZ_URL" == *"signoz.cloud"* ]]; then
            print_warning "MCP protocol validation failed (unexpected with real SigNoz - may indicate server bug)"
        else
            print_warning "MCP protocol validation failed (server may require real Signoz connection)"
        fi
        MCP_VALIDATION_STATUS="failed"
    fi
else
    print_section "🧪 MCP Protocol Validation"
    if [ "$QUICK" = true ]; then
        echo "Skipped (quick mode enabled for pre-commit)"
    else
        echo "Skipped (not testing full server - only testing: $TEST_PACKAGES)"
    fi
    MCP_VALIDATION_STATUS="skipped"
fi

# Only run tool validation for full server tests, not for individual packages, and not in quick mode
if [ "$QUICK" = false ] && ([ "$TEST_PACKAGES" = "./..." ] || [[ "$TEST_PACKAGES" == *"./cmd/server"* ]]); then
    print_section "🔧 Tool Validation"

    # Validate that all expected tools are present
    EXPECTED_TOOLS=(
        "signoz_test_connection"
        "signoz_fetch_dashboards"
        "signoz_fetch_dashboard_details"
        "signoz_fetch_dashboard_data"
        "signoz_fetch_apm_metrics"
        "signoz_fetch_services"
        "signoz_execute_clickhouse_query"
        "signoz_execute_builder_query"
        "signoz_fetch_traces_or_logs"
    )

    if [[ "$SIGNOZ_MCP_SIGNOZ_URL" == *"signoz.cloud"* ]]; then
        echo "Validating tool registration using real SigNoz instance..."
    else
        echo "Validating tool registration using test credentials..."
        echo "(Expect failures without real SigNoz connection)"
    fi
    # Create proper MCP sequence: initialize then tools/list
    TOOLS_VALIDATION_INPUT=$(cat << 'EOF'
{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test-client","version":"1.0.0"}}}
{"jsonrpc":"2.0","id":2,"method":"tools/list"}
EOF
)

    if TOOLS_OUTPUT=$(echo "$TOOLS_VALIDATION_INPUT" | SIGNOZ_MCP_SIGNOZ_URL="$SIGNOZ_MCP_SIGNOZ_URL" SIGNOZ_MCP_SIGNOZ_API_KEY="$SIGNOZ_MCP_SIGNOZ_API_KEY" ./bin/signoz-mcp-server stdio 2>/dev/null); then
        TOOLS_FOUND=0
        MISSING_TOOLS=()
        
        for tool in "${EXPECTED_TOOLS[@]}"; do
            if echo "$TOOLS_OUTPUT" | grep -q "\"$tool\""; then
                TOOLS_FOUND=$((TOOLS_FOUND + 1))
            else
                MISSING_TOOLS+=("$tool")
            fi
        done
        
        if [ $TOOLS_FOUND -eq ${#EXPECTED_TOOLS[@]} ]; then
            print_success "All $TOOLS_FOUND expected tools found"
            TOOL_VALIDATION_STATUS="passed"
        else
            if [ ${#MISSING_TOOLS[@]} -le 3 ]; then
                # Show individual missing tools if only a few
                for tool in "${MISSING_TOOLS[@]}"; do
                    print_warning "Tool not found: $tool"
                done
            else
                # Show summary - different message based on credential type
                if [[ "$SIGNOZ_MCP_SIGNOZ_URL" == *"signoz.cloud"* ]]; then
                    print_warning "${#MISSING_TOOLS[@]} tools not found (unexpected with real SigNoz connection)"
                else
                    print_warning "${#MISSING_TOOLS[@]} tools not found (expected without real SigNoz connection)"
                fi
                echo "Missing tools: ${MISSING_TOOLS[0]}, ${MISSING_TOOLS[1]}, ${MISSING_TOOLS[2]}, ... and $((${#MISSING_TOOLS[@]} - 3)) more"
            fi
            print_warning "Only $TOOLS_FOUND out of ${#EXPECTED_TOOLS[@]} tools found"
            TOOL_VALIDATION_STATUS="failed"
        fi
    else
        print_warning "Could not validate tools (may require real Signoz connection)"
        TOOL_VALIDATION_STATUS="failed"
    fi
else
    print_section "🔧 Tool Validation"
    if [ "$QUICK" = true ]; then
        echo "Skipped (quick mode enabled for pre-commit)"
    else
        echo "Skipped (not testing full server - only testing: $TEST_PACKAGES)"
    fi
    TOOL_VALIDATION_STATUS="skipped"
fi

print_section "📈 Test Summary"

echo "Test run completed successfully!"
echo
echo "Summary:"
echo "- Project built successfully"
echo "- Go vet passed"
echo "- All tests passed"

if [ "$BENCHMARKS" = true ]; then
    echo "- Benchmarks completed"
fi

if [ "$COVERAGE" = true ]; then
    echo "- Coverage report generated"
fi

# Determine if using real SigNoz credentials
USING_REAL_SIGNOZ=false
if [[ "$SIGNOZ_MCP_SIGNOZ_URL" == *"signoz.cloud"* ]]; then
    USING_REAL_SIGNOZ=true
fi

# Add MCP validation status to summary
case $MCP_VALIDATION_STATUS in
    "passed")
        echo "- MCP protocol compliance: ✓ Validated"
        ;;
    "failed")
        if [ "$USING_REAL_SIGNOZ" = true ]; then
            echo "- MCP protocol compliance: ⚠ Failed (unexpected with real SigNoz connection)"
        else
            echo "- MCP protocol compliance: ⚠ Failed (expected without real SigNoz connection)"
        fi
        ;;
    "skipped")
        echo "- MCP protocol compliance: ⏭ Skipped (package-specific test)"
        ;;
esac

# Add tool validation status to summary  
case $TOOL_VALIDATION_STATUS in
    "passed")
        echo "- Tool registration: ✓ All tools found"
        ;;
    "failed")
        if [ "$USING_REAL_SIGNOZ" = true ]; then
            echo "- Tool registration: ⚠ Failed (unexpected with real SigNoz connection)"
        else
            echo "- Tool registration: ⚠ Failed (expected without real SigNoz connection)"
        fi
        ;;
    "skipped")
        echo "- Tool registration: ⏭ Skipped (package-specific test)"
        ;;
esac

# Overall success determination
if [ "$MCP_VALIDATION_STATUS" = "failed" ] || [ "$TOOL_VALIDATION_STATUS" = "failed" ]; then
    if [ "$USING_REAL_SIGNOZ" = true ]; then
        print_warning "⚠ Tests passed, but server validation failed (unexpected with real SigNoz - may indicate server bug)"
    else
        print_warning "⚠ Tests passed, but server validation failed (expected without real SigNoz)"
    fi
else
    print_success "🎉 All checks passed!"
fi

# Clean up temporary files if not in coverage mode
if [ "$COVERAGE" != true ] && [ -f coverage.out ]; then
    rm coverage.out
fi

echo
echo "To run specific test categories:"
echo "  Unit tests only:     $0 --unit-only"
echo "  Integration tests:   $0 --integration"  
echo "  With coverage:       $0 --coverage"
echo "  With benchmarks:     $0 --benchmarks"
echo "  Quick mode:          $0 --quick"
echo "  Specific package:    $0 --package pkg/mcp"
echo