// Copyright (c) The OpenTofu Authors
// SPDX-License-Identifier: MPL-2.0
// Copyright (c) 2023 HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0

package plugin

import (
	"context"
	"errors"
	"fmt"

	plugin "github.com/hashicorp/go-plugin"
	"github.com/zclconf/go-cty/cty"
	ctyjson "github.com/zclconf/go-cty/cty/json"
	"github.com/zclconf/go-cty/cty/msgpack"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	"github.com/opentofu/opentofu/internal/logging"
	"github.com/opentofu/opentofu/internal/plugin/convert"
	"github.com/opentofu/opentofu/internal/plugin/validation"
	"github.com/opentofu/opentofu/internal/providers"
	"github.com/opentofu/opentofu/internal/tfdiags"
	proto "github.com/opentofu/opentofu/internal/tfplugin5"
)

var logger = logging.HCLogger()

// GRPCProviderPlugin implements plugin.GRPCPlugin for the go-plugin package.
type GRPCProviderPlugin struct {
	plugin.Plugin
	GRPCProvider func() proto.ProviderServer
}

var clientCapabilities = &proto.ClientCapabilities{
	// WriteOnlyAttributesAllowed indicates that the current system version
	// supports write-only attributes.
	// This enables the SDK to run specific validations and enable the
	// nullification of such configured attributes before returning the
	// response back to the system.
	WriteOnlyAttributesAllowed: true,
}

func (p *GRPCProviderPlugin) GRPCClient(ctx context.Context, broker *plugin.GRPCBroker, c *grpc.ClientConn) (any, error) {
	return &GRPCProvider{
		client: proto.NewProviderClient(c),
		ctx:    ctx,
	}, nil
}

func (p *GRPCProviderPlugin) GRPCServer(broker *plugin.GRPCBroker, s *grpc.Server) error {
	proto.RegisterProviderServer(s, p.GRPCProvider())
	return nil
}

// GRPCProvider handles the client, or core side of the plugin rpc connection.
// The GRPCProvider methods are mostly a translation layer between the
// tofu providers types and the grpc proto types, directly converting
// between the two.
type GRPCProvider struct {
	// PluginClient provides a reference to the plugin.Client which controls the plugin process.
	// This allows the GRPCProvider a way to shutdown the plugin process.
	PluginClient *plugin.Client

	// TestServer contains a grpc.Server to close when the GRPCProvider is being
	// used in an end to end test of a provider.
	TestServer *grpc.Server

	// Proto client use to make the grpc service calls.
	client proto.ProviderClient

	// this context is created by the plugin package, and is canceled when the
	// plugin process ends.
	//
	// THIS IS NOT THE RIGHT CONTEXT TO USE FOR GRPC CALLS! This represents
	// the overall context in which the plugin was launched and represents the
	// full runtime of the plugin process, whereas the [context.Context] passed
	// to individual [providers.Interface] methods is a more tightly-scoped
	// context focused on each individual request, and so is more appropriate
	// to use as the parent context for gRPC API calls.
	ctx context.Context

	// SchemaCache stores the schema for this provider. This is used to properly
	// serialize the requests for schemas.  This is shared between instances
	// of the provider.
	SchemaCache providers.SchemaCache

	// Keep track of if the proto schema fetch call has happend for GetProviderSchemaOptional
	// This allows caching to still function efficiently, without violating legacy provider's requirements
	hasFetchedSchema bool
}

var _ providers.Interface = new(GRPCProvider)

func (p *GRPCProvider) GetProviderSchema(ctx context.Context) (resp providers.GetProviderSchemaResponse) {
	logger.Trace("GRPCProvider: GetProviderSchema")

	schema := p.SchemaCache(func() providers.GetProviderSchemaResponse {
		return p.getProviderSchema(ctx)
	})

	if !p.hasFetchedSchema && !schema.ServerCapabilities.GetProviderSchemaOptional {
		// Force call, we only care that we are satisfying the legacy provider's schema call request
		_, _ = p.getProtoProviderSchema(ctx)
	}

	return schema
}

