// Copyright 2020-2024 Buf Technologies, Inc.
//
// 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
//
//      http://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.

// Package protoutil contains useful functions for interacting with descriptors.
// For now these include only functions for efficiently converting descriptors
// produced by the compiler to descriptor protos and functions for resolving
// "features" (a core concept of Protobuf Editions).
//
// Despite the fact that descriptor protos are mutable, calling code should NOT
// mutate any of the protos returned from this package. For efficiency, some
// values returned from this package may reference internal state of a compiler
// result, and mutating the proto could corrupt or invalidate parts of that
// result.
package protoutil

import (
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/reflect/protodesc"
	"google.golang.org/protobuf/reflect/protoreflect"
	"google.golang.org/protobuf/types/descriptorpb"
)

// DescriptorProtoWrapper is a protoreflect.Descriptor that wraps an
// underlying descriptor proto. It provides the same interface as
// Descriptor but with one extra operation, to efficiently query for
// the underlying descriptor proto.
//
// Descriptors that implement this will also implement another method
// whose specified return type is the concrete type returned by the
// AsProto method. The name of this method varies by the type of this
// descriptor:
//
//	 Descriptor Type        Other Method Name
//	---------------------+------------------------------------
//	 FileDescriptor      |  FileDescriptorProto()
//	 MessageDescriptor   |  MessageDescriptorProto()
//	 FieldDescriptor     |  FieldDescriptorProto()
//	 OneofDescriptor     |  OneofDescriptorProto()
//	 EnumDescriptor      |  EnumDescriptorProto()
//	 EnumValueDescriptor |  EnumValueDescriptorProto()
//	 ServiceDescriptor   |  ServiceDescriptorProto()
//	 MethodDescriptor    |  MethodDescriptorProto()
//
// For example, a DescriptorProtoWrapper that implements FileDescriptor
// returns a *descriptorpb.FileDescriptorProto value from its AsProto
// method and also provides a method with the following signature:
//
//	FileDescriptorProto() *descriptorpb.FileDescriptorProto
type DescriptorProtoWrapper interface {
	protoreflect.Descriptor
	// AsProto returns the underlying descriptor proto. The concrete
	// type of the proto message depends on the type of this
	// descriptor:
	//    Descriptor Type        Proto Message Type
	//   ---------------------+------------------------------------
	//    FileDescriptor      |  *descriptorpb.FileDescriptorProto
	//    MessageDescriptor   |  *descriptorpb.DescriptorProto
	//    FieldDescriptor     |  *descriptorpb.FieldDescriptorProto
	//    OneofDescriptor     |  *descriptorpb.OneofDescriptorProto
	//    EnumDescriptor      |  *descriptorpb.EnumDescriptorProto
	//    EnumValueDescriptor |  *descriptorpb.EnumValueDescriptorProto
	//    ServiceDescriptor   |  *descriptorpb.ServiceDescriptorProto
	//    MethodDescriptor    |  *descriptorpb.MethodDescriptorProto
	AsProto() proto.Message
}

// ProtoFromDescriptor extracts a descriptor proto from the given "rich"
// descriptor. For descriptors generated by the compiler, this is an
// inexpensive and non-lossy operation. Descriptors from other sources
// however may be expensive (to re-create a proto) and even lossy.
func ProtoFromDescriptor(d protoreflect.Descriptor) proto.Message {
	switch d := d.(type) {
	case protoreflect.FileDescriptor:
		return ProtoFromFileDescriptor(d)
	case protoreflect.MessageDescriptor:
		return ProtoFromMessageDescriptor(d)
	case protoreflect.FieldDescriptor:
		return ProtoFromFieldDescriptor(d)
	case protoreflect.OneofDescriptor:
		return ProtoFromOneofDescriptor(d)
	case protoreflect.EnumDescriptor:
		return ProtoFromEnumDescriptor(d)
	case protoreflect.EnumValueDescriptor:
		return ProtoFromEnumValueDescriptor(d)
	case protoreflect.ServiceDescriptor:
		return ProtoFromServiceDescriptor(d)
	case protoreflect.MethodDescriptor:
		return ProtoFromMethodDescriptor(d)
	default:
		// WTF??
		if res, ok := d.(DescriptorProtoWrapper); ok {
			return res.AsProto()
		}
		return nil
	}
}

// ProtoFromFileDescriptor extracts a descriptor proto from the given "rich"
// descriptor. For file descriptors generated by the compiler, this is an
// inexpensive and non-lossy operation. File descriptors from other sources
// however may be expensive (to re-create a proto) and even lossy.
func ProtoFromFileDescriptor(d protoreflect.FileDescriptor) *descriptorpb.FileDescriptorProto {
	if imp, ok := d.(protoreflect.FileImport); ok {
		d = imp.FileDescriptor
	}
	type canProto interface {
		FileDescriptorProto() *descriptorpb.FileDescriptorProto
	}
	if res, ok := d.(canProto); ok {
		return res.FileDescriptorProto()
	}
	if res, ok := d.(DescriptorProtoWrapper); ok {
		if fd, ok := res.AsProto().(*descriptorpb.FileDescriptorProto); ok {
			return fd
		}
	}
	return protodesc.ToFileDescriptorProto(d)
}

