package util

import (
	"context"
	"fmt"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"math"
	"time"
)

type MongoCollection struct {
	Collection *mongo.Collection
	Logger     func(ctx context.Context, log *MongoOperationLog)
}

type MongoOperationLog bson.M

// NewMongoCollection 新连接
func NewMongoCollection(cli *mongo.Client, db, col string, fn func(ctx context.Context, log *MongoOperationLog)) *MongoCollection {
	return &MongoCollection{
		Collection: cli.Database(db).Collection(col),
		Logger:     fn,
	}
}

// Names 获取库名和表名
func (mc *MongoCollection) Names() (string, string) {
	return mc.Collection.Database().Name(), mc.Collection.Name()
}

// NewOperationLog 一条日志
func (mc *MongoCollection) NewOperationLog(do string, set, whr, opt, res any, err error) *MongoOperationLog {
	db, col := mc.Names()
	return &MongoOperationLog{
		`db`:  db,
		`col`: col,
		`do`:  do,
		`set`: set,
		`whr`: whr,
		`opt`: opt,
		`res`: res,
		`err`: err,
		`at`:  time.Now().Unix(),
	}
}

// InsertOne 插入一条数据
func (mc *MongoCollection) InsertOne(ctx context.Context, item any, opts ...*options.InsertOneOptions) (newId any, err error) {
	defer func() {
		if mc.Logger != nil {
			go mc.Logger(ctx, mc.NewOperationLog(`InsertOne`, item, nil, opts, newId, err))
		}
	}()
	var res *mongo.InsertOneResult
	if res, err = mc.Collection.InsertOne(ctx, item, opts...); err != nil {
		return
	}
	newId = res.InsertedID
	return
}

// InsertMany 插入多条数据
func (mc *MongoCollection) InsertMany(ctx context.Context, items []any, opts ...*options.InsertManyOptions) (newIds []any, err error) {
	defer func() {
		if mc.Logger != nil {
			go mc.Logger(ctx, mc.NewOperationLog(`InsertMany`, items, nil, opts, newIds, err))
		}
	}()
	var res *mongo.InsertManyResult
	if res, err = mc.Collection.InsertMany(ctx, items, opts...); err != nil {
		return
	}
	newIds = res.InsertedIDs
	return
}

// UpdateByID 按ID更新数据
func (mc *MongoCollection) UpdateByID(ctx context.Context, id, set any, opts ...*options.UpdateOptions) (res *mongo.UpdateResult, err error) {
	defer func() {
		if mc.Logger != nil {
			go mc.Logger(ctx, mc.NewOperationLog(`UpdateByID`, set, id, opts, res, err))
		}
	}()
	res, err = mc.Collection.UpdateByID(ctx, id, bson.M{`$set`: set}, opts...)
	return
}

// UpdateOne 按条件更新一条数据
func (mc *MongoCollection) UpdateOne(ctx context.Context, where, set any, opts ...*options.UpdateOptions) (res *mongo.UpdateResult, err error) {
	defer func() {
		if mc.Logger != nil {
			go mc.Logger(ctx, mc.NewOperationLog(`UpdateOne`, set, where, opts, res, err))
		}
	}()
	res, err = mc.Collection.UpdateOne(ctx, where, bson.M{`$set`: set}, opts...)
	return
}

// UpdateMany 按条件更新多条数据
func (mc *MongoCollection) UpdateMany(ctx context.Context, where, set any, opts ...*options.UpdateOptions) (res *mongo.UpdateResult, err error) {
	defer func() {
		if mc.Logger != nil {
			go mc.Logger(ctx, mc.NewOperationLog(`UpdateMany`, set, where, opts, res, err))
		}
	}()
	res, err = mc.Collection.UpdateMany(ctx, where, bson.M{`$set`: set}, opts...)
	return
}

// FindOneAndUpdate 找到一条并更新
func (mc *MongoCollection) FindOneAndUpdate(ctx context.Context, where, set any, opts ...*options.FindOneAndUpdateOptions) (res bson.M, err error) {
	defer func() {
		if mc.Logger != nil {
			go mc.Logger(ctx, mc.NewOperationLog(`FindOneAndUpdate`, set, where, opts, res, err))
		}
	}()
	err = mc.Collection.FindOneAndUpdate(ctx, where, bson.M{`$set`: set}, opts...).Decode(&res)
	return
}