func (p *GRPCProvider) getProviderSchema(ctx context.Context) (resp providers.GetProviderSchemaResponse) {
	resp.ResourceTypes = make(map[string]providers.Schema)
	resp.DataSources = make(map[string]providers.Schema)
	resp.EphemeralResources = make(map[string]providers.Schema)
	resp.Functions = make(map[string]providers.FunctionSpec)

	protoResp, err := p.getProtoProviderSchema(ctx)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
		return resp
	}

	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))

	if resp.Diagnostics.HasErrors() {
		return resp
	}

	if protoResp.Provider == nil {
		resp.Diagnostics = resp.Diagnostics.Append(errors.New("missing provider schema"))
		return resp
	}

	// We want to allow "provider" blocks to work with ephemeral variables, so we
	// just mark its schema as able to get such values.
	resp.Provider = convert.ProtoToEphemeralProviderSchema(protoResp.Provider)
	if protoResp.ProviderMeta == nil {
		logger.Debug("No provider meta schema returned")
	} else {
		resp.ProviderMeta = convert.ProtoToProviderSchema(protoResp.ProviderMeta)
	}

	for name, res := range protoResp.ResourceSchemas {
		resp.ResourceTypes[name] = convert.ProtoToProviderSchema(res)
	}

	for name, data := range protoResp.DataSourceSchemas {
		resp.DataSources[name] = convert.ProtoToProviderSchema(data)
	}

	for name, data := range protoResp.EphemeralResourceSchemas {
		// Ephemeral resources should be able to work with ephemeral values by design.
		resp.EphemeralResources[name] = convert.ProtoToEphemeralProviderSchema(data)
	}

	for name, fn := range protoResp.Functions {
		resp.Functions[name] = convert.ProtoToFunctionSpec(fn)
	}

	identitySchemas, idsDiags := p.getResourceIdentitySchemas(ctx)
	if idsDiags.HasErrors() {
		// Identity schemas are an optional enhancement. A provider bug in
		// identity schemas should not prevent the provider from being used
		// for all other operations, so we log the error and continue
		// without identity schema support.
		logger.Warn("failed to fetch resource identity schemas, identity-based features will be unavailable", idsDiags.Err())
	} else {
		for name, idSchema := range identitySchemas {
			if resSchema, ok := resp.ResourceTypes[name]; ok {
				resSchema.IdentitySchema = idSchema.Body
				resSchema.IdentitySchemaVersion = idSchema.Version
				resp.ResourceTypes[name] = resSchema
			}
		}
	}

	if protoResp.ServerCapabilities != nil {
		resp.ServerCapabilities.PlanDestroy = protoResp.ServerCapabilities.PlanDestroy
		resp.ServerCapabilities.GetProviderSchemaOptional = protoResp.ServerCapabilities.GetProviderSchemaOptional
	}

	return resp
}

// Common code to fetch the raw schema data from the provider. This is called from
// multiple locations due to GetProviderSchemaOptional
func (p *GRPCProvider) getProtoProviderSchema(ctx context.Context) (*proto.GetProviderSchema_Response, error) {
	// Some providers may generate quite large schemas, and the internal default
	// grpc response size limit is 4MB. 64MB should cover most any use case, and
	// if we get providers nearing that we may want to consider a finer-grained
	// API to fetch individual resource schemas.
	// Note: this option is marked as EXPERIMENTAL in the grpc API. We keep
	// this for compatibility, but recent providers all set the max message
	// size much higher on the server side, which is the supported method for
	// determining payload size.
	const maxRecvSize = 64 << 20
	resp, err := p.client.GetSchema(ctx, new(proto.GetProviderSchema_Request), grpc.MaxRecvMsgSizeCallOption{MaxRecvMsgSize: maxRecvSize})

	// Mark that we have handled the internal requirement for legacy providers (!GetProviderSchemaOptional)
	p.hasFetchedSchema = true

	return resp, err
}

// getResourceIdentitySchemas should ONLY be called from GetProviderSchema, which
// merges the identity schemas into the cached provider schema response. All other
// callers should use resSchema.IdentitySchema from the GetProviderSchema response.
func (p *GRPCProvider) getResourceIdentitySchemas(ctx context.Context) (map[string]providers.ResourceIdentitySchema, tfdiags.Diagnostics) {
	logger.Trace("GRPCProvider: getResourceIdentitySchemas")
	var diags tfdiags.Diagnostics

	protoResponse, err := p.client.GetResourceIdentitySchemas(ctx, new(proto.GetResourceIdentitySchemas_Request))
	if err != nil {
		// Providers that don't implement identity schemas return Unimplemented;
		// this is expected and not an error — just return an empty response.
		if status.Code(err) == codes.Unimplemented {
			return nil, diags
		}
		diags = diags.Append(grpcErr(err))
		return nil, diags
	}

	diags = diags.Append(convert.ProtoToDiagnostics(protoResponse.Diagnostics))
	if diags.HasErrors() || len(protoResponse.IdentitySchemas) == 0 {
		return nil, diags
	}

	identitySchemas := make(map[string]providers.ResourceIdentitySchema, len(protoResponse.IdentitySchemas))
	for resourceType, schema := range protoResponse.IdentitySchemas {
		identitySchemas[resourceType] = *convert.ProtoToResourceIdentitySchema(schema)
	}
	return identitySchemas, diags
}

func (p *GRPCProvider) ValidateProviderConfig(ctx context.Context, r providers.ValidateProviderConfigRequest) (resp providers.ValidateProviderConfigResponse) {
	logger.Trace("GRPCProvider: ValidateProviderConfig")

	schema := p.GetProviderSchema(ctx)
	if schema.Diagnostics.HasErrors() {
		resp.Diagnostics = schema.Diagnostics
		return resp
	}

	ty := schema.Provider.Block.ImpliedType()

	mp, err := msgpack.Marshal(r.Config, ty)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}

	protoReq := &proto.PrepareProviderConfig_Request{
		Config: &proto.DynamicValue{Msgpack: mp},
	}

	protoResp, err := p.client.PrepareProviderConfig(ctx, protoReq)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
		return resp
	}

	config, err := decodeDynamicValue(protoResp.PreparedConfig, ty)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}
	resp.PreparedConfig = config

	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))
	return resp
}

