//
// This file is distributed under the MIT License. See LICENSE.md for details.
//

#include "llvm/ADT/SmallSet.h"
#include "llvm/Support/FormatVariadic.h"

#include "revng/ADT/RecursiveCoroutine.h"
#include "revng/Clift/CliftAttributes.h"
#include "revng/Clift/CliftDialect.h"
#include "revng/Clift/CliftTypes.h"
#include "revng/CliftImportModel/ImportModel.h"
#include "revng/Pipeline/Location.h"
#include "revng/Pipes/Ranks.h"

namespace {

namespace clift = mlir::clift;

template<typename Attribute>
using AttributeVector = llvm::SmallVector<Attribute, 16>;

class CliftConverter {
  mlir::MLIRContext *Context;
  llvm::function_ref<mlir::InFlightDiagnostic()> EmitError;

  llvm::DenseMap<uint64_t, clift::DefinedType> Cache;
  llvm::DenseMap<uint64_t, const model::TypeDefinition *> IncompleteTypes;

  llvm::SmallSet<uint64_t, 16> DefinitionGuardSet;

  class RecursiveDefinitionGuard {
    CliftConverter *Self = nullptr;
    uint64_t ID;

  public:
    explicit RecursiveDefinitionGuard(CliftConverter &Self, const uint64_t ID) {
      if (Self.DefinitionGuardSet.insert(ID).second) {
        this->Self = &Self;
        this->ID = ID;
      }
    }

    RecursiveDefinitionGuard(const RecursiveDefinitionGuard &) = delete;
    RecursiveDefinitionGuard &
    operator=(const RecursiveDefinitionGuard &) = delete;

    ~RecursiveDefinitionGuard() {
      if (Self != nullptr) {
        size_t const Erased = Self->DefinitionGuardSet.erase(ID);
        revng_assert(Erased == 1);
      }
    }

    explicit operator bool() const { return Self != nullptr; }
  };

public:
  explicit CliftConverter(mlir::MLIRContext &Context,
                          const model::Binary &Binary,
                          llvm::function_ref<mlir::InFlightDiagnostic()>
                            EmitError) :
    Context(&Context), EmitError(EmitError) {}

  CliftConverter(const CliftConverter &) = delete;
  CliftConverter &operator=(const CliftConverter &) = delete;

  ~CliftConverter() { revng_assert(DefinitionGuardSet.empty()); }

  clift::ValueType
  convertTypeDefinition(const model::TypeDefinition &ModelType) {
    const clift::ValueType T = fromTypeDefinition(ModelType,
                                                  /* RequireComplete = */ true);
    if (T and not processIncompleteTypes())
      return nullptr;
    return T;
  }

  clift::ValueType convertType(const model::Type &ModelType) {
    const clift::ValueType T = fromType(ModelType,
                                        /* RequireComplete = */ true);
    if (T and not processIncompleteTypes())
      return nullptr;
    return T;
  }

private:
  template<typename T, typename... ArgTypes>
  T make(const ArgTypes &...Args) {
    return T::getChecked(EmitError, Context, Args...);
  }

  template<typename KeyT>
  clift::MutableStringAttr
  makeNameAttr(llvm::StringRef Handle, llvm::StringRef Name = {}) {
    return clift::makeNameAttr<KeyT>(Context, Handle, Name);
  }

  static clift::PrimitiveKind
  getPrimitiveKind(const model::PrimitiveType &ModelType) {
    switch (ModelType.PrimitiveKind()) {
    case model::PrimitiveKind::Void:
      return clift::PrimitiveKind::VoidKind;
    case model::PrimitiveKind::Generic:
      return clift::PrimitiveKind::GenericKind;
    case model::PrimitiveKind::PointerOrNumber:
      return clift::PrimitiveKind::PointerOrNumberKind;
    case model::PrimitiveKind::Number:
      return clift::PrimitiveKind::NumberKind;
    case model::PrimitiveKind::Unsigned:
      return clift::PrimitiveKind::UnsignedKind;
    case model::PrimitiveKind::Signed:
      return clift::PrimitiveKind::SignedKind;
    case model::PrimitiveKind::Float:
      return clift::PrimitiveKind::FloatKind;

    case model::PrimitiveKind::Invalid:
    case model::PrimitiveKind::Count:
      revng_abort("These are invalid values. Something has gone wrong.");
    }
  }

