package io.apicurio.registry.maven;
   
   import com.google.protobuf.Descriptors;
   import com.squareup.wire.schema.internal.parser.ProtoFileElement;
   import io.apicurio.registry.content.ContentHandle;
   import io.apicurio.registry.content.TypedContent;
   import io.apicurio.registry.rest.client.RegistryClient;
   import io.apicurio.registry.rest.client.models.ArtifactReference;
   import io.apicurio.registry.types.ContentTypes;
  import io.apicurio.registry.utils.protobuf.schema.FileDescriptorUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.io.File;
  import java.io.FileNotFoundException;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Objects;
  import java.util.Set;
  import java.util.concurrent.ExecutionException;
  import java.util.stream.Collectors;
  
  public class ProtobufDirectoryParser extends AbstractDirectoryParser<Descriptors.FileDescriptor> {
  
      private static final String PROTO_SCHEMA_EXTENSION = ".proto";
      private static final Logger log = LoggerFactory.getLogger(ProtobufDirectoryParser.class);
  
      public ProtobufDirectoryParser(RegistryClient client) {
          super(client);
      }
  
      @Override
      public ParsedDirectoryWrapper<Descriptors.FileDescriptor> parse(File protoFile) {
  
          Set<File> protoFiles = Arrays
                  .stream(Objects.requireNonNull(protoFile.getParentFile()
                          .listFiles((dir, name) -> name.endsWith(PROTO_SCHEMA_EXTENSION))))
                  .filter(file -> !file.getName().equals(protoFile.getName())).collect(Collectors.toSet());
  
          try {
              final Map<String, String> requiredSchemaDefs = new HashMap<>();
              final Descriptors.FileDescriptor schemaDescriptor = FileDescriptorUtils
                      .parseProtoFileWithDependencies(protoFile, protoFiles, requiredSchemaDefs);
              assert allDependenciesHaveSamePackageName(requiredSchemaDefs, schemaDescriptor.getPackage())
                      : "All dependencies must have the same package name as the main proto file";
              Map<String, TypedContent> schemaContents = convertSchemaDefs(requiredSchemaDefs,
                      schemaDescriptor.getPackage());
              return new DescriptorWrapper(schemaDescriptor, schemaContents);
          } catch (Descriptors.DescriptorValidationException e) {
              throw new RuntimeException("Failed to read schema file: " + protoFile, e);
          } catch (FileDescriptorUtils.ReadSchemaException e) {
              log.warn(
                      "Error processing Avro schema with name {}. This usually means that the references are not ready yet to read it",
                      e.file());
              throw new RuntimeException(e.getCause());
          } catch (FileDescriptorUtils.ParseSchemaException e) {
              log.warn(
                      "Error processing Avro schema with name {}. This usually means that the references are not ready yet to parse it",
                      e.fileName());
              throw new RuntimeException(e.getCause());
          }
      }
  
      private static boolean allDependenciesHaveSamePackageName(Map<String, String> schemas,
              String mainProtoPackageName) {
          return schemas.keySet().stream().allMatch(fullDepName -> fullDepName.contains(mainProtoPackageName));
      }
  
      /**
       * Converts the schema definitions to a map of ContentHandle, stripping any package information from the
       * key, which is not needed for the schema registry, given that the dependent schemas are *always* in the
       * same package of the main proto file.
       */
      private Map<String, TypedContent> convertSchemaDefs(Map<String, String> requiredSchemaDefs,
              String mainProtoPackageName) {
          if (requiredSchemaDefs.isEmpty()) {
              return Map.of();
          }
          Map<String, TypedContent> schemaDefs = new HashMap<>(requiredSchemaDefs.size());
          for (Map.Entry<String, String> entry : requiredSchemaDefs.entrySet()) {
              String fileName = FileDescriptorUtils.extractProtoFileName(entry.getKey());
              TypedContent content = TypedContent.create(ContentHandle.create(entry.getValue()),
                      ContentTypes.APPLICATION_PROTOBUF);
              if (schemaDefs.put(fileName, content) != null) {
                  log.warn(
                          "There's a clash of dependency name, likely due to stripping the expected package name ie {}: dependencies: {}",
                          mainProtoPackageName,
                          Arrays.toString(requiredSchemaDefs.keySet().toArray(new Object[0])));
              }
          }
          return schemaDefs;
      }
  
      @Override
      public List<ArtifactReference> handleSchemaReferences(RegisterArtifact rootArtifact,
             Descriptors.FileDescriptor protoSchema, Map<String, TypedContent> fileContents)
             throws FileNotFoundException, InterruptedException, ExecutionException {
         Set<ArtifactReference> references = new HashSet<>();
         final Set<Descriptors.FileDescriptor> baseDeps = new HashSet<>(
                 Arrays.asList(FileDescriptorUtils.baseDependencies()));
         final ProtoFileElement rootSchemaElement = FileDescriptorUtils
                 .fileDescriptorToProtoFile(protoSchema.toProto());
 
         for (Descriptors.FileDescriptor dependency : protoSchema.getDependencies()) {
 
             List<ArtifactReference> nestedArtifactReferences = new ArrayList<>();
             String dependencyFullName = dependency.getPackage() + "/" + dependency.getName(); // FIXME find a
                                                                                               // better wat to
                                                                                               // do this
             if (!baseDeps.contains(dependency)
                     && rootSchemaElement.getImports().contains(dependencyFullName)) {
 
                 RegisterArtifact nestedArtifact = buildFromRoot(rootArtifact, dependencyFullName);
 
                 if (!dependency.getDependencies().isEmpty()) {
                     nestedArtifactReferences = handleSchemaReferences(nestedArtifact, dependency,
                             fileContents);
                 }
 
                 references.add(registerNestedSchema(dependencyFullName, nestedArtifactReferences,
                         nestedArtifact, fileContents.get(dependency.getName()).getContent().content()));
             }
         }
 
         return new ArrayList<>(references);
     }
 
     public static class DescriptorWrapper implements ParsedDirectoryWrapper<Descriptors.FileDescriptor> {
         final Descriptors.FileDescriptor fileDescriptor;
         final Map<String, TypedContent> schemaContents; // used to store the original file content to register
                                                         // the content as-is.
 
         public DescriptorWrapper(Descriptors.FileDescriptor fileDescriptor,
                 Map<String, TypedContent> schemaContents) {
             this.fileDescriptor = fileDescriptor;
             this.schemaContents = schemaContents;
         }
 
         @Override
         public Descriptors.FileDescriptor getSchema() {
             return fileDescriptor;
         }
 
         public Map<String, TypedContent> getSchemaContents() {
             return schemaContents;
         }
     }
 }