func (p *GRPCProvider) ValidateResourceConfig(ctx context.Context, r providers.ValidateResourceConfigRequest) (resp providers.ValidateResourceConfigResponse) {
	logger.Trace("GRPCProvider: ValidateResourceConfig")

	schema := p.GetProviderSchema(ctx)
	if schema.Diagnostics.HasErrors() {
		resp.Diagnostics = schema.Diagnostics
		return resp
	}

	resourceSchema, ok := schema.ResourceTypes[r.TypeName]
	if !ok {
		resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("unknown resource type %q", r.TypeName))
		return resp
	}

	mp, err := msgpack.Marshal(r.Config, resourceSchema.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}

	protoReq := &proto.ValidateResourceTypeConfig_Request{
		TypeName:           r.TypeName,
		Config:             &proto.DynamicValue{Msgpack: mp},
		ClientCapabilities: clientCapabilities,
	}

	protoResp, err := p.client.ValidateResourceTypeConfig(ctx, protoReq)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
		return resp
	}

	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))
	return resp
}

func (p *GRPCProvider) ValidateDataResourceConfig(ctx context.Context, r providers.ValidateDataResourceConfigRequest) (resp providers.ValidateDataResourceConfigResponse) {
	logger.Trace("GRPCProvider: ValidateDataResourceConfig")

	schema := p.GetProviderSchema(ctx)
	if schema.Diagnostics.HasErrors() {
		resp.Diagnostics = schema.Diagnostics
		return resp
	}

	dataSchema, ok := schema.DataSources[r.TypeName]
	if !ok {
		resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("unknown data source %q", r.TypeName))
		return resp
	}

	mp, err := msgpack.Marshal(r.Config, dataSchema.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}

	protoReq := &proto.ValidateDataSourceConfig_Request{
		TypeName: r.TypeName,
		Config:   &proto.DynamicValue{Msgpack: mp},
	}

	protoResp, err := p.client.ValidateDataSourceConfig(ctx, protoReq)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
		return resp
	}
	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))
	return resp
}

func (p *GRPCProvider) ValidateEphemeralConfig(ctx context.Context, r providers.ValidateEphemeralConfigRequest) (resp providers.ValidateEphemeralConfigResponse) {
	logger.Trace("GRPCProvider: ValidateEphemeralConfig")

	schema := p.GetProviderSchema(ctx)
	if schema.Diagnostics.HasErrors() {
		resp.Diagnostics = schema.Diagnostics
		return resp
	}

	ephemeralSchema, ok := schema.EphemeralResources[r.TypeName]
	if !ok {
		resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("unknown ephemeral resource %q", r.TypeName))
		return resp
	}

	mp, err := msgpack.Marshal(r.Config, ephemeralSchema.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}

	protoReq := &proto.ValidateEphemeralResourceConfig_Request{
		TypeName: r.TypeName,
		Config:   &proto.DynamicValue{Msgpack: mp},
	}

	protoResp, err := p.client.ValidateEphemeralResourceConfig(ctx, protoReq)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
		return resp
	}
	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))
	return resp
}

func (p *GRPCProvider) UpgradeResourceState(ctx context.Context, r providers.UpgradeResourceStateRequest) (resp providers.UpgradeResourceStateResponse) {
	logger.Trace("GRPCProvider: UpgradeResourceState")

	schema := p.GetProviderSchema(ctx)
	if schema.Diagnostics.HasErrors() {
		resp.Diagnostics = schema.Diagnostics
		return resp
	}

	resSchema, ok := schema.ResourceTypes[r.TypeName]
	if !ok {
		resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("unknown resource type %q", r.TypeName))
		return resp
	}

	protoReq := &proto.UpgradeResourceState_Request{
		TypeName: r.TypeName,
		Version:  r.Version,
		RawState: &proto.RawState{
			Json:    r.RawStateJSON,
			Flatmap: r.RawStateFlatmap,
		},
	}

	protoResp, err := p.client.UpgradeResourceState(ctx, protoReq)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
		return resp
	}
	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))

	ty := resSchema.Block.ImpliedType()
	resp.UpgradedState = cty.NullVal(ty)
	if protoResp.UpgradedState == nil {
		return resp
	}

	state, err := decodeDynamicValue(protoResp.UpgradedState, ty)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}
	resp.UpgradedState = state

	resp.Diagnostics = resp.Diagnostics.Append(validation.WriteOnlyAttributes(resSchema.Block, resp.UpgradedState, r.TypeName))

	return resp
}

func (p *GRPCProvider) UpgradeResourceIdentity(ctx context.Context, req providers.UpgradeResourceIdentityRequest) (resp providers.UpgradeResourceIdentityResponse) {
	logger.Trace("GRPCProvider: UpgradeResourceIdentity")

	protoReq := &proto.UpgradeResourceIdentity_Request{
		TypeName: req.TypeName,
		Version:  req.Version,
		RawIdentity: &proto.RawState{
			Json: req.RawIdentityJSON,
		},
	}

	protoResp, err := p.client.UpgradeResourceIdentity(ctx, protoReq)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
		return resp
	}
	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))

	if protoResp.UpgradedIdentity == nil || protoResp.UpgradedIdentity.IdentityData == nil {
		if !resp.Diagnostics.HasErrors() {
			resp.Diagnostics = resp.Diagnostics.Append(tfdiags.Sourceless(
				tfdiags.Error,
				"Provider returned no identity after upgrade",
				fmt.Sprintf("The provider returned a nil identity for %q after UpgradeResourceIdentity without reporting an error. This is a bug in the provider.", req.TypeName),
			))
		}
		return resp
	}

	// Now we need to look up the resource identity schema from the provider resource schemas.
	providerSchema := p.GetProviderSchema(ctx)
	if providerSchema.Diagnostics.HasErrors() {
		resp.Diagnostics = resp.Diagnostics.Append(providerSchema.Diagnostics)
		return resp
	}
	resSchema, ok := providerSchema.ResourceTypes[req.TypeName]
	if !ok || resSchema.IdentitySchema == nil {
		resp.Diagnostics = resp.Diagnostics.Append(tfdiags.Sourceless(
			tfdiags.Error,
			"No identity schema found",
			fmt.Sprintf("No identity schema available for resource type %q.", req.TypeName),
		))
		return resp
	}

	identity, err := decodeDynamicValue(protoResp.UpgradedIdentity.IdentityData, resSchema.IdentitySchema.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}
	resp.UpgradedIdentity = identity

	return resp
}

