// Copyright 2023 Google LLC
//
// 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
//
//     https://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 MPACT_SIM_DECODER_PROTO_FORMAT_VISITOR_H_
#define MPACT_SIM_DECODER_PROTO_FORMAT_VISITOR_H_

#include <deque>
#include <functional>
#include <list>
#include <memory>
#include <string>
#include <string_view>
#include <utility>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "antlr4-runtime/ParserRuleContext.h"
#include "mpact/sim/decoder/ProtoFormatLexer.h"
#include "mpact/sim/decoder/ProtoFormatParser.h"
#include "mpact/sim/decoder/antlr_parser_wrapper.h"
#include "mpact/sim/decoder/decoder_error_listener.h"
#include "mpact/sim/decoder/proto_constraint_expression.h"
#include "mpact/sim/decoder/proto_encoding_info.h"
#include "mpact/sim/decoder/proto_format_contexts.h"
#include "mpact/sim/decoder/proto_instruction_group.h"
#include "re2/re2.h"
#include "src/google/protobuf/descriptor.h"

// This file defines the visitor class used to generate C++ code for decoding
// instructions encoded in protobufs.

namespace mpact {
namespace sim {
namespace decoder {
namespace proto_fmt {

class ProtoInstructionEncoding;

using ::mpact::sim::decoder::proto_fmt::generated::ProtoFormatLexer;
using ::mpact::sim::decoder::proto_fmt::generated::ProtoFormatParser;

using ProtoFmtAntlrParserWrapper =
    decoder::AntlrParserWrapper<ProtoFormatParser, ProtoFormatLexer>;

class ProtoFormatVisitor {
 public:
  ProtoFormatVisitor() = default;
  ~ProtoFormatVisitor() = default;

  // This is the main function. It parses the .proto_fmt file and generates
  // the required C++ code.
  // Parameters:
  //   file_names: vector of source file names.
  //   decoder_name: the decoder for which to generate code.
  //   prefix: if non-empty, the prefix used in the generated file names.
  //   include_roots: vector of directories from which to resolve include files.
  //   proto_dirs: vector of directories from which to resolve proto files.
  //   proto_files: vector of .proto files to import.
  //   directory: target directory to generate the C++ files in.
  absl::Status Process(const std::vector<std::string> &file_names,
                       const std::string &decoder_name, std::string_view prefix,
                       const std::vector<std::string> &include_roots,
                       const std::vector<std::string> &proto_dirs,
                       const std::vector<std::string> &proto_files,
                       std::string_view directory);

  // Accessors for the error listener.
  decoder::DecoderErrorListener *error_listener() const {
    return error_listener_.get();
  }
  void set_error_listener(
      std::unique_ptr<decoder::DecoderErrorListener> listener) {
    error_listener_ = std::move(listener);
  }

 private:
  // This struct holds information about a range assignment in an instruction
  // generator.
  struct RangeAssignmentInfo {
    std::vector<std::string> range_names;
    std::list<RE2> range_regexes;
    std::vector<std::vector<std::string>> range_values;
  };

