package db

import (
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"gorm.io/gorm"
)

const (
	callBackBeforeName = "core:before"
	callBackAfterName  = "core:after"
	startTime          = "_start_time"

	gormErrorCounter       = "gorm_errors_total"
	gormExecutionHistogram = "gorm_execution_seconds"
)

var (
	durationHistogram *prometheus.HistogramVec
	errorCount        *prometheus.CounterVec
)

type MetricsPlugin struct{}

func (op *MetricsPlugin) Name() string {
	return "metricsPlugin"
}

func (op *MetricsPlugin) Initialize(db *gorm.DB) (err error) {
	_ = db.Callback().Create().Before("gorm:create").Register(callBackBeforeName, before)
	_ = db.Callback().Query().Before("gorm:query").Register(callBackBeforeName, before)
	_ = db.Callback().Delete().Before("gorm:delete").Register(callBackBeforeName, before)
	_ = db.Callback().Update().Before("gorm:update").Register(callBackBeforeName, before)
	_ = db.Callback().Row().Before("gorm:row").Register(callBackBeforeName, before)
	_ = db.Callback().Raw().Before("gorm:raw").Register(callBackBeforeName, before)

	_ = db.Callback().Create().After("gorm:create").Register(callBackAfterName, after)
	_ = db.Callback().Query().After("gorm:query").Register(callBackAfterName, after)
	_ = db.Callback().Delete().After("gorm:delete").Register(callBackAfterName, after)
	_ = db.Callback().Update().After("gorm:update").Register(callBackAfterName, after)
	_ = db.Callback().Row().After("gorm:row").Register(callBackAfterName, after)
	_ = db.Callback().Raw().After("gorm:raw").Register(callBackAfterName, after)

	durationHistogram = prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    gormExecutionHistogram,
			Help:    "record gorm execution time",
			Buckets: []float64{0.001, 0.01, 0.05, 0.1, 0.5, 1, 5, 15},
		}, []string{"table", "sql"},
	)

	errorCount = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Name: gormErrorCounter,
			Help: "record gorm error count",
		}, []string{"table", "sql", "error"},
	)

	// register metric
	prometheus.MustRegister(durationHistogram, errorCount)
	return
}

var _ gorm.Plugin = &MetricsPlugin{}

func before(db *gorm.DB) {
	db.InstanceSet(startTime, time.Now())
}

func after(db *gorm.DB) {
	_ts, isExist := db.InstanceGet(startTime)
	if !isExist {
		return
	}

	table := db.Statement.Table
	sql := db.Statement.SQL.String()
	err := db.Error
	if err != nil {
		errorCount.WithLabelValues(table, sql, err.Error()).Inc()
	}

	ts, ok := _ts.(time.Time)
	if !ok {
		return
	}
	duration := time.Since(ts).Seconds()
	durationHistogram.WithLabelValues(table, sql).Observe(duration)
}