func (p *GRPCProvider) ConfigureProvider(ctx context.Context, r providers.ConfigureProviderRequest) (resp providers.ConfigureProviderResponse) {
	logger.Trace("GRPCProvider: ConfigureProvider")

	schema := p.GetProviderSchema(ctx)
	if schema.Diagnostics.HasErrors() {
		resp.Diagnostics = schema.Diagnostics
		return resp
	}

	var mp []byte

	// we don't have anything to marshal if there's no config
	mp, err := msgpack.Marshal(r.Config, schema.Provider.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}

	protoReq := &proto.Configure_Request{
		TerraformVersion: r.TerraformVersion,
		Config: &proto.DynamicValue{
			Msgpack: mp,
		},
		ClientCapabilities: clientCapabilities,
	}

	protoResp, err := p.client.Configure(ctx, protoReq)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
		return resp
	}
	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))
	return resp
}

func (p *GRPCProvider) Stop(ctx context.Context) error {
	logger.Trace("GRPCProvider: Stop")
	// NOTE: The contract for providers.Interface guarantees that the ctx
	// passed to this function is never canceled, so we can safely use it
	// to make our "stop" request here. The context passed to other methods
	// _can_ be cancelled in some cases, so other methods should use
	// ctx = context.WithoutCancel(ctx) before making requests that need
	// to be able to terminate gracefully in response to "Stop".

	resp, err := p.client.Stop(ctx, new(proto.Stop_Request))
	if err != nil {
		return err
	}

	if resp.Error != "" {
		return errors.New(resp.Error)
	}
	return nil
}

func (p *GRPCProvider) ReadResource(ctx context.Context, r providers.ReadResourceRequest) (resp providers.ReadResourceResponse) {
	logger.Trace("GRPCProvider: ReadResource")

	schema := p.GetProviderSchema(ctx)
	if schema.Diagnostics.HasErrors() {
		resp.Diagnostics = schema.Diagnostics
		return resp
	}

	resSchema, ok := schema.ResourceTypes[r.TypeName]
	if !ok {
		resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("unknown resource type %s", r.TypeName))
		return resp
	}

	metaSchema := schema.ProviderMeta

	mp, err := msgpack.Marshal(r.PriorState, resSchema.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}

	protoReq := &proto.ReadResource_Request{
		TypeName:           r.TypeName,
		CurrentState:       &proto.DynamicValue{Msgpack: mp},
		Private:            r.Private,
		ClientCapabilities: clientCapabilities,
	}

	// Attach the identity if it is available
	if !r.PriorIdentity.IsNull() && resSchema.IdentitySchema != nil {
		identityMP, err := msgpack.Marshal(r.PriorIdentity, resSchema.IdentitySchema.ImpliedType())
		if err != nil {
			resp.Diagnostics = resp.Diagnostics.Append(err)
			return resp
		}
		protoReq.CurrentIdentity = &proto.ResourceIdentityData{
			IdentityData: &proto.DynamicValue{Msgpack: identityMP},
		}
	}

	if metaSchema.Block != nil {
		metaMP, err := msgpack.Marshal(r.ProviderMeta, metaSchema.Block.ImpliedType())
		if err != nil {
			resp.Diagnostics = resp.Diagnostics.Append(err)
			return resp
		}
		protoReq.ProviderMeta = &proto.DynamicValue{Msgpack: metaMP}
	}

	protoResp, err := p.client.ReadResource(ctx, protoReq)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
		return resp
	}
	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))

	state, err := decodeDynamicValue(protoResp.NewState, resSchema.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}
	resp.NewState = state
	resp.Private = protoResp.Private

	// Process identity from response
	if protoResp.NewIdentity != nil && protoResp.NewIdentity.IdentityData != nil && resSchema.IdentitySchema != nil {
		identity, err := decodeDynamicValue(protoResp.NewIdentity.IdentityData, resSchema.IdentitySchema.ImpliedType())
		if err != nil {
			resp.Diagnostics = resp.Diagnostics.Append(err)
			return resp
		}
		resp.NewIdentity = identity
	}

	resp.Diagnostics = resp.Diagnostics.Append(validation.WriteOnlyAttributes(resSchema.Block, resp.NewState, r.TypeName))

	return resp
}