  // Expand the given name according to any using statements that may apply.
  std::string Expand(std::string_view name) const;
  // Get the field descriptor for the named field from the message type while
  // building up a vector of one_of_fields that are on the path from the top
  // level message to the field.
  const google::protobuf::FieldDescriptor *GetField(
      const std::string &field_name,
      const google::protobuf::Descriptor *message_type,
      std::vector<const google::protobuf::FieldDescriptor *> &one_of_fields)
      const;
  // Get the descriptor for the named enumeration member.
  const google::protobuf::EnumValueDescriptor *GetEnumValueDescriptor(
      const std::string &full_name) const;
  // Get the numeric value of the named enumeration member.
  absl::StatusOr<int> GetEnumValue(const std::string &enum_name) const;
  // Helpful template function used to find a descriptor by name.
  template <typename T>
  const T *FindByName(
      const std::string &message_name, const std::string &name,
      std::function<const T *(const std::string &)> finder) const;
  // Visit the parse tree to catalog the different declarations.
  void PreProcessDeclarations(
      const std::vector<DeclarationCtx *> &declarations);
  // Process an include file declaration.
  void VisitIncludeFile(IncludeFileCtx *ctx);
  // Parse an included file.
  void ParseIncludeFile(antlr4::ParserRuleContext *ctx,
                        const std::string &file_name,
                        const std::vector<std::string> &dirs);
  // Creates the top level data structures and visits the declarations that
  // are necessary to generate the instruction decoder.
  std::unique_ptr<ProtoEncodingInfo> ProcessTopLevel(
      const std::string &decoder_name);
  // The following methods visit specific nodes in the parse tree.
  ProtoInstructionGroup *VisitInstructionGroupDef(
      InstructionGroupDefCtx *ctx, ProtoEncodingInfo *encoding_info);
  void VisitInstructionDef(InstructionDefCtx *ctx,
                           ProtoInstructionGroup *inst_group,
                           ProtoEncodingInfo *encoding_info);
  void VisitFieldConstraint(FieldConstraintCtx *ctx,
                            ProtoInstructionEncoding *inst_encoding,
                            const ProtoInstructionGroup *inst_group);
  ProtoConstraintExpression *VisitConstraintExpression(
      ConstraintExprCtx *ctx,
      const google::protobuf::FieldDescriptor *field_desc,
      const ProtoInstructionGroup *inst_group);
  ProtoConstraintExpression *VisitValue(ValueCtx *ctx);
  ProtoConstraintExpression *VisitNumber(NumberCtx *ctx);
  ProtoConstraintExpression *VisitQualifiedIdent(
      QualifiedIdentCtx *ctx,
      const google::protobuf::FieldDescriptor *field_desc,
      const ProtoInstructionGroup *inst_group);
  void VisitSetterGroupDef(SetterGroupDefCtx *ctx,
                           ProtoInstructionGroup *inst_group,
                           ProtoEncodingInfo *encoding_info);
  void VisitSetterDef(SetterDefCtx *ctx,
                      ProtoInstructionEncoding *inst_encoding,
                      ProtoInstructionGroup *inst_group,
                      ProtoEncodingInfo *encoding_info);
  void VisitSetterRef(SetterRefCtx *ctx,
                      ProtoInstructionEncoding *inst_encoding,
                      ProtoInstructionGroup *inst_group,
                      ProtoEncodingInfo *encoding_info);
  std::unique_ptr<ProtoEncodingInfo> VisitDecoderDef(DecoderDefCtx *ctx);
  // Process and generate instruction definitions from a generate statement.
  void ProcessInstructionDefGenerator(InstructionDefCtx *ctx,
                                      ProtoInstructionGroup *inst_group,
                                      ProtoEncodingInfo *encoding_info);
  std::string GenerateInstructionDefList(
      const std::vector<RangeAssignmentInfo *> &range_info_vec, int index,
      const std::string &template_str_in) const;
  // Process instruction groups.
  void ProcessSingleGroup(DecoderAttributeCtx *attr_ctx,
                          ProtoEncodingInfo *encoding_info,
                          absl::flat_hash_set<std::string> &group_name_set);
  void ProcessParentGroup(DecoderAttributeCtx *attr_ctx,
                          ProtoEncodingInfo *encoding_info,
                          absl::flat_hash_set<std::string> &group_name_set);

  // Called to generate and emit code for the decoder according to the parsed
  // input file.
  StringPair EmitCode(ProtoEncodingInfo *encoding_info);

  // Finders used to find specific object types from the proto2 pool.
  using FieldFinder = std::function<const google::protobuf::FieldDescriptor *(
      const std::string &)>;
  using MessageFinder =
      std::function<const google::protobuf::Descriptor *(const std::string &)>;
  using EnumTypeFinder = std::function<const google::protobuf::EnumDescriptor *(
      const std::string &)>;
  using EnumValueFinder =
      std::function<const google::protobuf::EnumValueDescriptor *(
          const std::string &)>;

  FieldFinder field_finder_;
  MessageFinder message_finder_;
  EnumTypeFinder enum_type_finder_;
  EnumValueFinder enum_value_finder_;

  // This stores a vector of include file root directories.
  std::vector<std::string> include_dir_vec_;
  // Keep track of files that are included in case there is recursive includes.
  std::deque<std::string> include_file_stack_;
  // Error listening object passed to the parser.
  std::unique_ptr<decoder::DecoderErrorListener> error_listener_ = nullptr;
  std::string decoder_name_;
  // Descriptor pool.
  const google::protobuf::DescriptorPool *descriptor_pool_;
  absl::flat_hash_map<std::string, const google::protobuf::FileDescriptor *>
      file_descriptor_map_;
  // Using decl map.
  absl::flat_hash_map<std::string, std::string> using_decl_map_;
  // Maps from identifiers to declaration contexts.
  absl::flat_hash_map<std::string, InstructionGroupDefCtx *> group_decl_map_;
  absl::flat_hash_map<std::string, DecoderDefCtx *> decoder_decl_map_;
  // AntlrParserWrapper vector.
  std::vector<ProtoFmtAntlrParserWrapper *> antlr_parser_wrappers_;
};

}  // namespace proto_fmt
}  // namespace decoder
}  // namespace sim
}  // namespace mpact

#endif  // MPACT_SIM_DECODER_PROTO_FORMAT_VISITOR_H_