package logger

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/jinzhu/gorm"
)

const (
	zapLoggerKey   = "zap_logger"
	contextGormKey = "added_value"
	loggerGormKey  = "gorm_logger"
)

// WithContext 用来将需要在sql日志记录的附加的值，将key设置为gorm-added
// 举例：context.WithValue(parentCtx, "gorm-added", <value>)
func WithContext(ctx context.Context, db *gorm.DB) *gorm.DB {
	if ctx == nil {
		return db
	}

	return db.Set(contextGormKey, ctx)
}

// 为数据库添加回调方法--添加数据库日志追踪
func AddGormCallbacks(db *gorm.DB, logger *Logger) {
	callbacks := newCallbacks()

	//设置logger用于记录数据库操作日志
	//此处不能使用Set方法存储
	//InstantSet方法才能影响当前的数据库标识
	db.InstantSet(zapLoggerKey, logger)

	registerCallbacks(db, "create", callbacks)
	registerCallbacks(db, "query", callbacks)
	registerCallbacks(db, "update", callbacks)
	registerCallbacks(db, "delete", callbacks)
	registerCallbacks(db, "row_query", callbacks)
}

type callbacks struct{}

func newCallbacks() *callbacks {
	return &callbacks{}
}

func (c *callbacks) beforeCreate(scope *gorm.Scope)   { c.before(scope) }
func (c *callbacks) afterCreate(scope *gorm.Scope)    { c.after(scope, "INSERT") }
func (c *callbacks) beforeQuery(scope *gorm.Scope)    { c.before(scope) }
func (c *callbacks) afterQuery(scope *gorm.Scope)     { c.after(scope, "SELECT") }
func (c *callbacks) beforeUpdate(scope *gorm.Scope)   { c.before(scope) }
func (c *callbacks) afterUpdate(scope *gorm.Scope)    { c.after(scope, "UPDATE") }
func (c *callbacks) beforeDelete(scope *gorm.Scope)   { c.before(scope) }
func (c *callbacks) afterDelete(scope *gorm.Scope)    { c.after(scope, "DELETE") }
func (c *callbacks) beforeRowQuery(scope *gorm.Scope) { c.before(scope) }
func (c *callbacks) afterRowQuery(scope *gorm.Scope)  { c.after(scope, "") }

func (c *callbacks) before(scope *gorm.Scope) {
	fields := Fields{"begin_time": time.Now().Format("2006-01-02 15:04:05.000")}
	ctx, ok := scope.Get(contextGormKey)
	if ok {
		fields["context_value"] = ctx.(context.Context).Value("gorm-added")
	}

	scope.Set(loggerGormKey, fields)
}

func (c *callbacks) after(scope *gorm.Scope, operation string) {
	zapLogger, ok := scope.Get(zapLoggerKey)
	if !ok {
		return
	}
	logger := zapLogger.(*Logger)
	if operation == "" {
		operation = strings.ToUpper(strings.Split(scope.SQL, " ")[0])
	}
	val, ok := scope.Get(loggerGormKey)
	if !ok {
		return
	}
	fields := val.(Fields)
	fields["DBStatement"] = scope.SQL
	fields["DB.SQLVar"] = scope.SQLVars
	scope.SQL = strings.Replace(scope.SQL, "?", "%v", -1)
	fields["DB.SQL"] = fmt.Sprintf(scope.SQL, scope.SQLVars...)
	fields["db.table"] = scope.TableName()
	fields["db.method"] = operation
	fields["db.hasErr"] = scope.HasError()
	fields["db.count"] = scope.DB().RowsAffected
	logger.WithFields(fields).Info("GORM TRACE")
}

func registerCallbacks(db *gorm.DB, name string, c *callbacks) {
	beforeName := fmt.Sprintf("tracing:%v_before", name)
	afterName := fmt.Sprintf("tracing:%v_after", name)
	gormCallbackName := fmt.Sprintf("gorm:%v", name)
	// gorm does some magic, if you pass CallbackProcessor here - nothing works
	switch name {
	case "create":
		db.Callback().Create().Before(gormCallbackName).Register(beforeName, c.beforeCreate)
		db.Callback().Create().After(gormCallbackName).Register(afterName, c.afterCreate)
	case "query":
		db.Callback().Query().Before(gormCallbackName).Register(beforeName, c.beforeQuery)
		db.Callback().Query().After(gormCallbackName).Register(afterName, c.afterQuery)
	case "update":
		db.Callback().Update().Before(gormCallbackName).Register(beforeName, c.beforeUpdate)
		db.Callback().Update().After(gormCallbackName).Register(afterName, c.afterUpdate)
	case "delete":
		db.Callback().Delete().Before(gormCallbackName).Register(beforeName, c.beforeDelete)
		db.Callback().Delete().After(gormCallbackName).Register(afterName, c.afterDelete)
	case "row_query":
		db.Callback().RowQuery().Before(gormCallbackName).Register(beforeName, c.beforeRowQuery)
		db.Callback().RowQuery().After(gormCallbackName).Register(afterName, c.afterRowQuery)
	}
}