func (p *GRPCProvider) PlanResourceChange(ctx context.Context, r providers.PlanResourceChangeRequest) (resp providers.PlanResourceChangeResponse) {
	logger.Trace("GRPCProvider: PlanResourceChange")

	schema := p.GetProviderSchema(ctx)
	if schema.Diagnostics.HasErrors() {
		resp.Diagnostics = schema.Diagnostics
		return resp
	}

	resSchema, ok := schema.ResourceTypes[r.TypeName]
	if !ok {
		resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("unknown resource type %q", r.TypeName))
		return resp
	}

	metaSchema := schema.ProviderMeta
	capabilities := schema.ServerCapabilities

	// If the provider doesn't support planning a destroy operation, we can
	// return immediately.
	if r.ProposedNewState.IsNull() && !capabilities.PlanDestroy {
		resp.PlannedState = r.ProposedNewState
		resp.PlannedPrivate = r.PriorPrivate
		return resp
	}

	priorMP, err := msgpack.Marshal(r.PriorState, resSchema.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}

	configMP, err := msgpack.Marshal(r.Config, resSchema.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}

	propMP, err := msgpack.Marshal(r.ProposedNewState, resSchema.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}

	protoReq := &proto.PlanResourceChange_Request{
		TypeName:           r.TypeName,
		PriorState:         &proto.DynamicValue{Msgpack: priorMP},
		Config:             &proto.DynamicValue{Msgpack: configMP},
		ProposedNewState:   &proto.DynamicValue{Msgpack: propMP},
		PriorPrivate:       r.PriorPrivate,
		ClientCapabilities: clientCapabilities,
	}

	if !r.PriorIdentity.IsNull() && resSchema.IdentitySchema != nil {
		identityMP, err := msgpack.Marshal(r.PriorIdentity, resSchema.IdentitySchema.ImpliedType())
		if err != nil {
			resp.Diagnostics = resp.Diagnostics.Append(err)
			return resp
		}
		protoReq.PriorIdentity = &proto.ResourceIdentityData{
			IdentityData: &proto.DynamicValue{Msgpack: identityMP},
		}
	}

	if metaSchema.Block != nil {
		metaMP, err := msgpack.Marshal(r.ProviderMeta, metaSchema.Block.ImpliedType())
		if err != nil {
			resp.Diagnostics = resp.Diagnostics.Append(err)
			return resp
		}
		protoReq.ProviderMeta = &proto.DynamicValue{Msgpack: metaMP}
	}

	protoResp, err := p.client.PlanResourceChange(ctx, protoReq)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
		return resp
	}
	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))

	if protoResp.PlannedIdentity != nil && protoResp.PlannedIdentity.IdentityData != nil && resSchema.IdentitySchema != nil {
		identity, err := decodeDynamicValue(protoResp.PlannedIdentity.IdentityData, resSchema.IdentitySchema.ImpliedType())
		if err != nil {
			resp.Diagnostics = resp.Diagnostics.Append(err)
			return resp
		}
		resp.PlannedIdentity = identity
	}

	state, err := decodeDynamicValue(protoResp.PlannedState, resSchema.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}
	resp.PlannedState = state

	for _, p := range protoResp.RequiresReplace {
		resp.RequiresReplace = append(resp.RequiresReplace, convert.AttributePathToPath(p))
	}

	resp.PlannedPrivate = protoResp.PlannedPrivate

	resp.LegacyTypeSystem = protoResp.LegacyTypeSystem

	resp.Diagnostics = resp.Diagnostics.Append(validation.WriteOnlyAttributes(resSchema.Block, resp.PlannedState, r.TypeName))

	return resp
}

func (p *GRPCProvider) ApplyResourceChange(ctx context.Context, r providers.ApplyResourceChangeRequest) (resp providers.ApplyResourceChangeResponse) {
	logger.Trace("GRPCProvider: ApplyResourceChange")

	schema := p.GetProviderSchema(ctx)
	if schema.Diagnostics.HasErrors() {
		resp.Diagnostics = schema.Diagnostics
		return resp
	}

	// Aside from fetching the schema above, the work of this function must
	// not be directly canceled by the incoming context's cancellation
	// signal, because we ask a provider plugin to gracefully cancel by
	// calling the Stop method and then its apply operation must be allowed
	// to run to completion to terminate gracefully if possible.
	ctx = context.WithoutCancel(ctx)

	resSchema, ok := schema.ResourceTypes[r.TypeName]
	if !ok {
		resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("unknown resource type %q", r.TypeName))
		return resp
	}

	metaSchema := schema.ProviderMeta

	priorMP, err := msgpack.Marshal(r.PriorState, resSchema.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}
	plannedMP, err := msgpack.Marshal(r.PlannedState, resSchema.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}
	configMP, err := msgpack.Marshal(r.Config, resSchema.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}

	protoReq := &proto.ApplyResourceChange_Request{
		TypeName:       r.TypeName,
		PriorState:     &proto.DynamicValue{Msgpack: priorMP},
		PlannedState:   &proto.DynamicValue{Msgpack: plannedMP},
		Config:         &proto.DynamicValue{Msgpack: configMP},
		PlannedPrivate: r.PlannedPrivate,
	}

	// Attach the proposed identity from the plan if it exists
	if !r.PlannedIdentity.IsNull() && resSchema.IdentitySchema != nil {
		identityMP, err := msgpack.Marshal(r.PlannedIdentity, resSchema.IdentitySchema.ImpliedType())
		if err != nil {
			resp.Diagnostics = resp.Diagnostics.Append(err)
			return resp
		}
		protoReq.PlannedIdentity = &proto.ResourceIdentityData{
			IdentityData: &proto.DynamicValue{Msgpack: identityMP},
		}
	}

	if metaSchema.Block != nil {
		metaMP, err := msgpack.Marshal(r.ProviderMeta, metaSchema.Block.ImpliedType())
		if err != nil {
			resp.Diagnostics = resp.Diagnostics.Append(err)
			return resp
		}
		protoReq.ProviderMeta = &proto.DynamicValue{Msgpack: metaMP}
	}

	protoResp, err := p.client.ApplyResourceChange(ctx, protoReq)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
		return resp
	}
	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))

	resp.Private = protoResp.Private

	if protoResp.NewIdentity != nil && protoResp.NewIdentity.IdentityData != nil && resSchema.IdentitySchema != nil {
		identity, err := decodeDynamicValue(protoResp.NewIdentity.IdentityData, resSchema.IdentitySchema.ImpliedType())
		if err != nil {
			resp.Diagnostics = resp.Diagnostics.Append(err)
			return resp
		}
		resp.NewIdentity = identity
	}

	state, err := decodeDynamicValue(protoResp.NewState, resSchema.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}
	resp.NewState = state

	resp.LegacyTypeSystem = protoResp.LegacyTypeSystem

	resp.Diagnostics = resp.Diagnostics.Append(validation.WriteOnlyAttributes(resSchema.Block, resp.NewState, r.TypeName))

	return resp
}