  auto getLocation(const model::TypeDefinition &T) {
    return pipeline::location(revng::ranks::TypeDefinition, T.key());
  }

  std::string getHandle(const model::TypeDefinition &T) {
    return getLocation(T).toString();
  }

  std::string getRegisterSetHandle(const model::RawFunctionDefinition &T) {
    return pipeline::locationString(revng::ranks::ArtificialStruct, T.key());
  }

  RecursiveCoroutine<clift::DefinedType>
  getTypeDefinition(const model::CABIFunctionDefinition &ModelType) {
    RecursiveDefinitionGuard Guard(*this, ModelType.ID());
    if (not Guard) {
      if (EmitError)
        EmitError() << "Recursive definition of CABIFunctionDefinition "
                    << ModelType.ID();
      rc_return nullptr;
    }

    AttributeVector<mlir::Type> ArgumentTypes;
    ArgumentTypes.reserve(ModelType.Arguments().size());

    for (const model::Argument &Argument : ModelType.Arguments()) {
      const auto Type = rc_recur fromType(*Argument.Type());
      if (not Type)
        rc_return nullptr;
      ArgumentTypes.push_back(Type);
    }

    mlir::Type ReturnType = nullptr;
    if (ModelType.ReturnType().isEmpty())
      ReturnType = rc_recur fromType(*model::PrimitiveType::makeVoid());
    else
      ReturnType = rc_recur fromType(*ModelType.ReturnType());

    if (not ReturnType)
      rc_return nullptr;

    auto Handle = getHandle(ModelType);
    auto NameAttr = makeNameAttr<clift::FunctionType>(Handle);
    rc_return make<clift::FunctionType>(llvm::StringRef(Handle),
                                        NameAttr,
                                        ReturnType,
                                        llvm::ArrayRef(ArgumentTypes));
  }

  RecursiveCoroutine<clift::DefinedType>
  getTypeDefinition(const model::EnumDefinition &ModelType) {
    RecursiveDefinitionGuard Guard(*this, ModelType.ID());
    if (not Guard) {
      if (EmitError)
        EmitError() << "Recursive definition of EnumDefinition "
                    << ModelType.ID();
      rc_return nullptr;
    }

    const auto UnderlyingType = rc_recur fromType(*ModelType.UnderlyingType());
    if (not UnderlyingType)
      rc_return nullptr;

    auto Location = getLocation(ModelType);

    AttributeVector<clift::EnumFieldAttr> Fields;
    Fields.reserve(ModelType.Entries().size());

    for (const model::EnumEntry &Entry : ModelType.Entries()) {
      auto Handle = Location.extend(revng::ranks::EnumEntry, Entry.Value())
                      .toString();

      auto NameAttr = makeNameAttr<clift::EnumFieldAttr>(Handle);
      auto Attr = make<clift::EnumFieldAttr>(llvm::StringRef(Handle),
                                             NameAttr,
                                             Entry.Value());

      if (not Attr)
        rc_return nullptr;

      Fields.push_back(Attr);
    }

    auto Handle = Location.toString();
    auto NameAttr = makeNameAttr<clift::EnumAttr>(Handle);
    auto Attr = make<clift::EnumAttr>(llvm::StringRef(Handle),
                                      NameAttr,
                                      UnderlyingType,
                                      llvm::ArrayRef(Fields));

    if (not Attr)
      rc_return nullptr;

    rc_return clift::EnumType::get(Attr);
  }

