package storage

import (
	"context"
	"path/filepath"
	"testing"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/testutil"
	dto "github.com/prometheus/client_model/go"
	"github.com/stretchr/testify/require"
	metametrics "github.com/tidbcloud/aws-shared-provider/pkg/metaserver/metrics"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

func TestMySQLMetricsPluginRecordsQueryResultsAndAvailability(t *testing.T) {
	resetMySQLMetrics(t)

	db, err := gorm.Open(
		sqlite.Open(filepath.Join(t.TempDir(), "mysql-metrics.db")),
		&gorm.Config{Logger: logger.Default.LogMode(logger.Silent)},
	)
	require.NoError(t, err)
	require.NoError(t, db.Exec("CREATE TABLE metric_rows (id INTEGER PRIMARY KEY)").Error)
	require.NoError(t, db.Exec("INSERT INTO metric_rows (id) VALUES (1)").Error)
	require.NoError(t, db.Use(&mysqlMetricsPlugin{}))

	var rows []struct {
		ID int
	}
	require.NoError(t, db.Table("metric_rows").Find(&rows).Error)
	require.Equal(t, float64(1), testutil.ToFloat64(
		metametrics.MySQLOperationsTotal.WithLabelValues(
			metametrics.MySQLOperationQuery,
			metametrics.ResultSuccess,
		),
	))
	require.Equal(t, uint64(1), mysqlOperationDurationCount(
		t,
		metametrics.MySQLOperationQuery,
		metametrics.ResultSuccess,
	))
	require.Equal(t, float64(1), testutil.ToFloat64(metametrics.MySQLLastOperationSuccess))

	require.Error(t, db.Table("missing_metric_rows").Find(&rows).Error)
	require.Equal(t, float64(1), testutil.ToFloat64(
		metametrics.MySQLOperationsTotal.WithLabelValues(
			metametrics.MySQLOperationQuery,
			metametrics.ResultFailure,
		),
	))
	require.Equal(t, uint64(1), mysqlOperationDurationCount(
		t,
		metametrics.MySQLOperationQuery,
		metametrics.ResultFailure,
	))
	require.Equal(t, float64(0), testutil.ToFloat64(metametrics.MySQLLastOperationSuccess))

	require.NoError(t, db.Table("metric_rows").Find(&rows).Error)
	require.Equal(t, float64(2), testutil.ToFloat64(
		metametrics.MySQLOperationsTotal.WithLabelValues(
			metametrics.MySQLOperationQuery,
			metametrics.ResultSuccess,
		),
	))
	require.Equal(t, float64(1), testutil.ToFloat64(metametrics.MySQLLastOperationSuccess))
}

func TestMySQLMetricsPluginIgnoresDryRunQueries(t *testing.T) {
	resetMySQLMetrics(t)

	db, err := gorm.Open(
		sqlite.Open(filepath.Join(t.TempDir(), "mysql-dry-run-metrics.db")),
		&gorm.Config{Logger: logger.Default.LogMode(logger.Silent)},
	)
	require.NoError(t, err)
	require.NoError(t, db.Use(&mysqlMetricsPlugin{}))

	var rows []struct {
		ID int
	}
	require.NoError(t, db.Session(&gorm.Session{DryRun: true}).
		Table("metric_rows").
		Find(&rows).Error)

	for _, result := range []string{metametrics.ResultSuccess, metametrics.ResultFailure} {
		require.Zero(t, testutil.ToFloat64(
			metametrics.MySQLOperationsTotal.WithLabelValues(
				metametrics.MySQLOperationQuery,
				result,
			),
		))
		require.Zero(t, mysqlOperationDurationCount(
			t,
			metametrics.MySQLOperationQuery,
			result,
		))
	}
	require.Zero(t, testutil.ToFloat64(metametrics.MySQLLastOperationSuccess))
}

func TestRunMySQLTransactionRecordsSuccess(t *testing.T) {
	resetMySQLMetrics(t)

	db, err := gorm.Open(
		sqlite.Open(filepath.Join(t.TempDir(), "mysql-transaction-success-metrics.db")),
		&gorm.Config{Logger: logger.Default.LogMode(logger.Silent)},
	)
	require.NoError(t, err)

	store := &memDBStorage{mysqlDB: db}
	require.NoError(t, store.runMySQLTransaction(context.Background(), func(*gorm.DB) error {
		return nil
	}))
	require.Equal(t, float64(1), testutil.ToFloat64(
		metametrics.MySQLOperationsTotal.WithLabelValues(
			metametrics.MySQLOperationTransaction,
			metametrics.ResultSuccess,
		),
	))
	require.Equal(t, uint64(1), mysqlOperationDurationCount(
		t,
		metametrics.MySQLOperationTransaction,
		metametrics.ResultSuccess,
	))
	require.Equal(t, float64(1), testutil.ToFloat64(metametrics.MySQLLastOperationSuccess))
}

func TestRunMySQLTransactionRecordsBeginFailure(t *testing.T) {
	resetMySQLMetrics(t)

	db, err := gorm.Open(
		sqlite.Open(filepath.Join(t.TempDir(), "mysql-transaction-metrics.db")),
		&gorm.Config{Logger: logger.Default.LogMode(logger.Silent)},
	)
	require.NoError(t, err)
	require.NoError(t, db.Use(&mysqlMetricsPlugin{}))
	sqlDB, err := db.DB()
	require.NoError(t, err)
	require.NoError(t, sqlDB.Close())

	store := &memDBStorage{mysqlDB: db}
	err = store.runMySQLTransaction(context.Background(), func(*gorm.DB) error {
		return nil
	})
	require.Error(t, err)
	require.Equal(t, float64(1), testutil.ToFloat64(
		metametrics.MySQLOperationsTotal.WithLabelValues(
			metametrics.MySQLOperationTransaction,
			metametrics.ResultFailure,
		),
	))
	require.Equal(t, uint64(1), mysqlOperationDurationCount(
		t,
		metametrics.MySQLOperationTransaction,
		metametrics.ResultFailure,
	))
	require.Equal(t, float64(0), testutil.ToFloat64(metametrics.MySQLLastOperationSuccess))
}

func TestRunMySQLTransactionRecordsCommitFailure(t *testing.T) {
	resetMySQLMetrics(t)

	db, err := gorm.Open(
		sqlite.Open(filepath.Join(t.TempDir(), "mysql-transaction-commit-metrics.db")),
		&gorm.Config{Logger: logger.Default.LogMode(logger.Silent)},
	)
	require.NoError(t, err)
	require.NoError(t, db.Exec("PRAGMA foreign_keys = ON").Error)
	require.NoError(t, db.Exec("CREATE TABLE parents (id INTEGER PRIMARY KEY)").Error)
	require.NoError(t, db.Exec(`
		CREATE TABLE children (
			id INTEGER PRIMARY KEY,
			parent_id INTEGER,
			FOREIGN KEY (parent_id) REFERENCES parents(id) DEFERRABLE INITIALLY DEFERRED
		)
	`).Error)
	require.NoError(t, db.Use(&mysqlMetricsPlugin{}))

	store := &memDBStorage{mysqlDB: db}
	err = store.runMySQLTransaction(context.Background(), func(tx *gorm.DB) error {
		return tx.Exec("INSERT INTO children (id, parent_id) VALUES (1, 404)").Error
	})
	require.ErrorContains(t, err, "FOREIGN KEY constraint failed")
	require.Equal(t, float64(1), testutil.ToFloat64(
		metametrics.MySQLOperationsTotal.WithLabelValues(
			metametrics.MySQLOperationTransaction,
			metametrics.ResultFailure,
		),
	))
	require.Equal(t, uint64(1), mysqlOperationDurationCount(
		t,
		metametrics.MySQLOperationTransaction,
		metametrics.ResultFailure,
	))
	require.Equal(t, float64(0), testutil.ToFloat64(metametrics.MySQLLastOperationSuccess))
}

func TestRunMySQLTransactionDoesNotRecordCallbackPanicAsSuccess(t *testing.T) {
	resetMySQLMetrics(t)

	db, err := gorm.Open(
		sqlite.Open(filepath.Join(t.TempDir(), "mysql-transaction-panic-metrics.db")),
		&gorm.Config{Logger: logger.Default.LogMode(logger.Silent)},
	)
	require.NoError(t, err)

	store := &memDBStorage{mysqlDB: db}
	require.PanicsWithValue(t, "transaction callback panic", func() {
		_ = store.runMySQLTransaction(context.Background(), func(*gorm.DB) error {
			panic("transaction callback panic")
		})
	})

	for _, result := range []string{metametrics.ResultSuccess, metametrics.ResultFailure} {
		require.Zero(t, testutil.ToFloat64(
			metametrics.MySQLOperationsTotal.WithLabelValues(
				metametrics.MySQLOperationTransaction,
				result,
			),
		))
		require.Zero(t, mysqlOperationDurationCount(
			t,
			metametrics.MySQLOperationTransaction,
			result,
		))
	}
	require.Zero(t, testutil.ToFloat64(metametrics.MySQLLastOperationSuccess))
}

func mysqlOperationDurationCount(t *testing.T, operation, result string) uint64 {
	t.Helper()

	collector, ok := metametrics.MySQLOperationDuration.
		WithLabelValues(operation, result).(prometheus.Metric)
	require.True(t, ok)
	var metric dto.Metric
	require.NoError(t, collector.Write(&metric))
	return metric.GetHistogram().GetSampleCount()
}

func resetMySQLMetrics(t *testing.T) {
	t.Helper()

	operations := []string{
		metametrics.MySQLOperationQuery,
		metametrics.MySQLOperationTransaction,
	}
	results := []string{metametrics.ResultSuccess, metametrics.ResultFailure}
	reset := func() {
		for _, operation := range operations {
			for _, result := range results {
				metametrics.MySQLOperationsTotal.DeleteLabelValues(operation, result)
				metametrics.MySQLOperationDuration.DeleteLabelValues(operation, result)
			}
		}
		metametrics.MySQLLastOperationSuccess.Set(0)
	}
	reset()
	t.Cleanup(reset)
}