func (p *GRPCProvider) ImportResourceState(ctx context.Context, r providers.ImportResourceStateRequest) (resp providers.ImportResourceStateResponse) {
	logger.Trace("GRPCProvider: ImportResourceState")

	schema := p.GetProviderSchema(ctx)
	if schema.Diagnostics.HasErrors() {
		resp.Diagnostics = schema.Diagnostics
		return resp
	}

	protoReq := &proto.ImportResourceState_Request{
		TypeName:           r.TypeName,
		ClientCapabilities: clientCapabilities,
	}

	// If importing by identity, encode the identity value
	if r.Target.IsIdentityBased() {
		resSchema, ok := schema.ResourceTypes[r.TypeName]
		if !ok || resSchema.IdentitySchema == nil {
			resp.Diagnostics = resp.Diagnostics.Append(tfdiags.Sourceless(
				tfdiags.Error,
				"No identity schema found",
				fmt.Sprintf("No identity schema available for resource type %q.", r.TypeName),
			))
			return resp
		}

		identityMP, err := msgpack.Marshal(r.Target.Identity, resSchema.IdentitySchema.ImpliedType())
		if err != nil {
			resp.Diagnostics = resp.Diagnostics.Append(err)
			return resp
		}

		protoReq.Identity = &proto.ResourceIdentityData{
			IdentityData: &proto.DynamicValue{Msgpack: identityMP},
		}
	} else {
		// We're ID based
		protoReq.Id = r.Target.ID
	}

	protoResp, err := p.client.ImportResourceState(ctx, protoReq)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
		return resp
	}
	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))

	for _, imported := range protoResp.ImportedResources {
		resource := providers.ImportedResource{
			TypeName: imported.TypeName,
			Private:  imported.Private,
		}

		resSchema, ok := schema.ResourceTypes[r.TypeName]
		if !ok {
			resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("unknown resource type %q", r.TypeName))
			continue
		}

		if imported.Identity != nil && imported.Identity.IdentityData != nil {
			if resSchema.IdentitySchema != nil {
				identity, err := decodeDynamicValue(imported.Identity.IdentityData, resSchema.IdentitySchema.ImpliedType())
				if err != nil {
					resp.Diagnostics = resp.Diagnostics.Append(err)
					continue
				}
				resource.Identity = identity
			}
		}

		state, err := decodeDynamicValue(imported.State, resSchema.Block.ImpliedType())
		if err != nil {
			resp.Diagnostics = resp.Diagnostics.Append(err)
			return resp
		}
		resource.State = state
		resp.ImportedResources = append(resp.ImportedResources, resource)
	}

	return resp
}

func (p *GRPCProvider) MoveResourceState(ctx context.Context, r providers.MoveResourceStateRequest) providers.MoveResourceStateResponse {
	var resp providers.MoveResourceStateResponse
	logger.Trace("GRPCProvider: MoveResourceState")

	schema := p.GetProviderSchema(ctx)
	if schema.Diagnostics.HasErrors() {
		resp.Diagnostics = schema.Diagnostics
		return resp
	}

	resourceSchema, ok := schema.ResourceTypes[r.TargetTypeName]
	if !ok {
		schema.Diagnostics = schema.Diagnostics.Append(fmt.Errorf("unknown data source %q", r.TargetTypeName))
		return resp
	}

	protoReq := &proto.MoveResourceState_Request{
		SourceProviderAddress: r.SourceProviderAddress,
		SourceTypeName:        r.SourceTypeName,
		SourceSchemaVersion:   int64(r.SourceSchemaVersion),
		SourceState: &proto.RawState{
			Json:    r.SourceStateJSON,
			Flatmap: r.SourceStateFlatmap,
		},
		SourcePrivate:  r.SourcePrivate,
		TargetTypeName: r.TargetTypeName,
	}

	protoResp, err := p.client.MoveResourceState(ctx, protoReq)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
		return resp
	}
	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))

	state, err := decodeDynamicValue(protoResp.TargetState, resourceSchema.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}
	resp.TargetState = state
	resp.TargetPrivate = protoResp.TargetPrivate

	resp.Diagnostics = resp.Diagnostics.Append(validation.WriteOnlyAttributes(resourceSchema.Block, resp.TargetState, r.TargetTypeName))

	return resp
}