// ProtoFromMessageDescriptor extracts a descriptor proto from the given "rich"
// descriptor. For message descriptors generated by the compiler, this is an
// inexpensive and non-lossy operation. Message descriptors from other sources
// however may be expensive (to re-create a proto) and even lossy.
func ProtoFromMessageDescriptor(d protoreflect.MessageDescriptor) *descriptorpb.DescriptorProto {
	type canProto interface {
		MessageDescriptorProto() *descriptorpb.DescriptorProto
	}
	if res, ok := d.(canProto); ok {
		return res.MessageDescriptorProto()
	}
	if res, ok := d.(DescriptorProtoWrapper); ok {
		if md, ok := res.AsProto().(*descriptorpb.DescriptorProto); ok {
			return md
		}
	}
	return protodesc.ToDescriptorProto(d)
}

// ProtoFromFieldDescriptor extracts a descriptor proto from the given "rich"
// descriptor. For field descriptors generated by the compiler, this is an
// inexpensive and non-lossy operation. Field descriptors from other sources
// however may be expensive (to re-create a proto) and even lossy.
func ProtoFromFieldDescriptor(d protoreflect.FieldDescriptor) *descriptorpb.FieldDescriptorProto {
	type canProto interface {
		FieldDescriptorProto() *descriptorpb.FieldDescriptorProto
	}
	if res, ok := d.(canProto); ok {
		return res.FieldDescriptorProto()
	}
	if res, ok := d.(DescriptorProtoWrapper); ok {
		if fd, ok := res.AsProto().(*descriptorpb.FieldDescriptorProto); ok {
			return fd
		}
	}
	return protodesc.ToFieldDescriptorProto(d)
}

// ProtoFromOneofDescriptor extracts a descriptor proto from the given "rich"
// descriptor. For oneof descriptors generated by the compiler, this is an
// inexpensive and non-lossy operation. Oneof descriptors from other sources
// however may be expensive (to re-create a proto) and even lossy.
func ProtoFromOneofDescriptor(d protoreflect.OneofDescriptor) *descriptorpb.OneofDescriptorProto {
	type canProto interface {
		OneofDescriptorProto() *descriptorpb.OneofDescriptorProto
	}
	if res, ok := d.(canProto); ok {
		return res.OneofDescriptorProto()
	}
	if res, ok := d.(DescriptorProtoWrapper); ok {
		if ood, ok := res.AsProto().(*descriptorpb.OneofDescriptorProto); ok {
			return ood
		}
	}
	return protodesc.ToOneofDescriptorProto(d)
}

// ProtoFromEnumDescriptor extracts a descriptor proto from the given "rich"
// descriptor. For enum descriptors generated by the compiler, this is an
// inexpensive and non-lossy operation. Enum descriptors from other sources
// however may be expensive (to re-create a proto) and even lossy.
func ProtoFromEnumDescriptor(d protoreflect.EnumDescriptor) *descriptorpb.EnumDescriptorProto {
	type canProto interface {
		EnumDescriptorProto() *descriptorpb.EnumDescriptorProto
	}
	if res, ok := d.(canProto); ok {
		return res.EnumDescriptorProto()
	}
	if res, ok := d.(DescriptorProtoWrapper); ok {
		if ed, ok := res.AsProto().(*descriptorpb.EnumDescriptorProto); ok {
			return ed
		}
	}
	return protodesc.ToEnumDescriptorProto(d)
}

// ProtoFromEnumValueDescriptor extracts a descriptor proto from the given "rich"
// descriptor. For enum value descriptors generated by the compiler, this is an
// inexpensive and non-lossy operation. Enum value descriptors from other sources
// however may be expensive (to re-create a proto) and even lossy.
func ProtoFromEnumValueDescriptor(d protoreflect.EnumValueDescriptor) *descriptorpb.EnumValueDescriptorProto {
	type canProto interface {
		EnumValueDescriptorProto() *descriptorpb.EnumValueDescriptorProto
	}
	if res, ok := d.(canProto); ok {
		return res.EnumValueDescriptorProto()
	}
	if res, ok := d.(DescriptorProtoWrapper); ok {
		if ed, ok := res.AsProto().(*descriptorpb.EnumValueDescriptorProto); ok {
			return ed
		}
	}
	return protodesc.ToEnumValueDescriptorProto(d)
}

// ProtoFromServiceDescriptor extracts a descriptor proto from the given "rich"
// descriptor. For service descriptors generated by the compiler, this is an
// inexpensive and non-lossy operation. Service descriptors from other sources
// however may be expensive (to re-create a proto) and even lossy.
func ProtoFromServiceDescriptor(d protoreflect.ServiceDescriptor) *descriptorpb.ServiceDescriptorProto {
	type canProto interface {
		ServiceDescriptorProto() *descriptorpb.ServiceDescriptorProto
	}
	if res, ok := d.(canProto); ok {
		return res.ServiceDescriptorProto()
	}
	if res, ok := d.(DescriptorProtoWrapper); ok {
		if sd, ok := res.AsProto().(*descriptorpb.ServiceDescriptorProto); ok {
			return sd
		}
	}
	return protodesc.ToServiceDescriptorProto(d)
}

// ProtoFromMethodDescriptor extracts a descriptor proto from the given "rich"
// descriptor. For method descriptors generated by the compiler, this is an
// inexpensive and non-lossy operation. Method descriptors from other sources
// however may be expensive (to re-create a proto) and even lossy.
func ProtoFromMethodDescriptor(d protoreflect.MethodDescriptor) *descriptorpb.MethodDescriptorProto {
	type canProto interface {
		MethodDescriptorProto() *descriptorpb.MethodDescriptorProto
	}
	if res, ok := d.(canProto); ok {
		return res.MethodDescriptorProto()
	}
	if res, ok := d.(DescriptorProtoWrapper); ok {
		if md, ok := res.AsProto().(*descriptorpb.MethodDescriptorProto); ok {
			return md
		}
	}
	return protodesc.ToMethodDescriptorProto(d)
}