/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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.
 */
#pragma once
#include <filesystem>
#include <map>
#include <memory>
#include <optional>
#include <set>
#include <string>
#include <unordered_map>
#include <vector>

class SymbolService;
struct irequest;

#include "oi/ContainerInfo.h"
#include "oi/Features.h"
#include "oi/FuncGen.h"
#include "oi/PaddingHunter.h"
#include "oi/TypeHierarchy.h"

extern "C" {
#include <drgn.h>
}

using namespace ObjectIntrospection;
namespace fs = std::filesystem;

struct ParentMember {
  drgn_type* type;
  uint64_t bit_offset;

  bool operator<(const ParentMember& parent) const {
    return (bit_offset < parent.bit_offset);
  }
};

class OICodeGen {
 public:
  struct Config {
    Config() = default;
    Config(const Config& other) = delete;
    Config& operator=(const Config& other) = delete;
    Config(Config&& other) = delete;
    Config& operator=(Config&& other) = delete;

    bool useDataSegment;
    FeatureSet features;
    std::set<fs::path> containerConfigPaths;
    std::set<std::string> defaultHeaders;
    std::set<std::string> defaultNamespaces;
    std::vector<std::pair<std::string, std::string>> membersToStub;
    std::vector<ContainerInfo> passThroughTypes;

    std::string toString() const;
    std::vector<std::string> toOptions() const;
  };

 private:
  // Private constructor. Please use the fallible `OICodeGen::buildFromConfig`
  // for the expected behaviour.
  OICodeGen(const Config&, SymbolService&);

 public:
  static std::unique_ptr<OICodeGen> buildFromConfig(const Config&,
                                                    SymbolService&);
  bool generate(std::string& code);

  [[deprecated("Use generate(std::string&) instead.")]] bool
  generateFunctionsForTypesDrgn(std::string& code) {
    return generate(code);
  }

  bool registerContainer(const fs::path&);

  // TODO: remove me once all the callsites are gone
  static void initializeCodeGen();

  drgn_qualified_type getRootType();
  void setRootType(drgn_qualified_type rt);
  void setLinkageName(std::string name) {
    linkageName = name;
  };
  TypeHierarchy getTypeHierarchy();
  std::map<std::string, PaddingInfo> getPaddingInfo();

  bool isContainer(drgn_type* type);

  bool buildName(drgn_type* type, std::string& text, std::string& outName);

  std::string typeToTransformedName(drgn_type* type);

  bool enumerateTypesRecurse(drgn_type* type);
  static std::string_view drgnKindStr(drgn_type* type);
  std::set<drgn_type*> processedTypes;
  bool isDynamic(drgn_type* type) const;

 private:
  const Config& config;
  FuncGen funcGen;

  using ContainerTypeMapEntry =
      std::pair<std::reference_wrapper<const ContainerInfo>,
                std::vector<drgn_qualified_type>>;

  using TemplateParamList =
      std::vector<std::pair<drgn_qualified_type, std::string>>;

  using SortedTypeDefMap = std::vector<std::pair<drgn_type*, drgn_type*>>;

  std::string rootTypeStr;
  std::string linkageName;
  std::map<drgn_type*, std::string> unnamedUnion;
  std::map<std::string, size_t> sizeMap;
  std::map<drgn_type*, ContainerTypeMapEntry> containerTypeMapDrgn;
  std::vector<std::unique_ptr<ContainerInfo>> containerInfoList;
  std::vector<drgn_type*> enumTypes;
  std::vector<std::string> knownTypes;
  drgn_qualified_type rootType;
  drgn_qualified_type rootTypeToIntrospect;

  std::map<std::string, std::string> typedefMap;
  std::map<drgn_type*, std::vector<ParentMember>> parentClasses;
  std::map<std::string, std::vector<drgn_type*>> childClasses;
  std::map<drgn_type*, std::vector<drgn_type*>> descendantClasses;

  SymbolService& symbols;

  size_t pad_index = 0;
  std::unordered_map<drgn_type*, std::pair<size_t, size_t>> paddingIndexMap;

