#!/bin/bash
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT

# CK Test - Build and test Composable Kernel in Docker

set -e
set -o pipefail

# Find script directory and load common utilities
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/common.sh"

# Initialize configuration
PROJECT_ROOT=$(get_project_root "${SCRIPT_DIR}")
CONTAINER_NAME=$(get_container_name "${PROJECT_ROOT}")

# Help message
show_help() {
    cat << EOF
CK Test - Build and test Composable Kernel in Docker

Usage: ck-test [options] <test_name> [test_options]

Options:
  -h, --help              Show this help message
  --name <name>           Specify container name
  --reconfigure           Reconfigure CMake before building
  --no-build              Skip building, run test directly

Arguments:
  test_name               Name of test executable (required)
  test_options            Additional options passed to test (e.g., --gtest_filter=*)

Environment:
  CK_CONTAINER_NAME - Override default container name
  GPU_TARGET        - Override GPU target detection (e.g., gfx950, gfx942)

Examples:
  ck-test test_amdgcn_mma
  ck-test test_amdgcn_mma --gtest_filter=*Fp16*
  ck-test --name my_container test_amdgcn_mma
  ck-test --reconfigure test_amdgcn_mma

EOF
}

# Parse arguments
test_name=""
reconfigure=false
no_build=false
test_options=()

while [[ $# -gt 0 ]]; do
    case $1 in
        -h|--help)
            show_help
            exit 0
            ;;
        --name)
            CONTAINER_NAME="$2"
            shift 2
            ;;
        --reconfigure)
            reconfigure=true
            shift
            ;;
        --no-build)
            no_build=true
            shift
            ;;
        --gtest_*|--help)
            test_options+=("$1")
            shift
            ;;
        *)
            if [ -z "$test_name" ]; then
                test_name="$1"
            else
                test_options+=("$1")
            fi
            shift
            ;;
    esac
done

# Validate test name
if [ -z "$test_name" ]; then
    echo "Error: test_name required"
    echo ""
    show_help
    exit 1
fi

# Ensure container is running
if ! container_is_running "${CONTAINER_NAME}"; then
    echo "Container '${CONTAINER_NAME}' not running. Starting..."
    "${SCRIPT_DIR}/ck-start" "${CONTAINER_NAME}"
    echo ""
fi

# Configure CMake if needed or requested
if [ "$reconfigure" = true ] || ! docker exec "${CONTAINER_NAME}" test -f /workspace/build/build.ninja 2>/dev/null; then
    echo "Detecting GPU target..."
    GPU_TARGET_DETECTED=$(detect_gpu_target "${CONTAINER_NAME}")

    if [ "$reconfigure" = true ]; then
        echo "Reconfiguring CMake from scratch for GPU target: ${GPU_TARGET_DETECTED}"
    else
        echo "Configuring build with CMake for GPU target: ${GPU_TARGET_DETECTED}"
    fi

    docker exec "${CONTAINER_NAME}" bash -c "
        cd /workspace || exit 1
        rm -rf /workspace/build
        mkdir /workspace/build
        cd /workspace/build || exit 1
        cmake .. -GNinja \
            -DGPU_TARGETS=${GPU_TARGET_DETECTED} \
            -DCMAKE_BUILD_TYPE=Release \
            -DCMAKE_CXX_COMPILER=/opt/rocm/llvm/bin/clang++ \
            -DBUILD_TESTING=ON 2>&1 | tail -30
    "
    echo ""
fi

# Build test if needed (unless --no-build is specified)
if [ "$no_build" = false ]; then
    if ! docker exec "${CONTAINER_NAME}" test -f "/workspace/build/bin/${test_name}" 2>/dev/null; then
        echo "Building ${test_name}..."
        docker exec "${CONTAINER_NAME}" bash -c "
            cd /workspace/build || exit 1
            ninja ${test_name} 2>&1
        "
        echo ""
    else
        echo "Test executable found, rebuilding to ensure latest version..."
        docker exec "${CONTAINER_NAME}" bash -c "
            cd /workspace/build || exit 1
            ninja ${test_name} 2>&1
        "
        echo ""
    fi
fi

# Run test
echo "Running: ${test_name} ${test_options[*]}"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

# Build the command with proper quoting
cmd="cd /workspace/build && ./bin/${test_name}"
for opt in "${test_options[@]}"; do
    cmd="${cmd} $(printf '%q' "$opt")"
done

docker exec "${CONTAINER_NAME}" bash -c "${cmd}"
exit_code=$?

echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
if [ $exit_code -eq 0 ]; then
    echo "Test completed successfully"
else
    echo "Test failed with exit code: ${exit_code}"
fi

exit $exit_code