// FindOneAndReplace 找到一条并替换
func (mc *MongoCollection) FindOneAndReplace(ctx context.Context, where, set any, opts ...*options.FindOneAndReplaceOptions) (res bson.M, err error) {
	defer func() {
		if mc.Logger != nil {
			go mc.Logger(ctx, mc.NewOperationLog(`FindOneAndReplace`, set, where, opts, res, err))
		}
	}()
	err = mc.Collection.FindOneAndReplace(ctx, where, bson.M{`$set`: set}, opts...).Decode(&res)
	return
}

// Drop 清空
func (mc *MongoCollection) Drop(ctx context.Context) (err error) {
	defer func() {
		if mc.Logger != nil {
			go mc.Logger(ctx, mc.NewOperationLog(`Drop`, nil, nil, nil, nil, err))
		}
	}()
	err = mc.Collection.Drop(ctx)
	return
}

// DeleteOne 删除一条记录
func (mc *MongoCollection) DeleteOne(ctx context.Context, where any, opts ...*options.DeleteOptions) (num int64, err error) {
	defer func() {
		if mc.Logger != nil {
			go mc.Logger(ctx, mc.NewOperationLog(`DeleteOne`, nil, where, opts, num, err))
		}
	}()
	var res *mongo.DeleteResult
	if res, err = mc.Collection.DeleteOne(ctx, where, opts...); err != nil {
		return
	}
	num = res.DeletedCount
	return
}

// DeleteMany 删除多条数据
func (mc *MongoCollection) DeleteMany(ctx context.Context, where any, opts ...*options.DeleteOptions) (num int64, err error) {
	defer func() {
		if mc.Logger != nil {
			go mc.Logger(ctx, mc.NewOperationLog(`DeleteMany`, nil, where, opts, num, err))
		}
	}()
	var res *mongo.DeleteResult
	if res, err = mc.Collection.DeleteMany(ctx, where, opts...); err != nil {
		return
	}
	num = res.DeletedCount
	return
}

// FindOneAndDelete 找到一条记录并删除
func (mc *MongoCollection) FindOneAndDelete(ctx context.Context, where any, opts ...*options.FindOneAndDeleteOptions) (res bson.M, err error) {
	defer func() {
		if mc.Logger != nil {
			go mc.Logger(ctx, mc.NewOperationLog(`FindOneAndDelete`, nil, where, opts, res, err))
		}
	}()
	err = mc.Collection.FindOneAndDelete(ctx, where, opts...).Decode(&res)
	return
}

// CursorWalk 逐条处理结果数据
func (mc *MongoCollection) CursorWalk(ctx context.Context, fn func(v bson.M) error, cursor *mongo.Cursor) error {
	defer func() {
		_ = cursor.Close(ctx)
	}()
	for cursor.Next(ctx) {
		var v bson.M
		if err := cursor.Decode(&v); err == nil {
			if err = fn(v); err != nil {
				return err
			}
		} else {
			return err
		}
	}
	return nil
}

// CursorAll 一次性拿到全部结果
func (mc *MongoCollection) CursorAll(ctx context.Context, cursor *mongo.Cursor) ([]bson.M, error) {
	defer func() {
		_ = cursor.Close(ctx)
	}()
	var v []bson.M
	if err := cursor.All(ctx, &v); err != nil {
		return v, err
	}
	return v, nil
}

// IndexesList 获取索引
func (mc *MongoCollection) IndexesList(ctx context.Context, opts ...*options.ListIndexesOptions) ([]bson.M, error) {
	if cursor, err := mc.Collection.Indexes().List(ctx, opts...); err != nil {
		return nil, err
	} else {
		return mc.CursorAll(ctx, cursor)
	}
}

// IndexesCreateOne 创建一条索引
func (mc *MongoCollection) IndexesCreateOne(ctx context.Context, model mongo.IndexModel, opts ...*options.CreateIndexesOptions) (res string, err error) {
	defer func() {
		if mc.Logger != nil {
			go mc.Logger(ctx, mc.NewOperationLog(`IndexesCreateOne`, model, nil, opts, res, err))
		}
	}()
	res, err = mc.Collection.Indexes().CreateOne(ctx, model, opts...)
	return
}

// IndexesCreateMany 创建多条索引
func (mc *MongoCollection) IndexesCreateMany(ctx context.Context, models []mongo.IndexModel, opts ...*options.CreateIndexesOptions) (res []string, err error) {
	defer func() {
		if mc.Logger != nil {
			go mc.Logger(ctx, mc.NewOperationLog(`IndexesCreateMany`, models, nil, opts, res, err))
		}
	}()
	res, err = mc.Collection.Indexes().CreateMany(ctx, models, opts...)
	return
}

