package io.apicurio.registry.serde.protobuf;

import com.google.protobuf.DescriptorProtos;
import com.google.protobuf.Descriptors;
import com.google.protobuf.Descriptors.DescriptorValidationException;
import com.google.protobuf.Descriptors.FileDescriptor;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.Message;
import com.squareup.wire.schema.internal.parser.MessageElement;
import com.squareup.wire.schema.internal.parser.ProtoFileElement;
import com.squareup.wire.schema.internal.parser.ProtoParser;
import io.apicurio.registry.resolver.ParsedSchema;
import io.apicurio.registry.resolver.ParsedSchemaImpl;
import io.apicurio.registry.resolver.SchemaParser;
import io.apicurio.registry.resolver.data.Record;
import io.apicurio.registry.types.ArtifactType;
import io.apicurio.registry.utils.IoUtil;
import io.apicurio.registry.utils.protobuf.schema.FileDescriptorUtils;
import io.apicurio.registry.utils.protobuf.schema.ProtobufSchema;

import java.util.ArrayList;
import java.util.Base64;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class ProtobufSchemaParser<U extends Message> implements SchemaParser<ProtobufSchema, U> {

    /**
     * @see io.apicurio.registry.resolver.SchemaParser#artifactType()
     */
    @Override
    public String artifactType() {
        return ArtifactType.PROTOBUF;
    }

    /**
     * @see io.apicurio.registry.resolver.SchemaParser#parseSchema(byte[], Map<String,
     *      ParsedSchema<ProtobufSchema>>)
     */
    @Override
    public ProtobufSchema parseSchema(byte[] rawSchema,
            Map<String, ParsedSchema<ProtobufSchema>> resolvedReferences) {
        try {
            // textual .proto file
            ProtoFileElement fileElem = ProtoParser.Companion.parse(FileDescriptorUtils.DEFAULT_LOCATION,
                    IoUtil.toString(rawSchema));
            Map<String, ProtoFileElement> dependencies = new HashMap<>();
            resolvedReferences.forEach((key, value) -> {
                dependencies.put(key, value.getParsedSchema().getProtoFileElement());
                if (value.hasReferences()) {
                    addReferencesToDependencies(value.getSchemaReferences(), dependencies);
                }
            });
            MessageElement firstMessage = FileDescriptorUtils.firstMessage(fileElem);
            if (firstMessage != null) {
                try {
                    final Descriptors.Descriptor descriptor = FileDescriptorUtils
                            .toDescriptor(firstMessage.getName(), fileElem, dependencies);
                    if (descriptor == null) {
                        // toDescriptor returns null when the message can't be resolved, typically due to
                        // missing imports (e.g. google/protobuf/timestamp.proto). Fall through to the
                        // alternative parsing method.
                        return getFileDescriptorFromElement(fileElem);
                    }
                    return new ProtobufSchema(descriptor.getFile(), fileElem);
                } catch (IllegalStateException ise) {
                    // If we fail to init the dynamic schema, try to get the descriptor from the proto element
                    return getFileDescriptorFromElement(fileElem);
                }
            } else {
                return getFileDescriptorFromElement(fileElem);
            }
        } catch (DescriptorValidationException pe) {
            throw new IllegalStateException("Error parsing protobuf schema ", pe);
        } catch (IllegalStateException illegalStateException) {
            // If we get here the server likely returned the full descriptor, try to parse it.
            return parseDescriptor(rawSchema, resolvedReferences);
        }
    }

    private ProtobufSchema parseDescriptor(byte[] rawSchema,
            Map<String, ParsedSchema<ProtobufSchema>> resolvedReferences) {
        // Try to parse the binary format, in case the server has returned the descriptor format.
        try {
            DescriptorProtos.FileDescriptorProto fileDescriptorProto = parseFileDescriptorProto(rawSchema);
            ProtoFileElement protoFileElement = FileDescriptorUtils
                    .fileDescriptorToProtoFile(fileDescriptorProto);

            if (resolvedReferences == null || resolvedReferences.isEmpty()) {
                return new ProtobufSchema(FileDescriptorUtils.protoFileToFileDescriptor(fileDescriptorProto),
                        protoFileElement);
            }

            // Build dependency map from resolved references
            Map<String, Descriptors.FileDescriptor> deps = new HashMap<>();
            for (Descriptors.FileDescriptor fd : FileDescriptorUtils.baseDependencies()) {
                deps.put(fd.getName(), fd);
            }
            resolvedReferences.forEach((name, ps) -> {
                deps.put(ps.getParsedSchema().getFileDescriptor().getName(),
                         ps.getParsedSchema().getFileDescriptor());
            });

            Descriptors.FileDescriptor[] depArray = deps.values().toArray(new Descriptors.FileDescriptor[0]);
            Descriptors.FileDescriptor fd = Descriptors.FileDescriptor.buildFrom(fileDescriptorProto, depArray);
            return new ProtobufSchema(fd, protoFileElement);
        } catch (InvalidProtocolBufferException | Descriptors.DescriptorValidationException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * Parse rawSchema as FileDescriptorProto, trying raw binary first,
     * then falling back to base64-decoded binary.
     */
    private DescriptorProtos.FileDescriptorProto parseFileDescriptorProto(byte[] rawSchema)
            throws InvalidProtocolBufferException {
        try {
            // First try: parse as raw binary
            return DescriptorProtos.FileDescriptorProto.parseFrom(rawSchema);
        } catch (InvalidProtocolBufferException e) {
            // Second try: decode from base64 and parse
            try {
                String content = IoUtil.toString(rawSchema);
                byte[] decoded = Base64.getDecoder().decode(content);
                return DescriptorProtos.FileDescriptorProto.parseFrom(decoded);
            } catch (IllegalArgumentException base64Error) {
                // Not valid base64, throw original error
                throw e;
            }
        }
    }

    private ProtobufSchema getFileDescriptorFromElement(ProtoFileElement fileElem)
            throws DescriptorValidationException {
        FileDescriptor fileDescriptor = FileDescriptorUtils.protoFileToFileDescriptor(fileElem);
        return new ProtobufSchema(fileDescriptor, fileElem);
    }

    private void addReferencesToDependencies(List<ParsedSchema<ProtobufSchema>> schemaReferences,
            Map<String, ProtoFileElement> dependencies) {
        schemaReferences.forEach(parsedSchema -> {
            dependencies.put(parsedSchema.referenceName(),
                    parsedSchema.getParsedSchema().getProtoFileElement());
            if (parsedSchema.hasReferences()) {
                addReferencesToDependencies(parsedSchema.getSchemaReferences(), dependencies);
            }
        });
    }

    /**
     * @see io.apicurio.registry.resolver.SchemaParser#getSchemaFromData(Record)
     */
    @Override
    public ParsedSchema<ProtobufSchema> getSchemaFromData(Record<U> data) {
        FileDescriptor schemaFileDescriptor = data.payload().getDescriptorForType().getFile();
        ProtoFileElement protoFileElement = toProtoFileElement(schemaFileDescriptor);
        ProtobufSchema protobufSchema = new ProtobufSchema(schemaFileDescriptor, protoFileElement);

        byte[] rawSchema = IoUtil.toBytes(protoFileElement.toSchema());

        return new ParsedSchemaImpl<ProtobufSchema>().setParsedSchema(protobufSchema)
                .setReferenceName(protobufSchema.getFileDescriptor().getName())
                .setSchemaReferences(handleDependencies(schemaFileDescriptor)).setRawSchema(rawSchema);
    }

    @Override
    public ParsedSchema<ProtobufSchema> getSchemaFromData(Record<U> data, boolean dereference) {
        return getSchemaFromData(data);
    }

    private List<ParsedSchema<ProtobufSchema>> handleDependencies(FileDescriptor fileDescriptor) {
        List<ParsedSchema<ProtobufSchema>> schemaReferences = new ArrayList<>();
        fileDescriptor.getDependencies().forEach(referenceFileDescriptor -> {

            ProtoFileElement referenceProtoFileElement = toProtoFileElement(referenceFileDescriptor);
            ProtobufSchema referenceProtobufSchema = new ProtobufSchema(referenceFileDescriptor,
                    referenceProtoFileElement);

            byte[] rawSchema = IoUtil.toBytes(referenceProtoFileElement.toSchema());

            ParsedSchema<ProtobufSchema> referencedSchema = new ParsedSchemaImpl<ProtobufSchema>()
                    .setParsedSchema(referenceProtobufSchema)
                    .setReferenceName(referenceProtobufSchema.getFileDescriptor().getName())
                    .setSchemaReferences(handleDependencies(referenceFileDescriptor)).setRawSchema(rawSchema);
            schemaReferences.add(referencedSchema);
        });

        return schemaReferences;
    }

    /**
     * This method converts the Descriptor to a ProtoFileElement that allows to get a textual representation
     * .proto file
     *
     * @param fileDescriptor
     * @return textual protobuf representation
     */
    public ProtoFileElement toProtoFileElement(FileDescriptor fileDescriptor) {
        return FileDescriptorUtils.fileDescriptorToProtoFile(fileDescriptor.toProto());
    }
}
