package logid

import (
	"context"

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

// TODO(chenfei): move into pkgv2/grpc/middleware

// UnaryServerInterceptor returns a new unary server interceptors that set logid into ctx
func UnaryServerInterceptor() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
		// try to get requestid from ctx for compatible with grpc-gateway
		md, ok := metadata.FromIncomingContext(ctx)
		if ok {
			if vals := md.Get(LogidHTTPHeaderKey); len(vals) == 1 {
				logid := vals[0]
				ctx = WithLogid(ctx, logid)
			}
		}

		ctx = injectCtxInfoIntoCtx(ctx, req)
		return handler(ctx, req)
	}
}

// StreamServerInterceptor returns a new streaming server interceptors that set logid into ctx
func StreamServerInterceptor() grpc.StreamServerInterceptor {
	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		wrapped := grpc_middleware.WrapServerStream(stream)
		md, ok := metadata.FromIncomingContext(wrapped.WrappedContext)
		if ok {
			if vals := md.Get(LogidHTTPHeaderKey); len(vals) == 1 {
				logid := vals[0]
				wrapped.WrappedContext = WithLogid(wrapped.WrappedContext, logid)
			}
		}

		return handler(srv, wrapped)
	}
}

func UnaryClientInterceptor() grpc.UnaryClientInterceptor {
	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
		injectCtxInfoIntoReq(ctx, req)
		return invoker(ctx, method, req, reply, cc, opts...)
	}
}

func injectCtxInfoIntoReq(ctx context.Context, req interface{}) {
	// FIXME(iven): Move logid to message header
	// pbutil.SetTagIntoReq(req, GRPCHeaderKey, GetLogid(ctx))
}

func injectCtxInfoIntoCtx(ctx context.Context, req interface{}) context.Context {
	if req == nil {
		return ctx
	}

	// FIXME(iven): Move logid to message header
	// requestid := pbutil.GetTagsFromReq(req)[GRPCHeaderKey]
	// if requestid != "" {
	// 	ctx = WithLogid(ctx, requestid)
	// }
	return ctx
}