// Copyright 2026 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

// Semaphore tests that require std::thread for cross-thread signaling.
// These are separated from semaphore_test.cc so platforms without C threading
// (wasm, bare-metal) can run the single-threaded semaphore tests without
// pulling in <thread>.

#include <chrono>
#include <thread>

#include "iree/hal/cts/util/test_base.h"

namespace iree::hal::cts {

using namespace std::chrono_literals;

class SemaphoreThreadTest : public CtsTestBase<> {};

// Tests waiting on a semaphore that signals past the desired value.
TEST_P(SemaphoreThreadTest, WaitLaterSignaledBeyond) {
  iree_hal_semaphore_t* semaphore = NULL;
  IREE_ASSERT_OK(
      iree_hal_semaphore_create(device_, IREE_HAL_QUEUE_AFFINITY_ANY, 2ull,
                                IREE_HAL_SEMAPHORE_FLAG_DEFAULT, &semaphore));

  std::thread thread([&]() {
    std::this_thread::sleep_for(100ms);
    // Signal beyond the desired value.
    IREE_ASSERT_OK(
        iree_hal_semaphore_signal(semaphore, 10ull, /*frontier=*/NULL));
  });

  IREE_ASSERT_OK(iree_hal_semaphore_wait(
      semaphore, 3ull, iree_make_deadline(IREE_TIME_INFINITE_FUTURE),
      IREE_ASYNC_WAIT_FLAG_NONE));
  thread.join();

  iree_hal_semaphore_release(semaphore);
}

TEST_P(SemaphoreThreadTest, WaitAnyLaterSignaled) {
  iree_hal_semaphore_t* semaphore_a = NULL;
  iree_hal_semaphore_t* semaphore_b = NULL;
  IREE_ASSERT_OK(
      iree_hal_semaphore_create(device_, IREE_HAL_QUEUE_AFFINITY_ANY, 0ull,
                                IREE_HAL_SEMAPHORE_FLAG_DEFAULT, &semaphore_a));
  IREE_ASSERT_OK(
      iree_hal_semaphore_create(device_, IREE_HAL_QUEUE_AFFINITY_ANY, 0ull,
                                IREE_HAL_SEMAPHORE_FLAG_DEFAULT, &semaphore_b));

  iree_hal_semaphore_t* semaphore_ptrs[] = {semaphore_a, semaphore_b};
  uint64_t payload_values[] = {1ull, 1ull};
  iree_hal_semaphore_list_t semaphore_list = {IREE_ARRAYSIZE(semaphore_ptrs),
                                              semaphore_ptrs, payload_values};

  std::thread thread([&]() {
    std::this_thread::sleep_for(100ms);
    IREE_ASSERT_OK(
        iree_hal_semaphore_signal(semaphore_b, 1ull, /*frontier=*/NULL));
  });

  IREE_ASSERT_OK(iree_hal_device_wait_semaphores(
      device_, IREE_ASYNC_WAIT_MODE_ANY, semaphore_list,
      iree_make_deadline(IREE_TIME_INFINITE_FUTURE),
      IREE_ASYNC_WAIT_FLAG_NONE));
  thread.join();

  iree_hal_semaphore_release(semaphore_a);
  iree_hal_semaphore_release(semaphore_b);
}

// Tests threading behavior by ping-ponging between the test main thread and
// a little thread.
TEST_P(SemaphoreThreadTest, PingPong) {
  iree_hal_semaphore_t* a2b = NULL;
  iree_hal_semaphore_t* b2a = NULL;
  IREE_ASSERT_OK(
      iree_hal_semaphore_create(device_, IREE_HAL_QUEUE_AFFINITY_ANY, 0ull,
                                IREE_HAL_SEMAPHORE_FLAG_DEFAULT, &a2b));
  IREE_ASSERT_OK(
      iree_hal_semaphore_create(device_, IREE_HAL_QUEUE_AFFINITY_ANY, 0ull,
                                IREE_HAL_SEMAPHORE_FLAG_DEFAULT, &b2a));
  std::thread thread([&]() {
    // Should advance right past this because the value is already set.
    IREE_ASSERT_OK(iree_hal_semaphore_wait(
        a2b, 0ull, iree_make_deadline(IREE_TIME_INFINITE_FUTURE),
        IREE_ASYNC_WAIT_FLAG_NONE));
    IREE_ASSERT_OK(iree_hal_semaphore_signal(b2a, 1ull, /*frontier=*/NULL));
    // Jump ahead (blocking at first).
    IREE_ASSERT_OK(iree_hal_semaphore_wait(
        a2b, 4ull, iree_make_deadline(IREE_TIME_INFINITE_FUTURE),
        IREE_ASYNC_WAIT_FLAG_NONE));
  });
  // Block until thread signals.
  IREE_ASSERT_OK(iree_hal_semaphore_wait(
      b2a, 1ull, iree_make_deadline(IREE_TIME_INFINITE_FUTURE),
      IREE_ASYNC_WAIT_FLAG_NONE));
  IREE_ASSERT_OK(iree_hal_semaphore_signal(a2b, 4ull, /*frontier=*/NULL));
  thread.join();

  iree_hal_semaphore_release(a2b);
  iree_hal_semaphore_release(b2a);
}

// Waiting the same value multiple times.
TEST_P(SemaphoreThreadTest, WaitOnTheSameValueMultipleTimes) {
  iree_hal_semaphore_t* semaphore = CreateSemaphore();
  std::thread thread([&]() {
    IREE_ASSERT_OK(iree_hal_semaphore_signal(semaphore, 1, /*frontier=*/NULL));
  });

  IREE_ASSERT_OK(iree_hal_semaphore_wait(
      semaphore, 1, iree_make_deadline(IREE_TIME_INFINITE_FUTURE),
      IREE_ASYNC_WAIT_FLAG_NONE));
  CheckSemaphoreValue(semaphore, 1);

  IREE_ASSERT_OK(iree_hal_semaphore_wait(
      semaphore, 1, iree_make_deadline(IREE_TIME_INFINITE_FUTURE),
      IREE_ASYNC_WAIT_FLAG_NONE));
  CheckSemaphoreValue(semaphore, 1);

  thread.join();

  iree_hal_semaphore_release(semaphore);
}

// Waiting for a finite amount of time.
TEST_P(SemaphoreThreadTest, WaitForFiniteTime) {
  auto generic_test_fn = [this](auto wait_fn) {
    iree_hal_semaphore_t* semaphore = this->CreateSemaphore();

    // Wait before signaling and make sure the semaphore value has not changed.
    IREE_ASSERT_STATUS_IS(IREE_STATUS_DEADLINE_EXCEEDED, wait_fn(semaphore));
    CheckSemaphoreValue(semaphore, 0);

    std::thread signaling_thread([&]() {
      IREE_ASSERT_OK(
          iree_hal_semaphore_signal(semaphore, 1, /*frontier=*/NULL));
    });

    // The semaphore must advance at some point.
    while (true) {
      iree_status_t status = wait_fn(semaphore);
      if (iree_status_is_deadline_exceeded(status)) {
        iree_status_ignore(status);
        continue;
      }
      IREE_ASSERT_OK(status);
      CheckSemaphoreValue(semaphore, 1);
      break;
    }

    signaling_thread.join();
    iree_hal_semaphore_release(semaphore);
  };

  // Immediate timeout.
  generic_test_fn([](iree_hal_semaphore_t* semaphore) {
    return iree_hal_semaphore_wait(semaphore, 1, iree_immediate_timeout(),
                                   IREE_ASYNC_WAIT_FLAG_NONE);
  });

  // Absolute timeout.
  generic_test_fn([](iree_hal_semaphore_t* semaphore) {
    return iree_hal_semaphore_wait(semaphore, 1,
                                   iree_make_deadline(iree_time_now() + 1),
                                   IREE_ASYNC_WAIT_FLAG_NONE);
  });

  // Relative timeout.
  generic_test_fn([](iree_hal_semaphore_t* semaphore) {
    return iree_hal_semaphore_wait(semaphore, 1, iree_make_timeout_ns(1),
                                   IREE_ASYNC_WAIT_FLAG_NONE);
  });
}

// Wait on all semaphores on multiple places simultaneously.
TEST_P(SemaphoreThreadTest, SimultaneousMultiWaitAll) {
  iree_hal_semaphore_t* semaphore1 = this->CreateSemaphore();
  iree_hal_semaphore_t* semaphore2 = this->CreateSemaphore();

  iree_hal_semaphore_t* semaphore_array[] = {semaphore1, semaphore2};
  uint64_t payload_array[] = {1, 1};
  iree_hal_semaphore_list_t semaphore_list = {IREE_ARRAYSIZE(semaphore_array),
                                              semaphore_array, payload_array};

  std::thread wait_thread1([&]() {
    IREE_ASSERT_OK(iree_hal_semaphore_list_wait(
        semaphore_list, iree_make_deadline(IREE_TIME_INFINITE_FUTURE),
        IREE_ASYNC_WAIT_FLAG_NONE));
  });

  std::thread wait_thread2([&]() {
    IREE_ASSERT_OK(iree_hal_semaphore_list_wait(
        semaphore_list, iree_make_deadline(IREE_TIME_INFINITE_FUTURE),
        IREE_ASYNC_WAIT_FLAG_NONE));
  });

  std::thread signal_thread([&]() {
    IREE_ASSERT_OK(
        iree_hal_semaphore_list_signal(semaphore_list, /*frontier=*/NULL));
  });

  wait_thread1.join();
  wait_thread2.join();
  signal_thread.join();

  CheckSemaphoreValue(semaphore1, 1);
  CheckSemaphoreValue(semaphore2, 1);

  iree_hal_semaphore_release(semaphore1);
  iree_hal_semaphore_release(semaphore2);
}

// Wait on a semaphore that is then failed.
TEST_P(SemaphoreThreadTest, WaitThenFail) {
  iree_hal_semaphore_t* semaphore = this->CreateSemaphore();

  iree_status_t status =
      iree_make_status(IREE_STATUS_CANCELLED, "WaitThenFail test.");
  std::thread signal_thread(
      [&]() { iree_hal_semaphore_fail(semaphore, iree_status_clone(status)); });

  iree_status_t wait_status = iree_hal_semaphore_wait(
      semaphore, 1, iree_make_deadline(IREE_TIME_INFINITE_FUTURE),
      IREE_ASYNC_WAIT_FLAG_NONE);
  IREE_EXPECT_STATUS_IS(IREE_STATUS_CANCELLED, wait_status);
  uint64_t value = 1234;
  iree_status_t query_status = iree_hal_semaphore_query(semaphore, &value);
  EXPECT_GE(value, IREE_HAL_SEMAPHORE_FAILURE_VALUE);
  IREE_EXPECT_STATUS_IS(iree_status_code(status), query_status);
  iree_status_ignore(status);

  signal_thread.join();
  iree_hal_semaphore_release(semaphore);
}

// Wait 2 semaphores then fail one of them.
TEST_P(SemaphoreThreadTest, MultiWaitThenFail) {
  iree_hal_semaphore_t* semaphore1 = this->CreateSemaphore();
  iree_hal_semaphore_t* semaphore2 = this->CreateSemaphore();

  iree_status_t status =
      iree_make_status(IREE_STATUS_CANCELLED, "MultiWaitThenFail test.");
  std::thread signal_thread([&]() {
    iree_hal_semaphore_fail(semaphore1, iree_status_clone(status));
  });

  iree_hal_semaphore_t* semaphore_array[] = {semaphore1, semaphore2};
  uint64_t payload_array[] = {1, 1};
  iree_hal_semaphore_list_t semaphore_list = {IREE_ARRAYSIZE(semaphore_array),
                                              semaphore_array, payload_array};
  iree_status_t wait_status = iree_hal_semaphore_list_wait(
      semaphore_list, iree_make_deadline(IREE_TIME_INFINITE_FUTURE),
      IREE_ASYNC_WAIT_FLAG_NONE);
  // multi_wait returns the actual failure code — the caller can follow up with
  // a query to get the full status with message/backtrace if needed.
  IREE_EXPECT_STATUS_IS(IREE_STATUS_CANCELLED, wait_status);
  uint64_t value = 1234;
  iree_status_t semaphore1_query_status =
      iree_hal_semaphore_query(semaphore1, &value);
  EXPECT_GE(value, IREE_HAL_SEMAPHORE_FAILURE_VALUE);
  IREE_EXPECT_STATUS_IS(iree_status_code(status), semaphore1_query_status);
  iree_status_ignore(status);

  // semaphore2 must not have changed.
  uint64_t semaphore2_value = 1234;
  IREE_ASSERT_OK(iree_hal_semaphore_query(semaphore2, &semaphore2_value));
  EXPECT_EQ(semaphore2_value, 0);

  signal_thread.join();
  iree_hal_semaphore_release(semaphore1);
  iree_hal_semaphore_release(semaphore2);
}

// Wait 2 semaphores using iree_hal_device_wait_semaphores then fail one.
TEST_P(SemaphoreThreadTest, DeviceMultiWaitThenFail) {
  iree_hal_semaphore_t* semaphore1 = this->CreateSemaphore();
  iree_hal_semaphore_t* semaphore2 = this->CreateSemaphore();

  iree_status_t status =
      iree_make_status(IREE_STATUS_CANCELLED, "DeviceMultiWaitThenFail test.");
  std::thread signal_thread([&]() {
    iree_hal_semaphore_fail(semaphore1, iree_status_clone(status));
  });

  iree_hal_semaphore_t* semaphore_array[] = {semaphore1, semaphore2};
  uint64_t payload_array[] = {1, 1};
  iree_hal_semaphore_list_t semaphore_list = {IREE_ARRAYSIZE(semaphore_array),
                                              semaphore_array, payload_array};
  iree_status_t wait_status = iree_hal_device_wait_semaphores(
      device_, IREE_ASYNC_WAIT_MODE_ANY, semaphore_list,
      iree_make_deadline(IREE_TIME_INFINITE_FUTURE), IREE_ASYNC_WAIT_FLAG_NONE);
  IREE_EXPECT_STATUS_IS(IREE_STATUS_CANCELLED, wait_status);
  uint64_t value = 1234;
  iree_status_t semaphore1_query_status =
      iree_hal_semaphore_query(semaphore1, &value);
  EXPECT_GE(value, IREE_HAL_SEMAPHORE_FAILURE_VALUE);
  IREE_EXPECT_STATUS_IS(iree_status_code(status), semaphore1_query_status);
  iree_status_ignore(status);

  // semaphore2 must not have changed.
  uint64_t semaphore2_value = 1234;
  IREE_ASSERT_OK(iree_hal_semaphore_query(semaphore2, &semaphore2_value));
  EXPECT_EQ(semaphore2_value, 0);

  signal_thread.join();
  iree_hal_semaphore_release(semaphore1);
  iree_hal_semaphore_release(semaphore2);
}

CTS_REGISTER_TEST_SUITE(SemaphoreThreadTest);

}  // namespace iree::hal::cts
