#!/bin/bash

# ============================================================================ #
# Copyright (c) 2022 - 2026 NVIDIA Corporation & Affiliates.                   #
# All rights reserved.                                                         #
#                                                                              #
# This source code and the accompanying materials are made available under     #
# the terms of the Apache License 2.0 which accompanies this distribution.     #
# ============================================================================ #

checkServerConnection() {
  PYTHONPATH=@CMAKE_BINARY_DIR@/python @Python_EXECUTABLE@ - << EOF
import socket
try:
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.connect(("localhost", 62453))
    s.close()
except Exception:
    exit(1)
EOF
}

# Launch the fake server
if [[ "$OSTYPE" == "darwin"* ]]; then
  LIB_EXT="dylib"
else
  LIB_EXT="so"
fi


PYTHONPATH=@CMAKE_BINARY_DIR@/python:@CMAKE_SOURCE_DIR@/python/tests/utils/mock_qpu @Python_EXECUTABLE@ @CMAKE_SOURCE_DIR@/unittests/nvqpp/backends/quake_backend/mock_server.py -device-lib @CMAKE_BINARY_DIR@/unittests/nvqpp/backends/quake_backend/libfake_quake_backend_device_funcs.$LIB_EXT &
# we'll need the process id to kill it
pid=$(echo "$!")
n=0
while ! checkServerConnection; do
  sleep 1
  n=$((n+1))
  if [ "$n" -eq "60" ]; then
    echo "Failed to start the server after 20 seconds"
    # kill the server
    kill -INT $pid
    exit 99
  fi
done

echo "Server started successfully"
test_err_sum=0
# Test the app nvq++ (remote execution on the mock server)
PATH=@CMAKE_BINARY_DIR@/bin:$PATH nvq++ --target quake_fake @CMAKE_SOURCE_DIR@/unittests/nvqpp/backends/quake_backend/test_app.cpp -o test_app
if [ $? -ne 0 ]; then
  echo ":x: nvq++ compilation failed"
  test_err_sum=$((test_err_sum+1))
else
  echo ":white_check_mark: Successfully compiled test_app with nvq++"
  # Run the test app
  ./test_app
  if [ $? -ne 0 ]; then
    echo ":x: test_app failed"
    test_err_sum=$((test_err_sum+1))
  else
    echo ":white_check_mark: Successfully ran test_app"
  fi
fi

# Loop remote execution test. CUDAQ_DUMP_JIT_IR validates that the client
# preserves cc.loop operations, and the mock server validates and executes the
# submitted payload.
PATH=@CMAKE_BINARY_DIR@/bin:$PATH nvq++ --target quake_fake @CMAKE_SOURCE_DIR@/unittests/nvqpp/backends/quake_backend/test_loop_app.cpp -o test_loop_app
if [ $? -ne 0 ]; then
  echo ":x: nvq++ compilation failed for test_loop_app"
  test_err_sum=$((test_err_sum+1))
else
  echo ":white_check_mark: Successfully compiled test_loop_app with nvq++"
  CUDAQ_DUMP_JIT_IR=1 ./test_loop_app > test_loop_app.out 2>&1
  if [ $? -ne 0 ]; then
    echo ":x: test_loop_app failed"
    cat test_loop_app.out
    test_err_sum=$((test_err_sum+1))
  else
    loop_count=$(grep -c "cc.loop" test_loop_app.out)
    # Expected loops: the nested for lowers to two cc.loop ops, and the while
    # lowers to one.
    if [ "$loop_count" -lt 3 ]; then
      echo ":x: test_loop_app client IR did not preserve expected cc.loop ops"
      cat test_loop_app.out
      test_err_sum=$((test_err_sum+1))
    else
      echo ":white_check_mark: Successfully ran test_loop_app with preserved client loops"
    fi
  fi
fi

# Emulation test
# Note: for emulation, we need to link in the device library. 
PATH=@CMAKE_BINARY_DIR@/bin:$PATH nvq++ --target quake_fake --emulate @CMAKE_SOURCE_DIR@/unittests/nvqpp/backends/quake_backend/test_app.cpp -o test_app @CMAKE_BINARY_DIR@/unittests/nvqpp/backends/quake_backend/libfake_quake_backend_device_funcs.$LIB_EXT
if [ $? -ne 0 ]; then
  echo ":x: nvq++ compilation failed"
  test_err_sum=$((test_err_sum+1))
else
  echo ":white_check_mark: Successfully compiled test_app (emulation) with nvq++"
  # Run the test app
  ./test_app
  if [ $? -ne 0 ]; then
    echo ":x: test_app failed"
    test_err_sum=$((test_err_sum+1))
  else
    echo ":white_check_mark: Successfully ran test_app (emulation)"
  fi
fi

# Loop emulation test. This covers structured control flow that survives the
# target pipeline before local wire-set-to-QIR lowering.
PATH=@CMAKE_BINARY_DIR@/bin:$PATH nvq++ --target quake_fake --emulate @CMAKE_SOURCE_DIR@/unittests/nvqpp/backends/quake_backend/test_loop_app.cpp -o test_loop_app_emulate
if [ $? -ne 0 ]; then
  echo ":x: nvq++ compilation failed for test_loop_app (emulation)"
  test_err_sum=$((test_err_sum+1))
else
  echo ":white_check_mark: Successfully compiled test_loop_app (emulation) with nvq++"
  ./test_loop_app_emulate
  if [ $? -ne 0 ]; then
    echo ":x: test_loop_app failed (emulation)"
    test_err_sum=$((test_err_sum+1))
  else
    echo ":white_check_mark: Successfully ran test_loop_app (emulation)"
  fi
fi

# Run the Python test
PYTHONPATH=@CMAKE_BINARY_DIR@/python @Python_EXECUTABLE@ @CMAKE_SOURCE_DIR@/unittests/nvqpp/backends/quake_backend/test_app.py
if [ $? -ne 0 ]; then
  echo ":x: Python test_app.py failed"
  test_err_sum=$((test_err_sum+1))
else
  echo ":white_check_mark: Successfully ran Python test_app.py"
fi

# Python loop remote execution test. CUDAQ_DUMP_JIT_IR validates that the Python
# client preserves cc.loop operations for the nop codegen path.
CUDAQ_DUMP_JIT_IR=1 PYTHONPATH=@CMAKE_BINARY_DIR@/python @Python_EXECUTABLE@ @CMAKE_SOURCE_DIR@/unittests/nvqpp/backends/quake_backend/test_loop_app.py > test_loop_app.py.out 2>&1
if [ $? -ne 0 ]; then
  echo ":x: Python test_loop_app.py failed"
  cat test_loop_app.py.out
  test_err_sum=$((test_err_sum+1))
else
  loop_count=$(grep -c "cc.loop" test_loop_app.py.out)
  if [ "$loop_count" -lt 3 ]; then
    echo ":x: Python test_loop_app.py client IR did not preserve expected cc.loop ops"
    cat test_loop_app.py.out
    test_err_sum=$((test_err_sum+1))
  else
    echo ":white_check_mark: Successfully ran Python test_loop_app.py with preserved client loops"
  fi
fi

# kill the server
kill -INT $pid
# return success / failure
if [ ! $test_err_sum -eq 0 ]; then
  echo "::error::${test_err_sum} tests failed."
  exit 1
fi