  RecursiveCoroutine<clift::ValueType>
  getRegisterSetType(const model::RawFunctionDefinition &ModelType) {
    auto Location = getLocation(ModelType);

    AttributeVector<clift::FieldAttr> Fields;
    Fields.reserve(ModelType.ReturnValues().size());

    uint64_t Offset = 0;
    for (const model::NamedTypedRegister &Register : ModelType.ReturnValues()) {
      const auto RegisterType = rc_recur fromType(*Register.Type());
      if (not RegisterType)
        rc_return nullptr;

      auto Handle = Location
                      .extend(revng::ranks::ReturnRegister, Register.Location())
                      .toString();

      auto NameAttr = makeNameAttr<clift::FieldAttr>(Handle);
      auto Attr = make<clift::FieldAttr>(llvm::StringRef(Handle),
                                         NameAttr,
                                         Offset,
                                         RegisterType);
      if (not Attr)
        rc_return nullptr;

      Fields.push_back(Attr);
      Offset += RegisterType.getByteSize();
    }

    auto Handle = Location.transmute(revng::ranks::ArtificialStruct).toString();
    auto NameAttr = makeNameAttr<clift::StructAttr>(Handle);
    auto Attr = make<clift::StructAttr>(llvm::StringRef(Handle),
                                        NameAttr,
                                        Offset,
                                        llvm::ArrayRef(Fields));

    if (not Attr)
      rc_return nullptr;

    rc_return clift::StructType::get(Attr);
  }

  RecursiveCoroutine<clift::DefinedType>
  getTypeDefinition(const model::RawFunctionDefinition &ModelType) {
    RecursiveDefinitionGuard Guard(*this, ModelType.ID());
    if (not Guard) {
      if (EmitError)
        EmitError() << "Recursive definition of RawFunctionDefinition "
                    << ModelType.ID();
      rc_return nullptr;
    }

    mlir::Type StackArgumentType;
    size_t ArgumentsCount = 0;

    if (not ModelType.StackArgumentsType().isEmpty()) {
      StackArgumentType = rc_recur fromType(*ModelType.StackArgumentsType());
      if (not StackArgumentType)
        rc_return nullptr;
      ++ArgumentsCount;
    }

    ArgumentsCount += ModelType.Arguments().size();
    AttributeVector<mlir::Type> ArgumentTypes;
    ArgumentTypes.reserve(ArgumentsCount);

    for (const model::NamedTypedRegister &Register : ModelType.Arguments()) {
      const auto Type = rc_recur fromType(*Register.Type());
      if (not Type)
        rc_return nullptr;
      ArgumentTypes.push_back(Type);
    }

    if (StackArgumentType)
      ArgumentTypes.push_back(StackArgumentType);

    clift::ValueType ReturnType;
    switch (ModelType.ReturnValues().size()) {
    case 0:
      ReturnType = make<clift::PrimitiveType>(clift::PrimitiveKind::VoidKind,
                                              /*Size=*/static_cast<uint64_t>(0),
                                              /*IsConst=*/false);
      break;

    case 1:
      ReturnType = rc_recur fromType(*ModelType.ReturnValues().begin()->Type());
      break;

    default: {
      ReturnType = clift::StructType::get(Context,
                                          getRegisterSetHandle(ModelType));

      const auto R = IncompleteTypes.try_emplace(ModelType.ID(), &ModelType);
      revng_assert(R.second && "Register set types are only visited once.");
    } break;
    }
    if (not ReturnType)
      rc_return nullptr;

    auto Handle = getHandle(ModelType);
    auto NameAttr = makeNameAttr<clift::FunctionType>(Handle);
    rc_return make<clift::FunctionType>(llvm::StringRef(Handle),
                                        NameAttr,
                                        mlir::Type(ReturnType),
                                        llvm::ArrayRef(ArgumentTypes));
  }

  RecursiveCoroutine<clift::DefinedType>
  getTypeDefinition(const model::StructDefinition &ModelType,
                    const bool RequireComplete) {
    if (not RequireComplete) {
      const auto T = clift::StructType::get(Context, getHandle(ModelType));
      if (not T.isComplete())
        IncompleteTypes.try_emplace(ModelType.ID(), &ModelType);
      rc_return T;
    }

    RecursiveDefinitionGuard Guard(*this, ModelType.ID());
    if (not Guard) {
      if (EmitError)
        EmitError() << "Recursive definition of StructTypeAttr "
                    << ModelType.ID();
      rc_return nullptr;
    }

    auto Location = getLocation(ModelType);

    AttributeVector<clift::FieldAttr> Fields;
    Fields.reserve(ModelType.Fields().size());

    for (const model::StructField &Field : ModelType.Fields()) {
      const auto FieldType = rc_recur fromType(*Field.Type(),
                                               /* RequireComplete = */ true);
      if (not FieldType)
        rc_return nullptr;

      auto Handle = Location.extend(revng::ranks::StructField, Field.Offset())
                      .toString();

      auto NameAttr = makeNameAttr<clift::FieldAttr>(Handle);
      auto Attr = make<clift::FieldAttr>(llvm::StringRef(Handle),
                                         NameAttr,
                                         Field.Offset(),
                                         FieldType);
      if (not Attr)
        rc_return nullptr;

      Fields.push_back(Attr);
    }

    auto Handle = Location.toString();
    auto NameAttr = makeNameAttr<clift::StructAttr>(Handle);
    auto Attr = make<clift::StructAttr>(llvm::StringRef(Handle),
                                        NameAttr,
                                        ModelType.Size(),
                                        llvm::ArrayRef(Fields));

    if (not Attr)
      rc_return nullptr;

    rc_return clift::StructType::get(Attr);
  }

