/*
 * Copyright (c) Advanced Micro Devices, Inc. All rights reserved.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

#include "rocm_smi/rocm_smi_dyn_gpu_metrics.h"

#include <cstddef>
#include <cstring>
#include <optional>
#include <shared_mutex>

#include "rocm_smi/rocm_smi.h"
#include "rocm_smi/rocm_smi_logger.h"
#include "rocm_smi/rocm_smi_main.h"
#include "rocm_smi/rocm_smi_utils.h"

namespace amd::smi {

using namespace details;

struct Cursor {
  const std::byte* byte_ptr;
  std::size_t remainder;
};

// Used when mismatch in schema to safely skip value
static inline bool skip_payload(Cursor& cur, AMDGpuMetricAttributeType_t t, uint64_t instances) {
  const std::size_t elem = get_metric_bytes(t);
  if (elem == 0 || instances > std::numeric_limits<size_t>::max() / elem) {
    return false;
  }

  const std::size_t bytes = static_cast<std::size_t>(instances) * elem;
  if (cur.remainder < bytes) {
    return false;
  }

  cur.byte_ptr += bytes;
  cur.remainder -= bytes;
  return true;
}

// Lookup a schema instance for (attr_id, attr_type)
// When schema_type_validate is false, returns schema entry for attr_id regardless of type mismatch
// This allows reading values with actual type from data, not schema type
static inline rsmi_status_t schema_lookup_instance(AMDGpuMetricAttributeId_t attr_id,
                                                   AMDGpuMetricAttributeType_t attr_type,
                                                   AMDGpuMetricAttributeInstance_t& schema_inst,
                                                   bool schema_type_validate = false) {
  if (const auto attr_id_itr = AMDGpuMetricsBaseSchema.find(attr_id);
      attr_id_itr != AMDGpuMetricsBaseSchema.end()) {
    const auto& inst = attr_id_itr->second.m_instance;
    if (!schema_type_validate || inst.m_attribute_type == attr_type) {
      // Log type mismatch when debug is enabled
      if (!schema_type_validate && inst.m_attribute_type != attr_type) {
        std::ostringstream ss;
        ss << __PRETTY_FUNCTION__ << " | Debug: Type mismatch for Attr ID: "
           << static_cast<std::underlying_type_t<AMDGpuMetricAttributeId_t>>(attr_id)
           << " | Schema Type: " << static_cast<int>(inst.m_attribute_type)
           << " | Driver Type: " << static_cast<int>(attr_type);
        LOG_DEBUG(ss);
      }
      // Override schema type with actual type from data
      schema_inst = inst;
      schema_inst.m_attribute_type = attr_type;
      return RSMI_STATUS_SUCCESS;
    }
    return RSMI_STATUS_NOT_SUPPORTED;
  }
  return RSMI_STATUS_NOT_FOUND;
}

static inline std::string attr_type_to_string(AMDGpuMetricAttributeType_t t) {
  switch (t) {
    case AMDGpuMetricAttributeType_t::TYPE_UINT8:
      return "TYPE_UINT8";
    case AMDGpuMetricAttributeType_t::TYPE_INT8:
      return "TYPE_INT8";
    case AMDGpuMetricAttributeType_t::TYPE_UINT16:
      return "TYPE_UINT16";
    case AMDGpuMetricAttributeType_t::TYPE_INT16:
      return "TYPE_INT16";
    case AMDGpuMetricAttributeType_t::TYPE_UINT32:
      return "TYPE_UINT32";
    case AMDGpuMetricAttributeType_t::TYPE_INT32:
      return "TYPE_INT32";
    case AMDGpuMetricAttributeType_t::TYPE_UINT64:
      return "TYPE_UINT64";
    case AMDGpuMetricAttributeType_t::TYPE_INT64:
      return "TYPE_INT64";
    default:
      return "UNKNOWN";
  }
}

template <class T>
static inline std::optional<T> read_scalar(Cursor& c) {
  // Ensure we can read safely
  if (c.remainder < sizeof(T)) {
    return std::nullopt;
  }
  T v{};
  std::memcpy(&v, c.byte_ptr, sizeof(T));
  c.byte_ptr += sizeof(T);
  c.remainder -= sizeof(T);
  return v;
}

template <class T>
static inline std::optional<std::vector<T>> read_vector(Cursor& c, std::size_t count) {
  static_assert(std::is_integral_v<T> && std::is_trivially_copyable_v<T>,
                "metrics expect integral element types");

  // Prevent size_t overflow
  if (count > SIZE_MAX / sizeof(T) || count == 0) {
    return std::nullopt;
  }

  // Ensure we can read entire array safely
  const std::size_t bytes = count * sizeof(T);
  if (c.remainder < bytes) {
    return std::nullopt;
  }

  std::vector<T> out;
  out.resize(count);
  std::memcpy(out.data(), c.byte_ptr, bytes);
  c.byte_ptr += bytes;
  c.remainder -= bytes;
  return out;
}

// Template to fill AMDGpuMetricAttributeValue_t with either a scalar<T> or vector<T>
template <typename T>
static inline std::optional<AMDGpuMetricAttributeValue_t> read_metric_value(Cursor& c,
                                                                            uint64_t instances) {
  if (instances == 1) {
    if (auto v = read_scalar<T>(c)) {
      return AMDGpuMetricAttributeValue_t{*v};
    }
    return std::nullopt;
  }
  if (auto vv = read_vector<T>(c, static_cast<std::size_t>(instances))) {
    return AMDGpuMetricAttributeValue_t{std::move(*vv)};
  }
  return std::nullopt;
}

auto AMDGpuDynamicMetrics_t::parse_from_buffer(const std::byte* data, std::size_t size) noexcept
    -> rsmi_status_t {
  std::ostringstream ss;
  rsmi_status_t status = RSMI_STATUS_SUCCESS;
  if (!data || (size < (sizeof(AMDGpuDynamicMetricsHeader_v1_t) + sizeof(uint32_t)))) {
    return RSMI_STATUS_INSUFFICIENT_SIZE;
  }

  // Grab header
  details::AMDGpuDynamicMetricsHeader_v1_t hdr{};
  std::memcpy(&hdr, data, sizeof(hdr));

  // Advance metrics pointer past header and keep track of remaining file size
  Cursor cur{(data + sizeof(hdr)), (size - sizeof(hdr))};

  // Grab attribute count, directly after header and increment
  auto attr_count_opt = read_scalar<uint32_t>(cur);
  if (!attr_count_opt) {
    return RSMI_STATUS_UNEXPECTED_SIZE;
  }
  uint32_t attr_count = *attr_count_opt;
  if (attr_count == 0 || attr_count > size) {
    return RSMI_STATUS_UNEXPECTED_SIZE;
  }
  std::string m_header_version_str = std::to_string(static_cast<uint32_t>(hdr.m_format_revision)) +
                                     "." +
                                     std::to_string(static_cast<uint32_t>(hdr.m_content_revision));
  ss << __PRETTY_FUNCTION__ << " | Info: Dynamic GPU Metrics"
     << " | Attr Count: " << attr_count << " | Header Version: " << m_header_version_str
     << " | Header Size: " << hdr.get_size() << " | Total Size: " << size << " |";
  LOG_TRACE(ss);

  details::AMDGpuMetricSchemaType_t metrics_data;
  metrics_data.reserve(attr_count);
  AMDGpuDynamicMetricsOffsetMap_t offsets;
  for (uint32_t i = 0; i < attr_count; ++i) {
    if (cur.remainder < sizeof(uint64_t)) {
      return RSMI_STATUS_UNEXPECTED_SIZE;
    }

    // Absolute offset for attribute start in file
    const std::size_t entry_start = static_cast<std::size_t>(cur.byte_ptr - data);

    // Read attribute instance and increment
    auto enc_opt = read_scalar<uint64_t>(cur);
    if (!enc_opt) {
      return RSMI_STATUS_UNEXPECTED_SIZE;
    }
    const uint64_t enc = *enc_opt;

    const auto dec = amdgpu_metrics_decode_attr(enc);

    const auto attr_type = static_cast<AMDGpuMetricAttributeType_t>(dec.m_attr_type);
    const auto attr_id = static_cast<AMDGpuMetricAttributeId_t>(dec.m_attr_id);
    const auto instances = static_cast<uint64_t>(dec.m_attr_instance);

    if (instances == 0) {
      return RSMI_STATUS_UNEXPECTED_SIZE;
    }

    // Schema lookup - use actual type from driver data, not schema type (schema_type_validate =
    // false) This allows different pmfw versions having different data types for same attribute ID
    AMDGpuMetricAttributeInstance_t inst{};
    status = schema_lookup_instance(attr_id, attr_type, inst, false);
    if (status != RSMI_STATUS_SUCCESS) {
      const auto attr_name = [&]() -> std::string {
        const auto it = AMDGpuMetricAttributeIdToString.find(attr_id);
        return (it != AMDGpuMetricAttributeIdToString.end()) ? it->second.m_short_info : "UNKNOWN";
      }();

      ss << __PRETTY_FUNCTION__ << " | Warn: schema lookup miss"
         << " | Attr Name: " << attr_name << " | Attr ID: "
         << static_cast<std::underlying_type_t<AMDGpuMetricAttributeId_t>>(attr_id)
         << " | Attr Type: " << attr_type_to_string(attr_type) << " ("
         << static_cast<std::underlying_type_t<AMDGpuMetricAttributeType_t>>(attr_type) << ")";

      if (status == RSMI_STATUS_NOT_SUPPORTED) {
        const auto expected_type = AMDGpuMetricsBaseSchema.at(attr_id).m_instance.m_attribute_type;
        ss << " | Size mismatch: schema expects " << attr_type_to_string(expected_type) << ", got "
           << attr_type_to_string(attr_type);
      }

      ss << " | Returning = " << getRSMIStatusString(status, false) << " |";
      LOG_TRACE(ss);

      if (!skip_payload(cur, attr_type, instances)) {
        return status;
      }
      continue;  // Do not emit row, go to next attribute
    }

    // Read scalar or all vector values after attribute instance
    AMDGpuMetricAttributeValue_t val{};

    std::optional<AMDGpuMetricAttributeValue_t> mv;
    switch (attr_type) {
      case AMDGpuMetricAttributeType_t::TYPE_UINT8: {
        mv = read_metric_value<std::uint8_t>(cur, instances);
        break;
      }
      case AMDGpuMetricAttributeType_t::TYPE_INT8: {
        mv = read_metric_value<std::int8_t>(cur, instances);
        break;
      }
      case AMDGpuMetricAttributeType_t::TYPE_UINT16: {
        mv = read_metric_value<std::uint16_t>(cur, instances);
        break;
      }
      case AMDGpuMetricAttributeType_t::TYPE_INT16: {
        mv = read_metric_value<std::int16_t>(cur, instances);
        break;
      }
      case AMDGpuMetricAttributeType_t::TYPE_UINT32: {
        mv = read_metric_value<std::uint32_t>(cur, instances);
        break;
      }
      case AMDGpuMetricAttributeType_t::TYPE_INT32: {
        mv = read_metric_value<std::int32_t>(cur, instances);
        break;
      }
      case AMDGpuMetricAttributeType_t::TYPE_UINT64: {
        mv = read_metric_value<std::uint64_t>(cur, instances);
        break;
      }
      case AMDGpuMetricAttributeType_t::TYPE_INT64: {
        mv = read_metric_value<std::int64_t>(cur, instances);
        break;
      }
      default:
        return RSMI_STATUS_INSUFFICIENT_SIZE;
    }

    if (!mv) {
      return RSMI_STATUS_UNEXPECTED_SIZE;
    }

    val = std::move(*mv);  // safely set val
    const uint32_t row_index = static_cast<uint32_t>(metrics_data.size());
    metrics_data.emplace_back(inst, val);
    offsets.try_emplace(entry_start, row_index);
  }

  {
    std::unique_lock<std::shared_mutex> lk(m_mutex);
    m_header = hdr;
    m_attr_count = attr_count;
    m_dynamic_metrics_data.swap(metrics_data);
    m_dynamic_metrics_data_offsets.swap(offsets);
  }
  return RSMI_STATUS_SUCCESS;
}

auto AMDGpuDynamicMetrics_t::parse_from_file(const std::string& metrics_file_path,
                                             std::size_t read_size) -> rsmi_status_t {
  AMDGPUMetricsDynDataBuffer_t buf;

  auto st = read_dynamic_gpu_metrics_file(metrics_file_path, read_size, buf);
  if (st != RSMI_STATUS_SUCCESS) {
    return st;
  }

  return parse_from_buffer(reinterpret_cast<const std::byte*>(buf.data()), buf.size());
}

rsmi_status_t read_dynamic_gpu_metrics_file(const std::string& metrics_file_path,
                                            const size_t read_size,
                                            AMDGPUMetricsDynDataBuffer_t& out) {
  // Clear output buffer and open file stream
  out.clear();
  std::ifstream gpu_metrics_file(metrics_file_path, std::ios::binary);
  if (!gpu_metrics_file.is_open()) {
    return RSMI_STATUS_NOT_FOUND;
  }

  if ((read_size <= 0)) {
    return RSMI_STATUS_UNEXPECTED_SIZE;
  }

  out.resize(read_size);
  gpu_metrics_file.read(reinterpret_cast<char*>(out.data()),
                        static_cast<std::streamsize>(read_size));

  const std::streamsize gpu_metrics_filesize = gpu_metrics_file.gcount();

  if (gpu_metrics_filesize <= 0) {
    out.clear();
    return RSMI_STATUS_NO_DATA;
  }

  out.resize(static_cast<std::size_t>(gpu_metrics_filesize));
  return RSMI_STATUS_SUCCESS;
}

}  // namespace amd::smi