// IndexesDropOne 删除一条索引
func (mc *MongoCollection) IndexesDropOne(ctx context.Context, name string, opts ...*options.DropIndexesOptions) (res bson.Raw, err error) {
	defer func() {
		if mc.Logger != nil {
			go mc.Logger(ctx, mc.NewOperationLog(`IndexesDropOne`, name, nil, opts, res, err))
		}
	}()
	res, err = mc.Collection.Indexes().DropOne(ctx, name, opts...)
	return
}

// IndexesDropAll 删除全部索引
func (mc *MongoCollection) IndexesDropAll(ctx context.Context, opts ...*options.DropIndexesOptions) (res bson.Raw, err error) {
	defer func() {
		if mc.Logger != nil {
			go mc.Logger(ctx, mc.NewOperationLog(`IndexesDropAll`, nil, nil, opts, res, err))
		}
	}()
	res, err = mc.Collection.Indexes().DropAll(ctx, opts...)
	return
}

// FindOne 查找一条记录
func (mc *MongoCollection) FindOne(ctx context.Context, where any, opts ...*options.FindOneOptions) (bson.M, error) {
	var v bson.M
	if err := mc.Collection.FindOne(ctx, where, opts...).Decode(&v); err != nil {
		return v, err
	} else {
		return v, nil
	}
}

// Find 按条件查找记录
func (mc *MongoCollection) Find(ctx context.Context, where any, opts ...*options.FindOptions) ([]bson.M, error) {
	if cursor, err := mc.Collection.Find(ctx, where, opts...); err != nil {
		return nil, err
	} else {
		return mc.CursorAll(ctx, cursor)
	}
}

// FindWalk 按条件查找并逐条处理
func (mc *MongoCollection) FindWalk(ctx context.Context, fn func(v bson.M) error, where any, opts ...*options.FindOptions) error {
	if cursor, err := mc.Collection.Find(ctx, where, opts...); err != nil {
		return err
	} else {
		return mc.CursorWalk(ctx, fn, cursor)
	}
}

// CountDocuments 统计记录数量
func (mc *MongoCollection) CountDocuments(ctx context.Context, where any, opts ...*options.CountOptions) (int64, error) {
	return mc.Collection.CountDocuments(ctx, where, opts...)
}

// PageWalk 分页取数据
func (mc *MongoCollection) PageWalk(
	ctx context.Context,
	fn func(v bson.M) error,
	page, pageSize int64,
	where, sort any,
	opts ...*options.FindOptions,
) (itemTotal, pageTotal int64, err error) {
	if pageSize <= 0 {
		pageSize = 10
	}
	if page < 1 {
		err = fmt.Errorf(`页码从1开始，当前是：%d`, page)
		return
	}
	if itemTotal, err = mc.CountDocuments(ctx, where); err != nil || itemTotal == 0 {
		return
	}
	if pageSize == 1 {
		pageTotal = itemTotal
	} else {
		pageTotal = int64(math.Ceil(float64(itemTotal) / float64(pageSize)))
	}
	if page > pageTotal {
		err = fmt.Errorf(`最大页：%d，当前页：%d，已超出`, pageTotal, page)
		return
	}
	opts = append(opts, options.Find().SetLimit(pageSize), options.Find().SetSkip((page-1)*pageSize))
	if sort != nil {
		opts = append(opts, options.Find().SetSort(sort))
	}
	err = mc.FindWalk(ctx, fn, where, opts...)
	return
}

// Distinct 查找并去重复
func (mc *MongoCollection) Distinct(ctx context.Context, name string, where any, opts ...*options.DistinctOptions) ([]any, error) {
	return mc.Collection.Distinct(ctx, name, where, opts...)
}

// BulkWrite 批量写入指令
func (mc *MongoCollection) BulkWrite(ctx context.Context, models []mongo.WriteModel, opts ...*options.BulkWriteOptions) (res *mongo.BulkWriteResult, err error) {
	defer func() {
		if mc.Logger != nil {
			go mc.Logger(ctx, mc.NewOperationLog(`BulkWrite`, models, nil, opts, res, err))
		}
	}()
	res, err = mc.Collection.BulkWrite(ctx, models, opts...)
	return
}

// AggregateWalk 参考：https://www.mongodb.com/docs/manual/reference/operator/aggregation-pipeline/#db-collection-aggregate-stages
func (mc *MongoCollection) AggregateWalk(ctx context.Context, fn func(v bson.M) error, pipeline mongo.Pipeline, opts ...*options.AggregateOptions) error {
	if cursor, err := mc.Collection.Aggregate(ctx, pipeline, opts...); err != nil {
		return err
	} else {
		return mc.CursorWalk(ctx, fn, cursor)
	}
}