  RecursiveCoroutine<clift::DefinedType>
  getTypeDefinition(const model::TypedefDefinition &ModelType,
                    const bool RequireComplete) {
    std::optional<RecursiveDefinitionGuard> Guard;

    if (RequireComplete) {
      Guard.emplace(*this, ModelType.ID());
      if (not *Guard) {
        if (EmitError)
          EmitError() << "Recursive definition of TypedefDefinition "
                      << ModelType.ID();
        rc_return nullptr;
      }
    }

    const auto UnderlyingType = rc_recur fromType(*ModelType.UnderlyingType(),
                                                  RequireComplete);
    if (not UnderlyingType)
      rc_return nullptr;

    auto Handle = getHandle(ModelType);
    auto NameAttr = makeNameAttr<clift::TypedefAttr>(Handle);
    auto Attr = make<clift::TypedefAttr>(llvm::StringRef(Handle),
                                         NameAttr,
                                         UnderlyingType);

    if (not Attr)
      rc_return nullptr;

    rc_return clift::TypedefType::get(Attr);
  }

  RecursiveCoroutine<clift::DefinedType>
  getTypeDefinition(const model::UnionDefinition &ModelType,
                    const bool RequireComplete) {
    if (not RequireComplete) {
      const auto T = clift::UnionType::get(Context, getHandle(ModelType));
      if (not T.isComplete())
        IncompleteTypes.try_emplace(ModelType.ID(), &ModelType);
      rc_return T;
    }

    RecursiveDefinitionGuard Guard(*this, ModelType.ID());
    if (not Guard) {
      if (EmitError)
        EmitError() << "Recursive definition of UnionTypeAttr "
                    << ModelType.ID();
      rc_return nullptr;
    }

    auto Location = getLocation(ModelType);

    AttributeVector<clift::FieldAttr> Fields;
    Fields.reserve(ModelType.Fields().size());

    for (const model::UnionField &Field : ModelType.Fields()) {
      const auto FieldType = rc_recur fromType(*Field.Type(),
                                               /* RequireComplete = */ true);
      if (not FieldType)
        rc_return nullptr;

      auto Handle = Location.extend(revng::ranks::UnionField, Field.Index())
                      .toString();

      auto NameAttr = makeNameAttr<clift::FieldAttr>(Handle);
      auto Attr = make<clift::FieldAttr>(llvm::StringRef(Handle),
                                         NameAttr,
                                         /*Offset=*/static_cast<uint64_t>(0),
                                         FieldType);
      if (not Attr)
        rc_return nullptr;
      Fields.push_back(Attr);
    }

    auto Handle = Location.toString();
    auto NameAttr = makeNameAttr<clift::UnionAttr>(Handle);
    auto Attr = make<clift::UnionAttr>(llvm::StringRef(Handle),
                                       NameAttr,
                                       llvm::ArrayRef(Fields));

    rc_return clift::UnionType::get(Attr);
  }

  RecursiveCoroutine<clift::DefinedType>
  getTypeDefinition(const model::TypeDefinition &T, bool &RequireComplete) {
    if (const auto *CFT = llvm::dyn_cast<model::CABIFunctionDefinition>(&T))
      rc_return getTypeDefinition(*CFT);

    if (const auto *RFT = llvm::dyn_cast<model::RawFunctionDefinition>(&T))
      rc_return getTypeDefinition(*RFT);

    if (const auto *Enum = llvm::dyn_cast<model::EnumDefinition>(&T))
      rc_return getTypeDefinition(*Enum);

    if (const auto *Struct = llvm::dyn_cast<model::StructDefinition>(&T))
      rc_return getTypeDefinition(*Struct, RequireComplete);

    if (const auto *Union = llvm::dyn_cast<model::UnionDefinition>(&T))
      rc_return getTypeDefinition(*Union, RequireComplete);

    if (const auto *Typedef = llvm::dyn_cast<model::TypedefDefinition>(&T))
      rc_return getTypeDefinition(*Typedef, RequireComplete);

    revng_abort("Unsupported type definition kind.");
  }

