package mongo

import (
	"context"
	"fmt"
	"github.com/qiniu/qmgo"
	mongoMiddleware "github.com/qiniu/qmgo/middleware"
	"go.mongodb.org/mongo-driver/mongo/readpref"
	"net/url"
	"strings"
)

type WrapperMongo struct {
	Client *qmgo.Client
	Name   string
}

type Setting struct {
	Name string `mapstructure:"name"`
	// shading或者replica_set，默认replica_set
	Type string `mapstructure:"type"`
	// 地址，不能为空
	Addr []string `mapstructure:"addr"`
	// primary、primaryPreferred、secondary、secondaryPreferred、nearest
	// 默认nearest
	ReadPreferenceMode string `mapstructure:"read_preference_mode"`
	//mongodb版本必须大于3.4才能设置，否侧报错，默认0即可
	MaxStalenessMS int64 `mapstructure:"max_staleness_ms"`
	// 最大连接数量，默认5
	MaxPoolSize uint64 `mapstructure:"max_pool_size"`
	// 最小连接数量，默认2
	MinPoolSize uint64 `mapstructure:"minPoolSize"`
	// 连接超时时间，默认2000ms
	ConnectTimeoutMs int64 `mapstructure:"connect_timeout_ms"`
	// 操作超时时间，默认200ms
	SocketTimeoutMS int64 `mapstructure:"socket_timeout_ms"`
	// 权限
	Auth *Auth `mapstructure:"auth"`
}

type Auth struct {
	// AuthMechanism: the mechanism to use for authentication. Supported values include "SCRAM-SHA-256", "SCRAM-SHA-1",
	// "MONGODB-CR", "PLAIN", "GSSAPI", "MONGODB-X509", and "MONGODB-AWS". This can also be set through the "authMechanism"
	// 一般情况 "SCRAM-SHA-256", "SCRAM-SHA-1"
	AuthMechanism string `mapstructure:"auth_mechanism"`
	// AuthSource: the name of the database to use for authentication. This defaults to "$external" for MONGODB-X509,
	// GSSAPI, and PLAIN and "admin" for all other mechanisms. This can also be set through the "authSource" URI option
	// 一般情况上为admin
	AuthSource string `mapstructure:"auth_source"`
	Username   string `mapstructure:"username"`
	Password   string `mapstructure:"password"`
	// PasswordSet对于对于GSSAPI生效，其它此字段被忽略
	PasswordSet bool `mapstructure:"password_set"`
}

// InitMongo 返回mongo连接池实例
// 用法：mongoPool.Client.Database("test_db").Collection("test_collection").
//			Find(ctx, bson.M{"_id": objId}).
//			One(&res)
// 依赖Qmgo，文档地址：https://github.com/qiniu/qmgo
func InitMongo(setting Setting, middleware ...callback) (*WrapperMongo, error) {
	var errMsg interface{}
	if setting.Name == "" || len(setting.Addr) == 0 {
		errMsg = "need name or addr"
		panic(errMsg)
	}

	mode, err := readpref.ModeFromString(setting.ReadPreferenceMode)
	if err != nil {
		mode = readpref.NearestMode
	}

	readPref := qmgo.ReadPref{
		Mode:           mode,
		MaxStalenessMS: setting.MaxStalenessMS,
	}

	if setting.ConnectTimeoutMs == 0 {
		setting.ConnectTimeoutMs = 2000
	}

	if setting.SocketTimeoutMS == 0 {
		setting.SocketTimeoutMS = 2000
	}

	if setting.MaxPoolSize == 0 {
		setting.MaxPoolSize = 5
	}

	if setting.MinPoolSize == 0 {
		setting.MaxPoolSize = 2
	}

	var auth *qmgo.Credential
	if setting.Auth != nil && setting.Auth.Username != "" {
		// 针对特殊字段，必须url转码才能正常
		setting.Auth.Username = url.QueryEscape(setting.Auth.Username)
		setting.Auth.Password = url.QueryEscape(setting.Auth.Password)
		// uri中已有，但不包含PasswordSet，没细测，保留最好
		auth = &qmgo.Credential{
			AuthMechanism: setting.Auth.AuthMechanism,
			AuthSource:    setting.Auth.AuthSource,
			Username:      setting.Auth.Username,
			Password:      setting.Auth.Password,
			PasswordSet:   setting.Auth.PasswordSet,
		}
	}
	_config := &qmgo.Config{
		Uri:              GetMongoUrl(setting),
		MaxPoolSize:      &setting.MaxPoolSize,
		MinPoolSize:      &setting.MinPoolSize,
		ConnectTimeoutMS: &setting.SocketTimeoutMS,
		SocketTimeoutMS:  &setting.SocketTimeoutMS,
		ReadPreference:   &readPref,
	}
	if auth != nil {
		_config.Auth = auth
	}

	ctx := context.Background()
	client, err := qmgo.NewClient(ctx, _config)
	if err != nil {
		return nil, err
	}

	for _, plugin := range middleware {
		mongoMiddleware.Register(plugin)
	}
	// 开启mon打点
	mongoMiddleware.Register(MonMongoPlugin)

	return &WrapperMongo{
		Client: client,
		Name:   setting.Name,
	}, nil
}

func GetMongoUrl(setting Setting) string {
	var paramStr []string
	joinPath := strings.Join(setting.Addr, ",")
	if setting.Auth != nil && setting.Auth.Username != "" {
		// 解决用户密码特殊字符必须放uri中和新老版本兼容问题
		uri := fmt.Sprintf("mongodb://%s:%s@%s/", setting.Auth.Username, setting.Auth.Password, joinPath)
		if setting.Auth.AuthMechanism != "" {
			paramStr = append(paramStr, "authMechanism="+setting.Auth.AuthMechanism)
		}
		if setting.Auth.AuthSource != "" {
			paramStr = append(paramStr, "authSource="+setting.Auth.AuthSource)
		}
		if len(paramStr) > 0 {
			uri += "?" + strings.Join(paramStr, "&")
		}
		return uri
	} else {
		return fmt.Sprintf("mongodb://%s", joinPath)
	}
}

// Close 关闭连接池
func (w *WrapperMongo) Close() {
	ctx := context.Background()
	_ = w.Client.Close(ctx)
}