package storage

import (
	"context"
	"fmt"
	"time"

	metametrics "github.com/tidbcloud/aws-shared-provider/pkg/metaserver/metrics"
	"gorm.io/gorm"
)

const (
	mysqlMetricsPluginName = "metaserver_mysql_metrics"
	mysqlQueryStartedAtKey = "metaserver:mysql_metrics:query:started_at"
)

type mysqlMetricsPlugin struct{}

func (*mysqlMetricsPlugin) Name() string {
	return mysqlMetricsPluginName
}

func (*mysqlMetricsPlugin) Initialize(db *gorm.DB) error {
	if err := db.Callback().Query().Before("gorm:query").
		Register("metaserver:mysql_metrics:before_query", beforeMySQLQuery); err != nil {
		return fmt.Errorf("register before-query mysql metrics callback: %w", err)
	}
	if err := db.Callback().Query().After("gorm:query").
		Register("metaserver:mysql_metrics:after_query", afterMySQLQuery); err != nil {
		return fmt.Errorf("register after-query mysql metrics callback: %w", err)
	}
	return nil
}

var _ gorm.Plugin = (*mysqlMetricsPlugin)(nil)

func beforeMySQLQuery(db *gorm.DB) {
	if db.DryRun {
		return
	}
	db.InstanceSet(mysqlQueryStartedAtKey, time.Now())
}

func afterMySQLQuery(db *gorm.DB) {
	if db.DryRun {
		return
	}
	startedAtValue, ok := db.InstanceGet(mysqlQueryStartedAtKey)
	if !ok {
		return
	}
	startedAt, ok := startedAtValue.(time.Time)
	if !ok {
		return
	}
	metametrics.ObserveMySQLOperation(metametrics.MySQLOperationQuery, startedAt, db.Error)
}

func (m *memDBStorage) runMySQLTransaction(
	ctx context.Context,
	fn func(*gorm.DB) error,
) error {
	startedAt := time.Now()
	err := m.mysqlDB.WithContext(ctx).Transaction(fn)
	metametrics.ObserveMySQLOperation(
		metametrics.MySQLOperationTransaction,
		startedAt,
		err,
	)
	return err
}