// Copyright 2024 The Centipede Authors.

//

// Licensed under the Apache License, Version 2.0 (the "License");

// you may not use this file except in compliance with the License.

// You may obtain a copy of the License at

//

//      https://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing, software

// distributed under the License is distributed on an "AS IS" BASIS,

// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

// See the License for the specific language governing permissions and

// limitations under the License.

#include "./common/bazel.h"

#include <algorithm>

#include <cstdlib>

#include <string>

#include "absl/log/check.h"

#include "absl/status/status.h"

#include "absl/strings/numbers.h"

#include "absl/strings/str_cat.h"

#include "absl/time/clock.h"

#include "absl/time/time.h"

namespace centipede {

namespace {

absl::Duration GetBazelTestTimeout() {

  const char *test_timeout_env = std::getenv("TEST_TIMEOUT");

  if (test_timeout_env == nullptr) return absl::InfiniteDuration();

  // When using `bazel run`, `TEST_TIMEOUT` is set but not used, and we should

  // ignore the timeout. We detect this by the presence of

  // `BUILD_WORKSPACE_DIRECTORY`

  // (https://bazel.build/docs/user-manual#running-executables).

  if (std::getenv("BUILD_WORKSPACE_DIRECTORY") != nullptr) {

    return absl::InfiniteDuration();

  }

  int timeout_s = 0;

  CHECK(absl::SimpleAtoi(test_timeout_env, &timeout_s))

      << "Failed to parse TEST_TIMEOUT: \"" << test_timeout_env << "\"";

  return absl::Seconds(timeout_s);

}

}  // namespace

TestShard GetBazelTestShard() {

  static TestShard cached_test_shard = [] {

    TestShard test_shard;

    if (const char *test_total_shards_env = std::getenv("TEST_TOTAL_SHARDS");

        test_total_shards_env != nullptr) {

      CHECK(absl::SimpleAtoi(test_total_shards_env, &test_shard.total_shards))

          << "Failed to parse TEST_TOTAL_SHARDS as an integer: \""

          << test_total_shards_env << "\"";

      CHECK_GT(test_shard.total_shards, 0)

          << "TEST_TOTAL_SHARDS must be greater than 0.";

    }

    if (const char *test_shard_index_env = std::getenv("TEST_SHARD_INDEX");

        test_shard_index_env != nullptr) {

      CHECK(absl::SimpleAtoi(test_shard_index_env, &test_shard.index))

          << "Failed to parse TEST_SHARD_INDEX as an integer: \""

          << test_shard_index_env << "\"";

      CHECK(0 <= test_shard.index && test_shard.index < test_shard.total_shards)

          << "TEST_SHARD_INDEX must be in the range [0, "

          << test_shard.total_shards << ").";

    }

    return test_shard;

  }();

  return cached_test_shard;

}

absl::Status VerifyBazelHasEnoughTimeToRunTest(absl::Time target_start_time,

                                               absl::Duration test_time_limit,

                                               int executed_tests_in_shard,

                                               int fuzz_test_count) {

  static const absl::Duration bazel_test_timeout = GetBazelTestTimeout();

  const int shard_count = GetBazelTestShard().total_shards;

  constexpr float kTimeoutSafetyFactor = 1.2;

  const auto required_test_time = kTimeoutSafetyFactor * test_time_limit;

  const auto remaining_duration =

      bazel_test_timeout - (absl::Now() - target_start_time);

  if (required_test_time <= remaining_duration) return absl::OkStatus();

  std::string error =

      "Cannot fuzz a fuzz test within the given timeout. Please ";

  if (executed_tests_in_shard == 0) {

    // Increasing number of shards won't help.

    const absl::Duration suggested_timeout =

        required_test_time * ((fuzz_test_count - 1) / shard_count + 1);

    absl::StrAppend(&error, "set the `timeout` to ", suggested_timeout,

                    " or reduce the fuzzing time, ");

  } else {

    constexpr int kMaxShardCount = 50;

    const int suggested_shard_count = std::min(

        (fuzz_test_count - 1) / executed_tests_in_shard + 1, kMaxShardCount);

    const int suggested_tests_per_shard =

        (fuzz_test_count - 1) / suggested_shard_count + 1;

    if (suggested_tests_per_shard > executed_tests_in_shard) {

      // We wouldn't be able to execute the suggested number of tests without

      // timeout. This case can only happen if we would in fact need more than

      // `kMaxShardCount` shards, indicating that there are simply too many fuzz

      // tests in a binary.

      CHECK_EQ(suggested_shard_count, kMaxShardCount);

      absl::StrAppend(&error,

                      "split the fuzz tests into several test binaries where "

                      "each binary has at most ",

                      executed_tests_in_shard * kMaxShardCount, "tests ",

                      "with `shard_count` = ", kMaxShardCount, ", ");

    } else {

      // In this case, `suggested_shard_count` must be greater than

      // `shard_count`, otherwise we would have already executed all the tests

      // without a timeout.

      CHECK_GT(suggested_shard_count, shard_count);

      absl::StrAppend(&error, "increase the `shard_count` to ",

                      suggested_shard_count, ", ");

    }

  }

  absl::StrAppend(&error, "to avoid this issue. ");

  absl::StrAppend(&error,

                  "(https://bazel.build/reference/be/"

                  "common-definitions#common-attributes-tests)");

  return absl::ResourceExhaustedError(error);

}

}  // namespace centipede