package runtimegrpc

import (
	"context"

	"github.com/pingcap-inc/nexuscrm-contract/wire-go/envelope"
	commonv1 "github.com/pingcap-inc/nexuscrm-contract/wire-go/gen/common/v1"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

// Every Runtime connection observes rejection of its shared cached identity.
// This never retries a data RPC or changes its error. The next caller exchanges
// identity through the existing bounded cache and in-flight coalescing owner.
func runtimeClientConnection(target string, transport credentials.TransportCredentials, coordinator *identityExchangeCoordinator) (*grpc.ClientConn, error) {
	return grpc.NewClient(target, grpc.WithTransportCredentials(transport), grpc.WithIdleTimeout(0),
		grpc.WithUnaryInterceptor(func(ctx context.Context, method string, request, reply any, connection *grpc.ClientConn, invoke grpc.UnaryInvoker, options ...grpc.CallOption) error {
			err := invoke(ctx, method, request, reply, connection, options...)
			coordinator.observeRejection(ctx, err)
			return err
		}),
		grpc.WithStreamInterceptor(func(ctx context.Context, description *grpc.StreamDesc, connection *grpc.ClientConn, method string, stream grpc.Streamer, options ...grpc.CallOption) (grpc.ClientStream, error) {
			result, err := stream(ctx, description, connection, method, options...)
			if err != nil {
				coordinator.observeRejection(ctx, err)
				return nil, err
			}
			return &identityObservingStream{ClientStream: result, requestContext: ctx, coordinator: coordinator}, nil
		}))
}

func (coordinator *identityExchangeCoordinator) observeRejection(ctx context.Context, err error) {
	if coordinator == nil || coordinator.bursts == nil || err == nil {
		return
	}
	failure, ok := status.FromError(err)
	if !ok || failure.Code() != codes.FailedPrecondition {
		return
	}
	values, _ := metadata.FromOutgoingContext(ctx)
	delegations := values.Get(envelope.RuntimeDelegationMetadataKey)
	if len(delegations) != 1 || delegations[0] == "" {
		return
	}
	for _, raw := range failure.Details() {
		detail, ok := raw.(*commonv1.RuntimeErrorDetail)
		if !ok || detail.GetCategory() != commonv1.ErrorCategory_ERROR_CATEGORY_STALE_STATE {
			continue
		}
		stale := detail.GetStaleState()
		if stale.GetAxis() != commonv1.StaleAxis_STALE_AXIS_SECURITY_EPOCH || stale.GetServedRevision() == 0 || stale.GetRequiredRevision() <= stale.GetServedRevision() {
			continue
		}
		// Compare the exact rejected opaque delegation, not just the cache key:
		// a late response must not evict a replacement another caller acquired.
		cache := coordinator.bursts
		cache.mu.Lock()
		for key, entry := range cache.entries {
			if entry.response.OpaqueDelegation == delegations[0] {
				delete(cache.entries, key)
			}
		}
		cache.mu.Unlock()
		return
	}
}

type identityObservingStream struct {
	grpc.ClientStream
	requestContext context.Context
	coordinator    *identityExchangeCoordinator
}

func (stream *identityObservingStream) SendMsg(message any) error {
	err := stream.ClientStream.SendMsg(message)
	stream.coordinator.observeRejection(stream.requestContext, err)
	return err
}

func (stream *identityObservingStream) RecvMsg(message any) error {
	err := stream.ClientStream.RecvMsg(message)
	stream.coordinator.observeRejection(stream.requestContext, err)
	return err
}
