package modshared

import (
	"context"

	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware/v2"
	"google.golang.org/grpc"
)

type rpcAPIKeyType int

const (
	rpcAPIKey rpcAPIKeyType = iota
)

type RPCAPIFactory func(ctx context.Context, method string) RPCAPI

func RPCAPIFromContext[T RPCAPI](ctx context.Context) T {
	rpcAPI, ok := ctx.Value(rpcAPIKey).(T)
	if !ok {
		// This is a programmer error, so panic.
		panic("RPCAPI not attached to context. Make sure you are using interceptors")
	}
	return rpcAPI
}

func InjectRPCAPI(ctx context.Context, rpcAPI RPCAPI) context.Context {
	return context.WithValue(ctx, rpcAPIKey, rpcAPI)
}

// UnaryRPCAPIInterceptor returns a new unary server interceptor that augments connection context with a RPCAPI.
func UnaryRPCAPIInterceptor(factory RPCAPIFactory) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
		return handler(augmentContextWithRPCAPI(ctx, info.FullMethod, factory), req)
	}
}

// StreamRPCAPIInterceptor returns a new stream server interceptor that augments connection context with a RPCAPI.
func StreamRPCAPIInterceptor(factory RPCAPIFactory) grpc.StreamServerInterceptor {
	return func(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		wrapper := grpc_middleware.WrapServerStream(ss)
		wrapper.WrappedContext = augmentContextWithRPCAPI(wrapper.WrappedContext, info.FullMethod, factory)
		return handler(srv, wrapper)
	}
}

func augmentContextWithRPCAPI(ctx context.Context, method string, factory RPCAPIFactory) context.Context {
	return InjectRPCAPI(ctx, factory(ctx, method))
}