func (p *GRPCProvider) ReadDataSource(ctx context.Context, r providers.ReadDataSourceRequest) (resp providers.ReadDataSourceResponse) {
	logger.Trace("GRPCProvider: ReadDataSource")

	schema := p.GetProviderSchema(ctx)
	if schema.Diagnostics.HasErrors() {
		resp.Diagnostics = schema.Diagnostics
		return resp
	}

	dataSchema, ok := schema.DataSources[r.TypeName]
	if !ok {
		schema.Diagnostics = schema.Diagnostics.Append(fmt.Errorf("unknown data source %q", r.TypeName))
	}

	metaSchema := schema.ProviderMeta

	config, err := msgpack.Marshal(r.Config, dataSchema.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}

	protoReq := &proto.ReadDataSource_Request{
		TypeName: r.TypeName,
		Config: &proto.DynamicValue{
			Msgpack: config,
		},
		ClientCapabilities: clientCapabilities,
	}

	if metaSchema.Block != nil {
		metaMP, err := msgpack.Marshal(r.ProviderMeta, metaSchema.Block.ImpliedType())
		if err != nil {
			resp.Diagnostics = resp.Diagnostics.Append(err)
			return resp
		}
		protoReq.ProviderMeta = &proto.DynamicValue{Msgpack: metaMP}
	}

	protoResp, err := p.client.ReadDataSource(ctx, protoReq)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
		return resp
	}
	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))

	state, err := decodeDynamicValue(protoResp.State, dataSchema.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}
	resp.State = state

	return resp
}

func (p *GRPCProvider) OpenEphemeralResource(ctx context.Context, r providers.OpenEphemeralResourceRequest) (resp providers.OpenEphemeralResourceResponse) {
	logger.Trace("GRPCProvider: OpenEphemeralResource")

	schema := p.GetProviderSchema(ctx)
	resp.Diagnostics = schema.Diagnostics
	if resp.Diagnostics.HasErrors() {
		return resp
	}

	ephemeralResourceSchema, ok := schema.EphemeralResources[r.TypeName]
	if !ok {
		resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("unknown ephemeral resource %q", r.TypeName))
		return resp
	}

	config, err := msgpack.Marshal(r.Config, ephemeralResourceSchema.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}

	protoReq := &proto.OpenEphemeralResource_Request{
		TypeName: r.TypeName,
		Config: &proto.DynamicValue{
			Msgpack: config,
		},
		ClientCapabilities: clientCapabilities,
	}

	protoResp, err := p.client.OpenEphemeralResource(ctx, protoReq)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
		return resp
	}
	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))

	result, err := decodeDynamicValue(protoResp.Result, ephemeralResourceSchema.Block.ImpliedType())
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(err)
		return resp
	}
	resp.Result = result
	resp.Private = protoResp.Private
	if protoResp.RenewAt != nil {
		renewAt := protoResp.RenewAt.AsTime()
		resp.RenewAt = &renewAt
	}

	return resp
}

func (p *GRPCProvider) RenewEphemeralResource(ctx context.Context, r providers.RenewEphemeralResourceRequest) (resp providers.RenewEphemeralResourceResponse) {
	logger.Trace("GRPCProvider: RenewEphemeralResource")

	protoReq := &proto.RenewEphemeralResource_Request{
		TypeName: r.TypeName,
		Private:  r.Private,
	}

	protoResp, err := p.client.RenewEphemeralResource(ctx, protoReq)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
		return resp
	}
	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))
	resp.Private = protoResp.Private
	if protoResp.RenewAt != nil {
		renewAt := protoResp.RenewAt.AsTime()
		resp.RenewAt = &renewAt
	}

	return resp
}

func (p *GRPCProvider) CloseEphemeralResource(ctx context.Context, r providers.CloseEphemeralResourceRequest) (resp providers.CloseEphemeralResourceResponse) {
	logger.Trace("GRPCProvider: CloseEphemeralResource")

	protoReq := &proto.CloseEphemeralResource_Request{
		TypeName: r.TypeName,
		Private:  r.Private,
	}

	protoResp, err := p.client.CloseEphemeralResource(ctx, protoReq)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
		return resp
	}
	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))

	return resp
}

func (p *GRPCProvider) GetFunctions(ctx context.Context) (resp providers.GetFunctionsResponse) {
	logger.Trace("GRPCProvider: GetFunctions")

	protoReq := &proto.GetFunctions_Request{}

	protoResp, err := p.client.GetFunctions(ctx, protoReq)
	if err != nil {
		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
		return resp
	}
	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))
	resp.Functions = make(map[string]providers.FunctionSpec)

	for name, fn := range protoResp.Functions {
		resp.Functions[name] = convert.ProtoToFunctionSpec(fn)
	}

	return resp
}

