// Copyright 2026 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.
//

#include "cbor/cbor_generator.h"

#include <memory>
#include <string>

#include "absl/base/no_destructor.h"
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/strings/ascii.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_replace.h"
#include "absl/strings/string_view.h"
#include "cbor/tags.pb.h"
#include "google/protobuf/compiler/code_generator.h"
#include "google/protobuf/compiler/cpp/names.h"
#include "google/protobuf/descriptor.h"
#include "google/protobuf/io/printer.h"
#include "google/protobuf/io/zero_copy_stream.h"

namespace cbor {
namespace {

std::string QualifiedCppName(const google::protobuf::Descriptor* descriptor) {
  return google::protobuf::compiler::cpp::QualifiedClassName(descriptor);
}

absl::string_view StripProtoExtension(absl::string_view filename) {
  if (absl::EndsWith(filename, ".proto")) {
    return filename.substr(0, filename.size() - 6);
  }
  return filename;
}

const cbor::Tags& GetTags(const google::protobuf::FieldDescriptor* field) {
  if (field->options().HasExtension(cbor::tags)) {
    return field->options().GetExtension(cbor::tags);
  }
  static const absl::NoDestructor<cbor::Tags> default_tags;
  return *default_tags;
}

std::string GetCborKey(const google::protobuf::FieldDescriptor* field) {
  const auto& tags = GetTags(field);
  if (!tags.name().empty()) {
    return std::string(tags.name());
  }
  return std::string(field->name());
}

const absl::flat_hash_map<google::protobuf::FieldDescriptor::Type,
                          absl::string_view>&
GetProtoTypeToCborType() {
  static const absl::NoDestructor<absl::flat_hash_map<
      google::protobuf::FieldDescriptor::Type, absl::string_view>>
      proto_type_to_cbor_type({
          {google::protobuf::FieldDescriptor::TYPE_MESSAGE, "Map"},
          {google::protobuf::FieldDescriptor::TYPE_STRING, "String"},
          {google::protobuf::FieldDescriptor::TYPE_UINT64, "Uint64"},
          {google::protobuf::FieldDescriptor::TYPE_INT64, "Int64"},
          {google::protobuf::FieldDescriptor::TYPE_UINT32, "Uint64"},
          {google::protobuf::FieldDescriptor::TYPE_INT32, "Int64"},
          {google::protobuf::FieldDescriptor::TYPE_BOOL, "Bool"},
          {google::protobuf::FieldDescriptor::TYPE_FLOAT, "Float"},
          {google::protobuf::FieldDescriptor::TYPE_DOUBLE, "Double"},
          {google::protobuf::FieldDescriptor::TYPE_BYTES, "ByteString"},
      });
  return *proto_type_to_cbor_type;
}

absl::string_view GetCborTypeName(
    const google::protobuf::FieldDescriptor* field) {
  const auto& tags = GetTags(field);
  if (tags.type() != cbor::Tags::TYPE_UNSPECIFIED) {
    switch (tags.type()) {
      case cbor::Tags::TYPE_UINT:
        return "Uint64";
      case cbor::Tags::TYPE_NINT:
        return "Int64";
      case cbor::Tags::TYPE_BYTE_STRING:
        return "ByteString";
      case cbor::Tags::TYPE_TEXT_STRING:
        return "String";
      default:
        break;
    }
  }
  auto it = GetProtoTypeToCborType().find(field->type());
  if (it != GetProtoTypeToCborType().end()) {
    return it->second;
  }
  return "";
}

void GenerateMessageHeader(const google::protobuf::Descriptor* desc,
                           google::protobuf::io::Printer& printer) {
  printer.Emit({{"name", QualifiedCppName(desc)}}, R"cc(
    absl::Status ToProto(const cbor::MapView& map, $name$* proto);
    absl::StatusOr<std::unique_ptr<cppbor::Map>> FromProto(const $name$& proto);
    absl::StatusOr<std::unique_ptr<cppbor::Map>> FromProto(
        const $name$& proto, FromProtoOptions options);
  )cc");
  for (int i = 0; i < desc->nested_type_count(); ++i) {
    GenerateMessageHeader(desc->nested_type(i), printer);
  }
}

void GenerateToProtoBody(const google::protobuf::Descriptor* desc,
                         google::protobuf::io::Printer& printer) {
  printer.Emit({{"name", QualifiedCppName(desc)}}, R"cc(
    absl::Status ToProto(const cbor::MapView& map, $name$* proto) {
  )cc");
  {
    auto indent = printer.WithIndent();
    for (int i = 0; i < desc->field_count(); ++i) {
      const auto* field = desc->field(i);
      const auto& tags = GetTags(field);
      if (tags.mode() == cbor::Tags::MODE_WRITE_ONLY ||
          tags.mode() == cbor::Tags::MODE_SKIP) {
        continue;
      }

      std::string key = GetCborKey(field);
      printer.Emit(R"cc(
        {
      )cc");
      {
        auto scope_indent = printer.WithIndent();
        if (field->is_repeated()) {
          if (field->type() ==
              google::protobuf::FieldDescriptor::TYPE_MESSAGE) {
            printer.Emit({{"field", field->name()}, {"key", key}}, R"cc(
              auto array_view = map.GetOptionalArray("$key$");
            )cc");
            if (tags.is_required()) {
              printer.Emit({{"key", key}}, R"cc(
                if (!array_view.has_value()) {
                  return absl::InvalidArgumentError(
                      "Missing required field: "
                      "$key$");
                }
              )cc");
            }
            if (tags.disallow_empty()) {
              printer.Emit({{"key", key}}, R"cc(
                if (array_view.has_value() && array_view->size() == 0) {
                  return absl::InvalidArgumentError(
                      "Field $key$ cannot be "
                      "empty");
                }
              )cc");
            }
            printer.Emit({{"field", field->name()}}, R"cc(
              if (array_view.has_value()) {
                for (uint32_t i = 0; i < array_view->size(); ++i) {
                  ABSL_ASSIGN_OR_RETURN(auto item, array_view->GetMap(i));
                  ABSL_RETURN_IF_ERROR(ToProto(std::move(item), proto->add_$field$()));
                }
              }
            )cc");
          } else if (!GetCborTypeName(field).empty()) {
            absl::string_view type = GetCborTypeName(field);
            printer.Emit(
                {{"field", field->name()}, {"key", key}, {"type", type}}, R"cc(
                  auto array_view = map.GetOptionalArray("$key$");
                )cc");
            if (tags.is_required()) {
              printer.Emit({{"key", key}}, R"cc(
                if (!array_view.has_value()) {
                  return absl::InvalidArgumentError(
                      "Missing required field: "
                      "$key$");
                }
              )cc");
            }
            if (tags.disallow_empty()) {
              printer.Emit({{"key", key}}, R"cc(
                if (array_view.has_value() && array_view->size() == 0) {
                  return absl::InvalidArgumentError(
                      "Field $key$ cannot be "
                      "empty");
                }
              )cc");
            }
            printer.Emit({{"field", field->name()}, {"type", type}}, R"cc(
              if (array_view.has_value()) {
                for (uint32_t i = 0; i < array_view->size(); ++i) {
                  ABSL_ASSIGN_OR_RETURN(auto item, array_view->Get$type$(i));
                  proto->add_$field$(std::move(item));
                }
              }
            )cc");
          }
        } else {
          if (field->type() ==
              google::protobuf::FieldDescriptor::TYPE_MESSAGE) {
            if (tags.is_required()) {
              printer.Emit({{"field", field->name()}, {"key", key}}, R"cc(
                auto map_view = map.GetOptionalMap("$key$");
                if (!map_view.has_value()) {
                  return absl::InvalidArgumentError(
                      "Missing required field: "
                      "$key$");
                }
                ABSL_RETURN_IF_ERROR(ToProto(*map_view, proto->mutable_$field$()));
              )cc");
            } else {
              printer.Emit({{"field", field->name()}, {"key", key}}, R"cc(
                auto map_view = map.GetOptionalMap("$key$");
                if (map_view.has_value()) {
                  ABSL_RETURN_IF_ERROR(ToProto(*map_view, proto->mutable_$field$()));
                }
              )cc");
            }
          } else if (!GetCborTypeName(field).empty()) {
            absl::string_view type = GetCborTypeName(field);
            if (tags.is_required()) {
              printer.Emit(
                  {{"field", field->name()}, {"key", key}, {"type", type}},
                  R"cc(
                    auto val = map.GetOptional$type$("$key$");
                    if (!val.has_value()) {
                      return absl::InvalidArgumentError(
                          "Missing required "
                          "field: $key$");
                    }
                    proto->set_$field$(*val);
                  )cc");
            } else {
              printer.Emit(
                  {{"field", field->name()}, {"key", key}, {"type", type}},
                  R"cc(
                    auto val = map.GetOptional$type$("$key$");
                    if (val.has_value()) {
                      proto->set_$field$(*val);
                    }
                  )cc");
            }
          }
        }
      }
      printer.Emit(R"cc(
        }
      )cc");
    }
    printer.Emit(R"cc(
      return absl::OkStatus();
    )cc");
  }
  printer.Emit(R"cc(
    }
  )cc");
  for (int i = 0; i < desc->nested_type_count(); ++i) {
    GenerateToProtoBody(desc->nested_type(i), printer);
  }
}

void GenerateFromProtoBody(const google::protobuf::Descriptor* desc,
                           google::protobuf::io::Printer& printer) {
  printer.Emit({{"name", QualifiedCppName(desc)}}, R"cc(
    absl::StatusOr<std::unique_ptr<cppbor::Map>> FromProto(
        const $name$& proto) {
      return FromProto(proto, FromProtoOptions{});
    }

    absl::StatusOr<std::unique_ptr<cppbor::Map>> FromProto(
        const $name$& proto, FromProtoOptions options) {
      auto map = std::make_unique<cppbor::Map>();
  )cc");
  {
    auto indent = printer.WithIndent();
    for (int i = 0; i < desc->field_count(); ++i) {
      const auto* field = desc->field(i);
      const auto& tags = GetTags(field);

      if (tags.mode() == cbor::Tags::MODE_READ_ONLY ||
          tags.mode() == cbor::Tags::MODE_SKIP) {
        continue;
      }

      std::string key = GetCborKey(field);
      if (field->is_repeated()) {
        if (tags.is_required() && tags.disallow_empty()) {
          printer.Emit({{"field", field->name()}, {"key", key}}, R"cc(
            if (!options.skip_validation && proto.$field$_size() == 0) {
              return absl::InvalidArgumentError("Field $key$ cannot be empty");
            }
          )cc");
        }
        if (tags.is_required()) {
          printer.Emit(R"cc(
            {
          )cc");
        } else {
          printer.Emit({{"field", field->name()}}, R"cc(
            if (proto.$field$_size() > 0) {
          )cc");
        }

        if (field->type() == google::protobuf::FieldDescriptor::TYPE_MESSAGE) {
          printer.Emit({{"field", field->name()}, {"key", key}}, R"cc(
            auto array = std::make_unique<cppbor::Array>();
            for (int i = 0; i < proto.$field$_size(); ++i) {
              ABSL_ASSIGN_OR_RETURN(auto item, FromProto(proto.$field$(i)));
              array->add(std::move(item));
            }
            map->add("$key$", std::move(array));
          )cc");
        } else if (field->type() ==
                   google::protobuf::FieldDescriptor::TYPE_BYTES) {
          printer.Emit({{"field", field->name()}, {"key", key}}, R"cc(
            auto array = std::make_unique<cppbor::Array>();
            for (int i = 0; i < proto.$field$_size(); ++i) {
              array->add(cppbor::Bstr(std::string(proto.$field$(i))));
            }
            map->add("$key$", std::move(array));
          )cc");
        } else if (field->type() ==
                   google::protobuf::FieldDescriptor::TYPE_STRING) {
          printer.Emit({{"field", field->name()}, {"key", key}}, R"cc(
            auto array = std::make_unique<cppbor::Array>();
            for (int i = 0; i < proto.$field$_size(); ++i) {
              array->add(cppbor::Tstr(std::string(proto.$field$(i))));
            }
            map->add("$key$", std::move(array));
          )cc");
        } else if (field->type() ==
                   google::protobuf::FieldDescriptor::TYPE_FLOAT) {
          printer.Emit({{"field", field->name()}, {"key", key}}, R"cc(
            auto array = std::make_unique<cppbor::Array>();
            for (int i = 0; i < proto.$field$_size(); ++i) {
              array->add(cppbor::Float(proto.$field$(i)));
            }
            map->add("$key$", std::move(array));
          )cc");
        } else if (field->type() ==
                   google::protobuf::FieldDescriptor::TYPE_DOUBLE) {
          printer.Emit({{"field", field->name()}, {"key", key}}, R"cc(
            auto array = std::make_unique<cppbor::Array>();
            for (int i = 0; i < proto.$field$_size(); ++i) {
              array->add(cppbor::Double(proto.$field$(i)));
            }
            map->add("$key$", std::move(array));
          )cc");
        } else {
          printer.Emit({{"field", field->name()}, {"key", key}}, R"cc(
            auto array = std::make_unique<cppbor::Array>();
            for (int i = 0; i < proto.$field$_size(); ++i) {
              array->add(proto.$field$(i));
            }
            map->add("$key$", std::move(array));
          )cc");
        }

        printer.Emit(R"cc(
          }
        )cc");
      } else {
        if (tags.is_required()) {
          printer.Emit({{"field", field->name()}, {"key", key}}, R"cc(
            if (!options.skip_validation && !proto.has_$field$()) {
              return absl::InvalidArgumentError(
                  "Missing required field: "
                  "$key$");
            }
          )cc");
        }
        if (field->type() == google::protobuf::FieldDescriptor::TYPE_MESSAGE) {
          printer.Emit({{"field", field->name()}, {"key", key}}, R"cc(
            if (proto.has_$field$()) {
              ABSL_ASSIGN_OR_RETURN(auto item, FromProto(proto.$field$()));
              map->add("$key$", std::move(item));
            }
          )cc");
        } else if (field->type() ==
                   google::protobuf::FieldDescriptor::TYPE_BYTES) {
          printer.Emit({{"field", field->name()}, {"key", key}}, R"cc(
            if (proto.has_$field$()) {
              map->add("$key$", cppbor::Bstr(std::string(proto.$field$())));
            }
          )cc");
        } else if (field->type() ==
                   google::protobuf::FieldDescriptor::TYPE_STRING) {
          printer.Emit({{"field", field->name()}, {"key", key}}, R"cc(
            if (proto.has_$field$()) {
              map->add("$key$", cppbor::Tstr(std::string(proto.$field$())));
            }
          )cc");
        } else if (field->type() ==
                   google::protobuf::FieldDescriptor::TYPE_FLOAT) {
          printer.Emit({{"field", field->name()}, {"key", key}}, R"cc(
            if (proto.has_$field$()) {
              map->add("$key$", cppbor::Float(proto.$field$()));
            }
          )cc");
        } else if (field->type() ==
                   google::protobuf::FieldDescriptor::TYPE_DOUBLE) {
          printer.Emit({{"field", field->name()}, {"key", key}}, R"cc(
            if (proto.has_$field$()) {
              map->add("$key$", cppbor::Double(proto.$field$()));
            }
          )cc");
        } else {
          printer.Emit({{"field", field->name()}, {"key", key}}, R"cc(
            if (proto.has_$field$()) {
              map->add("$key$", proto.$field$());
            }
          )cc");
        }
      }
    }
    printer.Emit(R"cc(
      return map;
    )cc");
  }
  printer.Emit(R"cc(
    }
  )cc");
  for (int i = 0; i < desc->nested_type_count(); ++i) {
    GenerateFromProtoBody(desc->nested_type(i), printer);
  }
}

void CollectUsedDeps(
    const google::protobuf::Descriptor* descriptor,
    const google::protobuf::FileDescriptor* current_file,
    absl::flat_hash_set<const google::protobuf::FileDescriptor*>& used_deps) {
  for (int i = 0; i < descriptor->field_count(); ++i) {
    const auto* f = descriptor->field(i);
    if (f->type() == google::protobuf::FieldDescriptor::TYPE_MESSAGE) {
      if (GetTags(f).mode() == cbor::Tags::MODE_SKIP) {
        continue;
      }
      const auto* m = f->message_type();
      if (m->file() != current_file) {
        used_deps.insert(m->file());
      }
    }
  }
  for (int i = 0; i < descriptor->nested_type_count(); ++i) {
    CollectUsedDeps(descriptor->nested_type(i), current_file, used_deps);
  }
}

void GenerateHeader(const google::protobuf::FileDescriptor* file,
                    absl::string_view base_name,
                    google::protobuf::compiler::GeneratorContext* context) {
  std::unique_ptr<google::protobuf::io::ZeroCopyOutputStream> h_stream(
      context->Open(absl::StrCat(base_name, ".cbor.h")));
  google::protobuf::io::Printer h_printer(h_stream.get());

  std::string guard = absl::AsciiStrToUpper(absl::StrReplaceAll(
      absl::StrCat(base_name, "_CBOR_H_"), {{"/", "_"}, {".", "_"}}));
  h_printer.Emit({{"guard", guard}}, R"cc(
#ifndef $guard$
#define $guard$
  )cc");

  if (file->message_type_count() > 0) {
    h_printer.Emit(R"cc(
#include <memory>

#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "cbor/cbor.h"
#include "cbor/options.h"
#include "cppbor/cppbor.h"
    )cc");
  }

  h_printer.Emit({{"base", base_name}}, R"cc(
#include "$base$.pb.h"
  )cc");

  for (int i = 0; i < file->public_dependency_count(); ++i) {
    absl::string_view pub_base =
        StripProtoExtension(file->public_dependency(i)->name());
    h_printer.Emit({{"pub", pub_base}}, R"cc(
#include "$pub$.cbor.h"  // IWYU pragma: export
    )cc");
  }

  h_printer.Emit(R"cc(
    namespace cbor {
  )cc");

  for (int i = 0; i < file->message_type_count(); ++i) {
    GenerateMessageHeader(file->message_type(i), h_printer);
  }

  h_printer.Emit({{"guard", guard}}, R"cc(
    }  // namespace cbor

#endif  // $guard$
  )cc");
}

void GenerateCc(const google::protobuf::FileDescriptor* file,
                absl::string_view base_name,
                google::protobuf::compiler::GeneratorContext* context) {
  std::unique_ptr<google::protobuf::io::ZeroCopyOutputStream> cc_stream(
      context->Open(absl::StrCat(base_name, ".cbor.cc")));
  google::protobuf::io::Printer cc_printer(cc_stream.get());

  cc_printer.Emit({{"base", base_name}}, R"cc(
#include "$base$.cbor.h"
  )cc");

  if (file->message_type_count() > 0) {
    cc_printer.Emit({{"base", base_name}}, R"cc(
#include "$base$.pb.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "cbor/cbor.h"
#include "cppbor/cppbor.h"
    )cc");
  }

  absl::flat_hash_set<const google::protobuf::FileDescriptor*> used_deps;
  for (int i = 0; i < file->message_type_count(); ++i) {
    CollectUsedDeps(file->message_type(i), file, used_deps);
  }

  for (const auto* dep : used_deps) {
    absl::string_view dep_base = StripProtoExtension(dep->name());
    cc_printer.Emit({{"dep", dep_base}}, R"cc(
#include "$dep$.cbor.h"
    )cc");
  }

  cc_printer.Emit(R"cc(
    namespace cbor {
  )cc");

  for (int i = 0; i < file->message_type_count(); ++i) {
    GenerateToProtoBody(file->message_type(i), cc_printer);
    GenerateFromProtoBody(file->message_type(i), cc_printer);
  }

  cc_printer.Emit(R"cc(
    }  // namespace cbor
  )cc");
}

}  // namespace

bool CborGenerator::Generate(
    const google::protobuf::FileDescriptor* file, const std::string& parameter,
    google::protobuf::compiler::GeneratorContext* context,
    std::string* error) const {
  absl::string_view base_name = StripProtoExtension(file->name());

  GenerateHeader(file, base_name, context);
  GenerateCc(file, base_name, context);

  return true;
}

}  // namespace cbor