package client

import (
	"context"
	"google.golang.org/grpc"
	"sync"
)

type GrpcDirectClientFactory struct {
	defaultOpts []grpc.DialOption
	clientCache *sync.Map
}

func NewGrpcDirectClientFactory(defaultOpts ...grpc.DialOption) *GrpcDirectClientFactory {
	return &GrpcDirectClientFactory{
		defaultOpts: defaultOpts,
		clientCache: &sync.Map{},
	}
}

func (f *GrpcDirectClientFactory) NewConn(ctx context.Context, addr string, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
	dialOpts := make([]grpc.DialOption, 0, len(f.defaultOpts)+len(opts))
	dialOpts = append(dialOpts, f.defaultOpts...)
	dialOpts = append(dialOpts, opts...)

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

func (f *GrpcDirectClientFactory) GetConn(ctx context.Context, addr string, opts ...grpc.DialOption) (conn *grpc.ClientConn, err error) {
	val, ok := f.clientCache.Load(addr)
	if ok {
		conn = val.(*grpc.ClientConn)
		return
	}
	conn, err = f.NewConn(ctx, addr, opts...)
	if err != nil {
		return
	}
	f.clientCache.Store(addr, conn)
	return
}