package common

import (
	"context"
	"strings"
	"time"

	"braces.dev/errtrace"
	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
	"github.com/sercand/kuberesolver/v5"
	tigrpc "github.com/tidbcloud/tiinf/transport/tigrpc/client"
	"google.golang.org/grpc"
	"google.golang.org/protobuf/encoding/protojson"

	"github.com/tidbcloud/tidb-management-service/tidb-mgmt-portal/pkg/apimiddleware/grpcerror"
	utils "github.com/tidbcloud/tidb-management-service/utilities/grpc"
)

func CreateGrpcConn(address string, timeout time.Duration) (*grpc.ClientConn, error) {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	opts := []tigrpc.Option{
		tigrpc.WithPrometheus(),
		tigrpc.WithTracing(),
		tigrpc.WithUnaryInterceptors(utils.ClientPayloadLogInterceptor, utils.TimeoutInterceptor(timeout)),
	}
	if strings.HasPrefix(address, "kubernetes://") {
		kuberesolver.RegisterInCluster()
		opts = append(opts, tigrpc.WithLoadBalance())
	}

	conn, err := tigrpc.NewClientContext(ctx, address, opts...)
	if err != nil {
		return nil, errtrace.Wrap(err)
	}

	return conn, nil
}

// MakeGatewayMux returns a mux for a GRPC gateway
func MakeGatewayMux() *runtime.ServeMux {
	return runtime.NewServeMux(
		runtime.WithMarshalerOption(runtime.MIMEWildcard, &runtime.JSONPb{
			MarshalOptions: protojson.MarshalOptions{
				// UseProtoNames:   true,
				EmitDefaultValues: true,
			},
			UnmarshalOptions: protojson.UnmarshalOptions{
				DiscardUnknown: true,
			},
		}),
		runtime.WithErrorHandler(grpcerror.GRPCGatewayErrorHandler),
	)
}