package mongo

import (
	"reflect"

	"github.com/shopspring/decimal"
	"github.com/zeromicro/go-zero/core/stores/mon"
	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
)

type TypeCodec struct {
	ValueType reflect.Type
	Encoder   bson.ValueEncoder
	Decoder   bson.ValueDecoder
}

// WithTypeCodec registers TypeCodecs to convert custom types.
func WithTypeCodec(typeCodecs ...TypeCodec) mon.Option {
	return func(c *options.ClientOptions) {
		registry := bson.NewRegistry()
		for _, v := range typeCodecs {
			registry.RegisterTypeEncoder(v.ValueType, v.Encoder)
			registry.RegisterTypeDecoder(v.ValueType, v.Decoder)
		}
		c.SetRegistry(registry)
	}
}

// SetCustomDecimalType force convert primitive.Decimal128 to decimal.Decimal.
func SetCustomDecimalType() mon.Option {
	return WithTypeCodec(TypeCodec{
		ValueType: reflect.TypeOf(decimal.Decimal{}),
		Encoder:   &MgoDecimal{},
		Decoder:   &MgoDecimal{},
	})
}

func InitMongoOptions(cfg Conf) mon.Option {
	return func(opts *options.ClientOptions) {
		opts.SetMaxPoolSize(cfg.MaxPoolSize)
		opts.SetMinPoolSize(cfg.MinPoolSize)
		opts.SetMaxConnIdleTime(cfg.MaxConnIdleTime)
		opts.SetCompressors([]string{"snappy"})
	}
}
