/***************************************************************************************************
 * Copyright (c) 2023 - 2025 Hygon Information Technology Co., Ltd. All rights reserved.
 * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: BSD-3-Clause
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice, this
 * list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 * this list of conditions and the following disclaimer in the documentation
 * and/or other materials provided with the distribution.
 *
 * 3. Neither the name of the copyright holder nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 **************************************************************************************************/

#include <hip/hip_runtime_api.h>

#include "hytlass_unit_test.h"

/////////////////////////////////////////////////////////////////////////////////////////////////

/// Gets a HIP device
hipDeviceProp_t GetHipDevice() {

  hipError_t err;

  int hipDeviceId;
  err = hipGetDevice(&hipDeviceId);
  if (hipSuccess != err) {
    std::cerr << "*** Error: Could not detect active GPU device ID"
              << " [" << hipGetErrorString(err) << "]" << std::endl;
    exit(1);
  }

  hipDeviceProp_t deviceProperties;
  err = hipGetDeviceProperties(&deviceProperties, hipDeviceId);

  return deviceProperties;
}

/// Prints device properties
std::ostream &operator<<(std::ostream &out, hipDeviceProp_t const &deviceProperties) {

  int deviceMajorMinor = deviceProperties.major * 10 + deviceProperties.minor;
  if (deviceMajorMinor) {
    int32_t clock_MHz = deviceProperties.clockRate / 1000;
    out << "GPU(compute_"
      << deviceMajorMinor << ", "
      << deviceProperties.multiProcessorCount << " SMs @ " << clock_MHz << " MHz)";
  }
  else {
    out << "No HIP device.";
  }

  return out;
}
/////////////////////////////////////////////////////////////////////////////////////////////////
int GetArchCode() {
  auto deviceProps = GetHipDevice();
  std::string name(deviceProps.gcnArchName);
  if (name.rfind("gfx", 0) == 0) {
    name = name.substr(3);
  }
  std::string archCode;
  for (char c : name) {
    if (isdigit(c)) archCode.push_back(c);
    else break;
  }

  return archCode.empty() ? -1 : std::stoi(archCode);
}
/////////////////////////////////////////////////////////////////////////////////////////////////

/// Sets flags for Unit test
void FilterArchitecture() {
  // Default flags can be overwritten by --gtest_filter from commandline

  int const kMaxDevice = 999;

  hipError_t err;

  int hipDeviceId;
  err = hipGetDevice(&hipDeviceId);
  if (hipSuccess != err) {
    std::cerr << "*** Error: Could not detect active GPU device ID"
              << " [" << hipGetErrorString(err) << "]" << std::endl;
    exit(1);
  }

  hipDeviceProp_t deviceProperties;
  err = hipGetDeviceProperties(&deviceProperties, hipDeviceId);
  if (hipSuccess != err) {
    std::cerr << "*** Error: Could not get device properties for GPU " << hipDeviceId << " ["
              << hipGetErrorString(err) << "]" << std::endl;
    exit(1);
  }

  int deviceMajorMinor = GetArchCode();

  struct {

    /// Unit test filter string
    char const *filter;

    /// Minimum compute capability for the kernels in the named test
    int min_compute_capability;

    /// Maximum compute capability for which the kernels are enabled 
    int max_compute_capability;
  } 
  test_filters[] = {
    { "GFX906*",                      906, kMaxDevice},
    { "GFX928*",                      928, kMaxDevice},
    { "GFX936*",                      936, kMaxDevice},
    { 0, 0, false }
  };


  // Set negative test filters
  std::stringstream ss;
  ss << "-";
  for (int i = 0, j = 0; test_filters[i].filter; ++i) {

    if (deviceMajorMinor < test_filters[i].min_compute_capability ||
        deviceMajorMinor > test_filters[i].max_compute_capability) {

      ss << (j++ ? ":" : "") << test_filters[i].filter;
    }
  }

  ::testing::GTEST_FLAG(filter) = ss.str();
}

/////////////////////////////////////////////////////////////////////////////////////////////////

int HytlassUnitTestProblemCount() {
    if(const char* problem_count = std::getenv("HYTLASS_UNIT_TEST_PROBLEM_COUNT")) {

        return std::stoi(problem_count);
    } 

    return 0;
}

/////////////////////////////////////////////////////////////////////////////////////////////////
