package db

import (
	"context"
	"errors"
	"fmt"
	"os"
	"runtime/debug"
	"strings"
	"time"

	"91porn-server/common/log"

	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.mongodb.org/mongo-driver/x/mongo/driver/connstring"
)

var dataSource = make(map[string]*MongoDB)
var registerPool []Register

// DBmap
type DBmap struct {
	Key string
	URL string
}

type Register struct {
	Key   string
	Table []string
}

// MongoOptions 数据库配置
type MongoOptions struct {
	URL string `json:"url"` // 服务器连接地址
}

// MongoDB MongoDB
type MongoDB struct {
	db *mongo.Database
}

func (m *MongoDB) Tool() *MongoTool {
	return m.ToolCtx(context.Background())
}

func (m *MongoDB) ToolCtx(ctx context.Context) *MongoTool {
	return &MongoTool{db: m.db, ctx: ctx}
}

// Coll 获取表名
func (m *MongoDB) Coll(name string) *MongoTool {
	t := m.Tool()
	return t.Coll(name)
}

// CollCtx 获取表名 从外部传入ctx
func (m *MongoDB) CollCtx(ctx context.Context, name string) *MongoTool {
	t := m.ToolCtx(ctx)
	return t.Coll(name)
}

// MongoTool mongo官方库事务封装
type MongoTool struct {
	db   *mongo.Database
	ctx  context.Context // 当前使用的ctx
	coll *mongo.Collection
}

func (m *MongoTool) Coll(name string) *MongoTool {
	opt := options.CollectionOptions{Registry: registry}
	m.coll = m.db.Collection(name, &opt)
	return m
}

// Trans 开启事务处理包裹处理,里面处理的全是利用的事务的ctx
func (m *MongoDB) Trans(fn func(*MongoTool) error, opts ...*TransOpts) error {
	return m.TransCtx(context.Background(), fn, opts...)
}

// TransCtx 外部传入ctx
func (m *MongoDB) TransCtx(ctx context.Context, fn func(*MongoTool) error, opts ...*TransOpts) error {
	t := m.ToolCtx(ctx)
	return m.db.Client().UseSession(ctx, func(sessionContext mongo.SessionContext) error {
		if err := sessionContext.StartTransaction(); err != nil {
			return err
		}
		defer func() {
			if err := recover(); err != nil {
				log.Warn(fmt.Sprintf("caught panic during transaction, aborting. err: %+v, stack: %s", err, debug.Stack()))
				if err := sessionContext.AbortTransaction(sessionContext); err != nil {
					log.Warn("mongo AbortTransaction panic err", log.E(err))
				}
			}
			sessionContext.EndSession(sessionContext)
		}()
		t.ctx = sessionContext
		if err := runTransactionWithRetry(t, fn, MergeTransOpts(opts)); err != nil {
			if strings.Contains(err.Error(), "NoSuchTransaction") {
				log.Warn("NoSuchTransaction error, return")
				return err
			}
			if err := sessionContext.AbortTransaction(sessionContext); err != nil {
				log.Warn("mongo AbortTransaction err", log.E(err))
			}
			log.Warn("caught exception during transaction, aborting.", log.E(err))
			sessionContext.EndSession(sessionContext)
			return err
		}
		return commitWithRetry(sessionContext)
	})
}

// runTransactionWithRetry is an example function demonstrating transaction retry logic.
func runTransactionWithRetry(t *MongoTool, txnFn func(t *MongoTool) error, opts *TransOpts) error {
	//no set ReEntryCount is loop retry
	for {
		err := txnFn(t) // Performs transaction.
		if err == nil {
			return nil
		}
		log.Warn("Transaction aborted. Caught exception during transaction.", log.E(err))
		// If transient error, retry the whole transaction
		if strings.Contains(err.Error(), "NoSuchTransaction") {
			log.Info("NoSuchTransaction error,return and break retry loop")
			return err
		}
		cmdErr, ok := err.(mongo.CommandError)
		if ok && cmdErr.HasErrorLabel("TransientTransactionError") {
			if opts != nil {
				if opts.ReEntryCount != nil {
					if *opts.ReEntryCount <= 0 {
						return cmdErr
					}
					*opts.ReEntryCount--
				}
			}
			log.Info("TransientTransactionError, retrying transaction...")
			continue
		}
		// else return err
		return err
	}
}