  RecursiveCoroutine<clift::ValueType>
  fromTypeDefinition(const model::TypeDefinition &ModelType,
                     bool RequireComplete = false,
                     const bool Const = false) {
    if (Const) {
      auto Type = rc_recur fromTypeDefinition(ModelType,
                                              RequireComplete,
                                              /*Const=*/false);
      rc_return Type.addConst();
    }

    if (const auto It = Cache.find(ModelType.ID()); It != Cache.end())
      rc_return It->second;

    if (not ModelType.verify()) {
      if (EmitError)
        EmitError() << "Invalid model type definition";

      rc_return nullptr;
    }

    auto Type = rc_recur getTypeDefinition(ModelType, RequireComplete);

    if (Type and RequireComplete) {
      auto [Iterator, Inserted] = Cache.try_emplace(ModelType.ID(), Type);
      revng_assert(Inserted);
    }

    rc_return Type;
  }

  RecursiveCoroutine<clift::ValueType> fromType(const model::Type &ModelType,
                                                bool RequireComplete = false) {
    if (not ModelType.verify()) {
      if (EmitError)
        EmitError() << "Invalid model type";

      rc_return nullptr;
    }

    if (const auto &P = llvm::dyn_cast<model::PrimitiveType>(&ModelType)) {
      rc_return make<clift::PrimitiveType>(getPrimitiveKind(*P),
                                           P->Size(),
                                           P->IsConst());

    } else if (const auto &D = llvm::dyn_cast<model::DefinedType>(&ModelType)) {
      rc_return fromTypeDefinition(D->unwrap(), RequireComplete, D->IsConst());

    } else if (const auto &A = llvm::dyn_cast<model::ArrayType>(&ModelType)) {
      rc_return make<clift::ArrayType>(rc_recur fromType(*A->ElementType(),
                                                         RequireComplete),
                                       A->ElementCount());

    } else if (const auto &P = llvm::dyn_cast<model::PointerType>(&ModelType)) {
      // If there's a pointer in the way, the base type does not have to be
      // complete.
      RequireComplete = false;

      rc_return make<clift::PointerType>(rc_recur fromType(*P->PointeeType(),
                                                           RequireComplete),
                                         P->PointerSize(),
                                         P->IsConst());

    } else {
      if (EmitError)
        EmitError() << "Unknown model type";

      rc_return nullptr;
    }
  }

  bool processIncompleteTypes() {
    while (not IncompleteTypes.empty()) {
      const auto Iterator = IncompleteTypes.begin();
      const model::TypeDefinition &ModelType = *Iterator->second;
      IncompleteTypes.erase(Iterator);

      clift::ValueType CompleteType;
      if (auto RFT = llvm::dyn_cast<model::RawFunctionDefinition>(&ModelType)) {
        CompleteType = getRegisterSetType(*RFT);
      } else {
        CompleteType = fromTypeDefinition(ModelType, /*RequireComplete=*/true);
      }

      if (not CompleteType)
        return false;
    }

    return true;
  }
};

} // namespace

clift::ValueType
clift::importModelType(llvm::function_ref<mlir::InFlightDiagnostic()> EmitError,
                       mlir::MLIRContext &Context,
                       const model::TypeDefinition &ModelType,
                       const model::Binary &Binary) {
  return CliftConverter(Context, Binary, EmitError)
    .convertTypeDefinition(ModelType);
}

clift::ValueType
clift::importModelType(llvm::function_ref<mlir::InFlightDiagnostic()> EmitError,
                       mlir::MLIRContext &Context,
                       const model::Type &ModelType,
                       const model::Binary &Binary) {
  return CliftConverter(Context, Binary, EmitError).convertType(ModelType);
}