func (p *GRPCProvider) CallFunction(ctx context.Context, r providers.CallFunctionRequest) (resp providers.CallFunctionResponse) {
	logger.Trace("GRPCProvider: CallFunction")

	schema := p.GetProviderSchema(ctx)
	if schema.Diagnostics.HasErrors() {
		// This should be unreachable
		resp.Error = schema.Diagnostics.Err()
		return resp
	}

	spec, ok := schema.Functions[r.Name]
	if !ok {
		funcs := p.GetFunctions(ctx)
		if funcs.Diagnostics.HasErrors() {
			// This should be unreachable
			resp.Error = funcs.Diagnostics.Err()
			return resp
		}
		spec, ok = funcs.Functions[r.Name]
		if !ok {
			// This should be unreachable
			resp.Error = fmt.Errorf("invalid CallFunctionRequest: function %s not defined in provider schema", r.Name)
			return resp
		}
	}

	protoReq := &proto.CallFunction_Request{
		Name:      r.Name,
		Arguments: make([]*proto.DynamicValue, len(r.Arguments)),
	}

	// Translate the arguments
	// As this is functionality is always sitting behind cty/function.Function, we skip some validation
	// checks of from the function and param spec.  We still include basic validation to prevent panics,
	// just in case there are bugs in cty.  See context_functions_test.go for explicit testing of argument
	// handling and short-circuiting.
	if len(r.Arguments) < len(spec.Parameters) {
		// This should be unreachable
		resp.Error = fmt.Errorf("invalid CallFunctionRequest: function %s expected %d parameters and got %d instead", r.Name, len(spec.Parameters), len(r.Arguments))
		return resp
	}

	for i, arg := range r.Arguments {
		var paramSpec providers.FunctionParameterSpec
		if i < len(spec.Parameters) {
			paramSpec = spec.Parameters[i]
		} else {
			// We are past the end of spec.Parameters, this is either variadic or an error
			if spec.VariadicParameter != nil {
				paramSpec = *spec.VariadicParameter
			} else {
				// This should be unreachable
				resp.Error = fmt.Errorf("invalid CallFunctionRequest: too many arguments passed to non-variadic function %s", r.Name)
				return resp
			}
		}

		if !paramSpec.AllowUnknownValues && !arg.IsWhollyKnown() {
			// Unlike the standard in cty, AllowUnknownValues == false does not just apply to
			// the root of the value (IsKnown) and instead also applies to values inside collections
			// and structures (IsWhollyKnown).
			// This is documented in the tfplugin proto file comments.
			//
			// The standard cty logic can be found in:
			// https://github.com/zclconf/go-cty/blob/ea922e7a95ba2be57897697117f318670e066d22/cty/function/function.go#L288-L290
			resp.Result = cty.UnknownVal(spec.Return)
			return resp
		}

		if arg.IsNull() {
			if paramSpec.AllowNullValue {
				continue
			} else {
				resp.Error = &providers.CallFunctionArgumentError{
					Text:             fmt.Sprintf("parameter %s is null, which is not allowed for function %s", paramSpec.Name, r.Name),
					FunctionArgument: i,
				}
			}
		}

		encodedArg, err := msgpack.Marshal(arg, paramSpec.Type)
		if err != nil {
			resp.Error = err
			return resp
		}

		protoReq.Arguments[i] = &proto.DynamicValue{
			Msgpack: encodedArg,
		}
	}

	protoResp, err := p.client.CallFunction(ctx, protoReq)
	if err != nil {
		resp.Error = err
		return resp
	}

	if protoResp.Error != nil {
		err := &providers.CallFunctionArgumentError{
			Text: protoResp.Error.Text,
		}
		if protoResp.Error.FunctionArgument != nil {
			err.FunctionArgument = int(*protoResp.Error.FunctionArgument)
		}
		resp.Error = err
		return resp
	}

	resp.Result, resp.Error = decodeDynamicValue(protoResp.Result, spec.Return)
	return resp
}

// closing the grpc connection is final, and tofu will call it at the end of every phase.
func (p *GRPCProvider) Close(ctx context.Context) error {
	logger.Trace("GRPCProvider: Close")

	// Make sure to stop the server if we're not running within go-plugin.
	if p.TestServer != nil {
		p.TestServer.Stop()
	}

	// Check this since it's not automatically inserted during plugin creation.
	// It's currently only inserted by the command package, because that is
	// where the factory is built and is the only point with access to the
	// plugin.Client.
	if p.PluginClient == nil {
		logger.Debug("provider has no plugin.Client")
		return nil
	}

	p.PluginClient.Kill()
	return nil
}

// Decode a DynamicValue from either the JSON or MsgPack encoding.
func decodeDynamicValue(v *proto.DynamicValue, ty cty.Type) (cty.Value, error) {
	// always return a valid value
	var err error
	res := cty.NullVal(ty)
	if v == nil {
		return res, nil
	}

	switch {
	case len(v.Msgpack) > 0:
		res, err = msgpack.Unmarshal(v.Msgpack, ty)
	case len(v.Json) > 0:
		res, err = ctyjson.Unmarshal(v.Json, ty)
	}
	return res, err
}
