package client

// this package wraps some utilities for grpc server and client

import (
	"context"

	grpcZap "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap"
	"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
	grpcPrometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/pingcap/log"
	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
	"go.opentelemetry.io/otel/trace"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials/insecure"
)

type client struct {
	unaryInterceptors  []grpc.UnaryClientInterceptor
	streamInterceptors []grpc.StreamClientInterceptor
	grpcDialOptions    []grpc.DialOption

	loadBalance bool
}

type Option func(c *client)

func WithPrometheus() Option {
	return func(c *client) {
		grpcPrometheus.EnableClientHandlingTimeHistogram()
		c.unaryInterceptors = append(c.unaryInterceptors, grpcPrometheus.UnaryClientInterceptor)
		c.streamInterceptors = append(c.streamInterceptors, grpcPrometheus.StreamClientInterceptor)
	}
}

func WithTracing(opts ...otelgrpc.Option) Option {
	return func(c *client) {
		c.unaryInterceptors = append(c.unaryInterceptors, otelgrpc.UnaryClientInterceptor(opts...))
		c.streamInterceptors = append(c.streamInterceptors, otelgrpc.StreamClientInterceptor(opts...))
	}
}

func TraceIDMessageProducer(ctx context.Context, msg string, level zapcore.Level, code codes.Code, err error, duration zapcore.Field) {
	// re-extract logger from newCtx, as it may have extra fields that changed in the holder.
	spanCtx := trace.SpanContextFromContext(ctx)
	if spanCtx.HasTraceID() {
		ctxzap.Extract(ctx).Check(level, msg).Write(
			zap.Error(err),
			zap.String("grpc.code", code.String()),
			duration,
			zap.String("trace_id", spanCtx.TraceID().String()),
		)
	} else {
		ctxzap.Extract(ctx).Check(level, msg).Write(
			zap.Error(err),
			zap.String("grpc.code", code.String()),
			duration,
		)
	}
}

func WithZapLogger(logger *zap.Logger, opts ...grpcZap.Option) Option {
	return func(c *client) {
		c.unaryInterceptors = append(c.unaryInterceptors, grpcZap.UnaryClientInterceptor(logger, opts...))
		c.streamInterceptors = append(c.streamInterceptors, grpcZap.StreamClientInterceptor(logger, opts...))
	}
}

func WithUnaryInterceptors(interceptors ...grpc.UnaryClientInterceptor) Option {
	return func(c *client) {
		c.unaryInterceptors = append(c.unaryInterceptors, interceptors...)
	}
}

func WithStreamInterceptors(interceptors ...grpc.StreamClientInterceptor) Option {
	return func(c *client) {
		c.streamInterceptors = append(c.streamInterceptors, interceptors...)
	}
}

func WithLoadBalance() Option {
	return func(c *client) {
		c.loadBalance = true
	}
}

func WithDefault() Option {
	return func(c *client) {
		zapOpts := grpcZap.WithMessageProducer(TraceIDMessageProducer)
		options := []Option{WithPrometheus(), WithTracing(), WithZapLogger(log.L(), zapOpts)}

		for _, opt := range options {
			opt(c)
		}
	}
}

func WithGrpcOption(opt grpc.DialOption) Option {
	return func(c *client) {
		c.grpcDialOptions = append(c.grpcDialOptions, opt)
	}
}

func NewClientConn(addr string, opts ...Option) (*grpc.ClientConn, error) {
	c := &client{}
	for _, opt := range opts {
		opt(c)
	}

	dialOpts := []grpc.DialOption{
		grpc.WithTransportCredentials(insecure.NewCredentials()), // don't support TLS connection yet
		grpc.WithChainUnaryInterceptor(c.unaryInterceptors...),
		grpc.WithChainStreamInterceptor(c.streamInterceptors...),
	}
	dialOpts = append(dialOpts, c.grpcDialOptions...)

	if c.loadBalance {
		dialOpts = append(dialOpts, grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`))
	}

	return grpc.Dial(addr, dialOpts...)
}

func NewClientContext(ctx context.Context, addr string, opts ...Option) (*grpc.ClientConn, error) {
	c := &client{}
	for _, opt := range opts {
		opt(c)
	}

	dialOpts := []grpc.DialOption{
		grpc.WithTransportCredentials(insecure.NewCredentials()), // don't support TLS connection yet
		grpc.WithChainUnaryInterceptor(c.unaryInterceptors...),
		grpc.WithChainStreamInterceptor(c.streamInterceptors...),
	}
	dialOpts = append(dialOpts, c.grpcDialOptions...)

	if c.loadBalance {
		dialOpts = append(dialOpts, grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`))
	}

	return grpc.DialContext(ctx, addr, dialOpts...)
}