  std::map<drgn_type*, drgn_type*> typedefTypes;
  std::map<drgn_type*, std::vector<DrgnClassMemberInfo>> classMembersMap;
  std::map<drgn_type*, std::vector<DrgnClassMemberInfo>> classMembersMapCopy;
  std::map<drgn_type*, std::string> typeToNameMap;
  std::map<std::string, drgn_type*> nameToTypeMap;
  std::set<drgn_type*> funcDefTypeList;
  std::vector<drgn_type*> structDefType;
  std::set<drgn_type*> knownDummyTypeList;
  std::map<drgn_type*, drgn_type*> pointerToTypeMap;
  std::set<drgn_type*> thriftIssetStructTypes;
  std::vector<drgn_type*> topoSortedStructTypes;
  std::vector<std::pair<std::string, std::string>> membersToStub;

  ContainerInfoRefSet containerTypesFuncDef;

  std::map<std::string, PaddingInfo> paddedStructs;

  std::map<drgn_type*, std::vector<DrgnClassMemberInfo>>& getClassMembersMap();

  class DrgnString {
    struct FreeDeleter {
      void operator()(void* allocated) {
        free(allocated);
      }
    };

   public:
    std::string_view contents;
    DrgnString(char* data, size_t length)
        : contents{data, length}, _data{data} {
    }
    DrgnString() = delete;

   private:
    std::unique_ptr<char, FreeDeleter> _data;
  };

  bool feature(Feature f) const {
    return config.features[f];
  }
  static void prependQualifiers(enum drgn_qualifiers, std::string& sb);
  static std::string stripFullyQualifiedName(
      const std::string& fullyQualifiedName);
  std::string stripFullyQualifiedNameWithSeparators(
      const std::string& fullyQualifiedname);
  static void removeTemplateParamAtIndex(std::vector<std::string>& params,
                                         const size_t index);
  std::unordered_map<drgn_type*, DrgnString> fullyQualifiedNames;
  std::optional<const std::string_view> fullyQualifiedName(drgn_type* type);
  static SortedTypeDefMap getSortedTypeDefMap(
      const std::map<drgn_type*, drgn_type*>& typedefTypeMap);

  std::optional<std::reference_wrapper<const ContainerInfo>> getContainerInfo(
      drgn_type* type);
  void printAllTypes();
  void printAllTypeNames();

  static void addPaddingForBaseClass(drgn_type* type,
                                     std::vector<std::string>& def);
  void addTypeToName(drgn_type* type, std::string name);
  bool generateNamesForTypes();
  bool generateJitCode(std::string& code);
  bool generateStructDefs(std::string& code);
  bool generateStructDef(drgn_type* e, std::string& code);
  bool getDrgnTypeName(drgn_type* type, std::string& outName);

  bool getDrgnTypeNameInt(drgn_type* type, std::string& outName);
  bool recordChildren(drgn_type* type);
  bool enumerateChildClasses();
  bool populateDefsAndDecls();
  static void memberTransformName(
      std::map<std::string, std::string>& templateTransformMap,
      std::string& typeName);

  bool getMemberDefinition(drgn_type* type);
  bool isKnownType(const std::string& type);
  bool isKnownType(const std::string& type, std::string& matched);

  static bool getTemplateParams(
      drgn_type* type,
      size_t numTemplateParams,
      std::vector<std::pair<drgn_qualified_type, std::string>>& v);

  bool enumerateTemplateParamIdxs(drgn_type* type,
                                  const ContainerInfo& containerInfo,
                                  const std::vector<size_t>& paramIdxs,
                                  bool& ifStub);
  bool getContainerTemplateParams(drgn_type* type, bool& ifStub);
  void enumerateDescendants(drgn_type* type, drgn_type* baseType);
  void getFuncDefinitionStr(std::string& code,
                            drgn_type* type,
                            const std::string& typeName);
  std::optional<uint64_t> getDrgnTypeSize(drgn_type* type);

  std::optional<std::string> getNameForType(drgn_type* type);

  static std::string preProcessUniquePtr(drgn_type* type, std::string name);
  std::string transformTypeName(drgn_type* type, std::string& text);
  static std::string templateTransformType(const std::string& typeName);
  static std::string structNameTransformType(const std::string& typeName);

