package mysql

import (
	"testing"
	"time"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/shiningrush/fastflow/pkg/mod"
	"github.com/stretchr/testify/require"
	gormDriver "gorm.io/driver/mysql"
	"gorm.io/gorm"
)

const (
	queryHasCmdColumn = `
		SELECT COUNT(*) FROM information_schema.columns
		WHERE table_schema = DATABASE()
		AND table_name = 'dag_instances'
		AND column_name = 'has_cmd'
	`
	queryHasCmdIndex = `
		SELECT COUNT(*) FROM information_schema.statistics
		WHERE table_schema = DATABASE()
		AND table_name = 'dag_instances'
		AND index_name = 'idx_worker_has_cmd'
	`
	addHasCmdColumn = `
		ALTER TABLE dag_instances
		ADD COLUMN has_cmd TINYINT AS (cmd IS NOT NULL) STORED
	`
	addHasCmdIndex = `
		CREATE INDEX idx_worker_has_cmd ON dag_instances(worker, has_cmd)
	`
)

func newMockStore(t *testing.T) (*Store, sqlmock.Sqlmock) {
	t.Helper()

	sqlDB, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual))
	require.NoError(t, err)
	t.Cleanup(func() {
		mock.ExpectClose()
		require.NoError(t, sqlDB.Close())
	})

	db, err := gorm.Open(gormDriver.New(gormDriver.Config{
		Conn:                      sqlDB,
		SkipInitializeWithVersion: true,
	}), &gorm.Config{})
	require.NoError(t, err)

	return &Store{
		opt:    &StoreOption{Timeout: time.Second},
		gormDB: db,
	}, mock
}

func TestListDagInstanceWithoutFilterTagsUsesHasCmdIndexColumn(t *testing.T) {
	store, mock := newMockStore(t)

	mock.ExpectBegin()
	mock.ExpectQuery("SELECT * FROM `dag_instances` WHERE worker = ? AND has_cmd = 1").
		WithArgs("dataflow-workflow-0").
		WillReturnRows(sqlmock.NewRows([]string{"id"}))
	mock.ExpectCommit()

	instances, err := store.ListDagInstanceWithoutFilterTags(&mod.ListDagInstanceInput{
		Worker: "dataflow-workflow-0",
		HasCmd: true,
	})
	require.NoError(t, err)
	require.Empty(t, instances)
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestMigrateHasCmdColumnCreatesMissingColumnAndIndex(t *testing.T) {
	store, mock := newMockStore(t)

	mock.ExpectQuery(queryHasCmdColumn).
		WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
	mock.ExpectExec(addHasCmdColumn).WillReturnResult(sqlmock.NewResult(0, 0))
	mock.ExpectQuery(queryHasCmdIndex).
		WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
	mock.ExpectExec(addHasCmdIndex).WillReturnResult(sqlmock.NewResult(0, 0))

	require.NoError(t, store.migrateHasCmdColumn(store.gormDB))
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestMigrateHasCmdColumnSkipsExistingColumnAndIndex(t *testing.T) {
	store, mock := newMockStore(t)

	mock.ExpectQuery(queryHasCmdColumn).
		WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
	mock.ExpectQuery(queryHasCmdIndex).
		WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))

	require.NoError(t, store.migrateHasCmdColumn(store.gormDB))
	require.NoError(t, mock.ExpectationsWereMet())
}