// commitWithRetry is an example function demonstrating transaction retry logic.
func commitWithRetry(sess mongo.SessionContext) error {
	for {
		err := sess.CommitTransaction(sess)
		switch e := err.(type) {
		case nil:
			log.Info("Transaction committed.")
			return nil
		case mongo.CommandError:
			// Can retry commit
			if e.HasErrorLabel("UnknownTransactionCommitResult") {
				log.Info("UnknownTransactionCommitResult, retrying commit operation...")
				continue
			}
			log.Info("Error during commit...")
			return e
		default:
			log.Info("Error during commit...")
			return e
		}
	}
}

// Disconnect 关闭数据库连接
func (m *MongoDB) Disconnect() error {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	log.Info("closing  mongodb  connections")
	defer cancel()
	if err := m.db.Client().Disconnect(ctx); err != nil {
		log.Warn(fmt.Sprintf("close mongo connections err: %+v", err))
		return err
	}
	return nil
}

// InitMongo  初始化Mongo
func InitMongo(murl string) (*MongoDB, error) {
	cs, err := connstring.Parse(murl)
	if err != nil {
		log.Error("mongo URL parse fail", log.Any("url", murl), log.E(err))
		return nil, err
	}
	db := cs.Database
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	// Connect to MongoDB
	mongoClient, err := mongo.Connect(ctx, options.Client().ApplyURI(murl))
	if err != nil {
		log.Error("mongodb connect fail", log.E(err))
		return nil, err
	}
	// Check the connection
	if err = mongoClient.Ping(context.Background(), nil); err != nil {
		log.Error("mongodb connect ping is fail")
		return nil, err
	}
	mongoDataBase := mongoClient.Database(db)
	return &MongoDB{db: mongoDataBase}, nil
}

func InitDS(dbmap []DBmap, register []Register) map[string]*MongoDB {
	if len(dbmap) == 0 {
		log.Error("dbmap must be not empty ")
		os.Exit(1)
	}
	for _, v := range dbmap {
		db, err := InitMongo(v.URL)
		if err != nil {
			log.Error(fmt.Sprintf("[DB-%s] start up error", v.Key), log.E(err))
			os.Exit(1)
		}
		log.Info(fmt.Sprintf("[DB-%s] connect is successfully", v.Key))
		dataSource[v.Key] = db
	}
	registerPool = register
	return dataSource
}

func Init(tableName string) *MongoDB {
	if tableName == "" {
		panic(errors.New("table name must not be empty"))
	}
	return Selector(tableName)
}

// 初始化只读数据
func InitRead(tableName string) *MongoDB {
	if tableName == "" {
		panic(errors.New("table name must not be empty"))
	}
	return SelectorRead(tableName)
}

func CloseDS() {
	if len(dataSource) == 0 {
		log.Warn("dataSource empty ")
		return
	}
	for k, v := range dataSource {
		if err := v.Disconnect(); err != nil {
			log.Error(fmt.Sprintf("[DB-%s]Mongo Disconnect error", k), log.E(err))
			continue
		}
		log.Info(fmt.Sprintf("[DB-%s]Mongo Disconnect OK", k))
	}
}

// BaseDAO 如果ctx为nil 则表示不使用事务,如果ctx不为空则表示使用事务
func BaseDAO(tableName string, ctx context.Context) *MongoTool {
	if tableName == "" {
		panic(errors.New("table name must not be empty"))
	}
	mongdb := Selector(tableName)
	if ctx == nil {
		return &MongoTool{db: mongdb.db, ctx: context.Background(), coll: mongdb.db.Collection(tableName)}
	}
	return &MongoTool{db: mongdb.db, ctx: ctx, coll: mongdb.db.Collection(tableName)}
}

func Selector(tableName string) *MongoDB {
	var isExist = false
	var db *MongoDB
	for _, r := range registerPool {
		if strings.HasPrefix(r.Key, "Read") {
			continue
		}
		for _, v := range r.Table {
			if v == tableName {
				isExist = true
				db = dataSource[r.Key]
				break
			}
		}
	}
	if !isExist {
		log.Warn("current table not register,please register it first", log.Any("table", tableName))
		panic(errors.New("current table not register"))
	}
	return db
}

func SelectorRead(tableName string) *MongoDB {
	var isExist = false
	var db *MongoDB
	for _, r := range registerPool {
		if !strings.HasPrefix(r.Key, "Read") {
			continue
		}
		for _, v := range r.Table {
			if v == tableName {
				isExist = true
				db = dataSource[r.Key]
				break
			}
		}
	}
	if !isExist {
		log.Warn("current table not register,please register it first", log.Any("table", tableName))
		panic(errors.New("current table not register"))
	}
	return db
}