  bool addPadding(uint64_t padding_bits, std::string& code);
  static void deduplicateMemberName(
      std::unordered_map<std::string, int>& memberNames,
      std::string& memberName);
  std::optional<uint64_t> generateMember(
      const DrgnClassMemberInfo& m,
      std::unordered_map<std::string, int>& memberNames,
      uint64_t currOffsetBits,
      std::string& code,
      bool isInUnion);
  bool generateParent(drgn_type* p,
                      std::unordered_map<std::string, int>& memberNames,
                      uint64_t& currOffsetBits,
                      std::string& code,
                      size_t offsetToNextMember);
  std::optional<uint64_t> getAlignmentRequirements(drgn_type* e);
  bool generateStructMembers(drgn_type* e,
                             std::unordered_map<std::string, int>& memberNames,
                             std::string& code,
                             uint64_t& out_offset_bits,
                             PaddingInfo& paddingInfo,
                             bool& violatesAlignmentRequirement,
                             size_t offsetToNextMember);
  void getFuncDefClassMembers(std::string& code,
                              drgn_type* type,
                              std::unordered_map<std::string, int>& memberNames,
                              bool skipPadding = false);
  bool isDrgnSizeComplete(drgn_type* type);

  static bool getEnumUnderlyingTypeStr(drgn_type* e,
                                       std::string& enumUnderlyingTypeStr);

  bool ifEnumerateClass(const std::string& typeName);

  bool enumerateClassParents(drgn_type* type, const std::string& typeName);
  bool enumerateClassMembers(drgn_type* type,
                             const std::string& typeName,
                             bool& isStubbed);
  bool enumerateClassTemplateParams(drgn_type* type,
                                    const std::string& typeName,
                                    bool& isStubbed);
  bool ifGenerateMemberDefinition(const std::string& typeName);
  bool generateMemberDefinition(drgn_type* type, std::string& typeName);
  std::optional<std::pair<std::string_view, std::string_view>> isMemberToStub(
      const std::string& type, const std::string& member);
  std::optional<std::string_view> isTypeToStub(const std::string& typeName);
  bool isTypeToStub(drgn_type* type, const std::string& typeName);
  bool isEmptyClassOrFunctionType(drgn_type* type, const std::string& typeName);
  bool enumerateClassType(drgn_type* type);
  bool enumerateTypeDefType(drgn_type* type);
  bool enumerateEnumType(drgn_type* type);
  bool enumeratePointerType(drgn_type* type);
  bool enumeratePrimitiveType(drgn_type* type);
  bool enumerateArrayType(drgn_type* type);

  bool isUnnamedStruct(drgn_type* type);

  std::string getAnonName(drgn_type*, const char*);
  std::string getStructName(drgn_type* type) {
    return getAnonName(type, "__anon_struct_");
  }
  std::string getUnionName(drgn_type* type) {
    return getAnonName(type, "__anon_union_");
  }
  static void declareThriftStruct(std::string& code, std::string_view name);

  bool isNumMemberGreaterThanZero(drgn_type* type);
  void getClassMembersIncludingParent(drgn_type* type,
                                      std::vector<DrgnClassMemberInfo>& out);
  bool staticAssertMemberOffsets(
      const std::string& struct_name,
      drgn_type* struct_type,
      std::string& assert_str,
      std::unordered_map<std::string, int>& member_names,
      uint64_t base_offset = 0);
  bool addStaticAssertsForType(drgn_type* type,
                               bool generateAssertsForOffsets,
                               std::string& code);
  bool buildNameInt(drgn_type* type,
                    std::string& nameWithoutTemplate,
                    std::string& outName);
  void replaceTemplateOperator(
      std::vector<std::pair<drgn_qualified_type, std::string>>& template_params,
      std::vector<std::string>& template_params_strings,
      size_t index);
  void replaceTemplateParameters(
      drgn_type* type,
      std::vector<std::pair<drgn_qualified_type, std::string>>& template_params,
      std::vector<std::string>& template_params_strings,
      const std::string& nameWithoutTemplate);
};