/*
 * SPDX-FileCopyrightText: Copyright (c) 2021-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * 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
 *
 * http://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.
 */
#ifndef NVIDIA_GXF_STD_TYPE_REGISTRY_HPP
#define NVIDIA_GXF_STD_TYPE_REGISTRY_HPP

#include <map>
#include <set>
#include <shared_mutex>  // NOLINT
#include <string>
#include <string_view>
#include <unordered_map>
#include <utility>

#include "gxf/core/expected.hpp"

namespace nvidia {
namespace gxf {

// A type registry which maps a C++ type to it's name and which can be used to
// track base classes.
class TypeRegistry {
 public:
  Expected<void> add(gxf_tid_t tid, const char* component_type_name);

  Expected<void> add_base(const char* component_type_name, const char* base_type_name);

  Expected<gxf_tid_t> id_from_name(const char* component_type_name) const;

  Expected<bool> is_base(gxf_tid_t derived, gxf_tid_t base) const;

  Expected<const char*> name(gxf_tid_t tid) const;

  // Clears the is_base cache. Call this if the type hierarchy changes after initial setup.
  void clearIsBaseCache();

 private:
  // Internal helper that performs the actual is_base computation without caching.
  // Must be called with mutex_ held in shared mode.
  Expected<bool> is_base_uncached(gxf_tid_t derived, gxf_tid_t base) const;

  // Hash function for (derived, base) TID pairs used in is_base cache
  struct TidPairHash {
    size_t operator()(const std::pair<gxf_tid_t, gxf_tid_t>& p) const noexcept {
      // Combine all four hash values for good distribution
      size_t h1 = p.first.hash1 ^ (p.first.hash2 << 1);
      size_t h2 = p.second.hash1 ^ (p.second.hash2 << 1);
      return h1 ^ (h2 << 1);
    }
  };

  // Equality function for TID pairs
  struct TidPairEqual {
    bool operator()(const std::pair<gxf_tid_t, gxf_tid_t>& a,
                    const std::pair<gxf_tid_t, gxf_tid_t>& b) const noexcept {
      return a.first.hash1 == b.first.hash1 && a.first.hash2 == b.first.hash2 &&
             a.second.hash1 == b.second.hash1 && a.second.hash2 == b.second.hash2;
    }
  };

  // Transparent hash functor for string lookups - allows lookup with const char*
  // without allocating std::string.
  //
  // Note: Heterogeneous lookup for std::unordered_map with is_transparent requires
  // C++20 (P0919R3). With C++20 enabled, find() can accept const char* directly
  // without allocating a temporary std::string.
  struct StringHash {
    using is_transparent = void;  // Enables heterogeneous lookup
    size_t operator()(std::string_view sv) const noexcept {
      return std::hash<std::string_view>{}(sv);
    }
    size_t operator()(const std::string& s) const noexcept {
      return std::hash<std::string_view>{}(s);
    }
    size_t operator()(const char* s) const noexcept {
      return std::hash<std::string_view>{}(s);
    }
  };

  // Transparent equality functor for string lookups
  struct StringEqual {
    using is_transparent = void;  // Enables heterogeneous lookup
    bool operator()(std::string_view a, std::string_view b) const noexcept {
      return a == b;
    }
  };

  std::map<gxf_tid_t, std::string> names_;
  // Using unordered_map with transparent hash for O(1) lookup without string allocation.
  // Heterogeneous lookup allows find() with const char* without constructing std::string.
  std::unordered_map<std::string, gxf_tid_t, StringHash, StringEqual> tids_;
  std::map<gxf_tid_t, std::set<gxf_tid_t>> bases_;
  mutable std::shared_timed_mutex mutex_;

  // Cache for is_base query results: (derived, base) -> is_base_result
  // This cache is populated lazily and provides O(1) lookup for repeated queries.
  mutable std::unordered_map<std::pair<gxf_tid_t, gxf_tid_t>, bool,
                             TidPairHash, TidPairEqual> is_base_cache_;
  mutable std::shared_mutex is_base_cache_mutex_;
};

}  // namespace gxf
}  // namespace nvidia

#endif
