package nodegroup

import (
	"context"
	"errors"
	"strings"
	"testing"
	"time"

	"github.com/DATA-DOG/go-sqlmock"
	mysqlDriver "github.com/go-sql-driver/mysql"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	dbmocks "github.com/tidbcloud/cluster-service-ng/pkg/externalmocks/dbmocks"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
)

func newLockMockDB(t *testing.T) (*gorm.DB, sqlmock.Sqlmock) {
	t.Helper()
	sqlDB, mockDB, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
	require.NoError(t, err)

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

	t.Cleanup(func() {
		mockDB.ExpectClose()
		require.NoError(t, sqlDB.Close())
	})

	return gdb, mockDB
}

func TestDataplaneNodeInstanceGroupLockRepo_TryAcquire(t *testing.T) {
	ctx := context.Background()
	lease := time.Minute
	now := time.Now()
	tests := []struct {
		name      string
		dp        string
		spec      string
		primary   string
		owner     string
		setupMock func(*gorm.DB, sqlmock.Sqlmock, *dbmocks.Provider)
		wantCode  codes.Code
		wantNil   bool
		validate  func(*testing.T, *entity.DataplaneNodeInstanceGroupLock)
	}{
		{
			name:     "invalid args",
			dp:       "",
			spec:     "",
			primary:  "",
			owner:    "",
			wantCode: codes.InvalidArgument,
		},
		{
			name:    "create new lock",
			dp:      "dp-1",
			spec:    "8C16G",
			primary: "m5.2xlarge",
			owner:   "owner-1",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_locks.*").WithArgs("dp-1", "8c16g", "m5.2xlarge", 1).WillReturnRows(sqlmock.NewRows([]string{"dataplane_name"}))
				mockDB.ExpectExec("INSERT .*dataplane_node_instance_group_locks.*").WillReturnResult(sqlmock.NewResult(1, 1))
			},
			wantCode: codes.OK,
			validate: func(t *testing.T, lock *entity.DataplaneNodeInstanceGroupLock) {
				t.Helper()
				require.NotNil(t, lock)
				assert.Equal(t, "dp-1", lock.DataplaneName)
				assert.Equal(t, "8c16g", lock.NodeSpec)
				assert.Equal(t, "m5.2xlarge", lock.PrimaryInstance)
				assert.Equal(t, "owner-1", lock.OwnerID)
				assert.Equal(t, int64(1), lock.FencingToken)
				assert.True(t, lock.ExpiresAt.After(now))
			},
		},
		{
			name:    "renew owned lock",
			dp:      "dp-1",
			spec:    "8C16G",
			primary: "m5.2xlarge",
			owner:   "owner-1",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				rows := sqlmock.NewRows([]string{
					"dataplane_name",
					"node_spec",
					"primary_instance",
					"owner_id",
					"fencing_token",
					"expires_at",
					"created_at",
					"updated_at",
				}).AddRow("dp-1", "8c16g", "m5.2xlarge", "owner-1", int64(2), now.Add(time.Minute), now, now)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_locks.*").WithArgs("dp-1", "8c16g", "m5.2xlarge", 1).WillReturnRows(rows)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_locks.*").WillReturnResult(sqlmock.NewResult(0, 1))
			},
			wantCode: codes.OK,
			validate: func(t *testing.T, lock *entity.DataplaneNodeInstanceGroupLock) {
				t.Helper()
				require.NotNil(t, lock)
				assert.Equal(t, int64(3), lock.FencingToken)
				assert.Equal(t, "owner-1", lock.OwnerID)
				assert.True(t, lock.ExpiresAt.After(now))
			},
		},
		{
			name:    "lock held and not expired",
			dp:      "dp-1",
			spec:    "8C16G",
			primary: "m5.2xlarge",
			owner:   "owner-2",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				rows := sqlmock.NewRows([]string{
					"dataplane_name",
					"node_spec",
					"primary_instance",
					"owner_id",
					"fencing_token",
					"expires_at",
					"created_at",
					"updated_at",
				}).AddRow("dp-1", "8c16g", "m5.2xlarge", "owner-1", int64(2), now.Add(time.Minute), now, now)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_locks.*").WithArgs("dp-1", "8c16g", "m5.2xlarge", 1).WillReturnRows(rows)
			},
			wantCode: codes.OK,
			wantNil:  true,
		},
		{
			name:    "take expired lock",
			dp:      "dp-1",
			spec:    "8C16G",
			primary: "m5.2xlarge",
			owner:   "owner-2",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				rows := sqlmock.NewRows([]string{
					"dataplane_name",
					"node_spec",
					"primary_instance",
					"owner_id",
					"fencing_token",
					"expires_at",
					"created_at",
					"updated_at",
				}).AddRow("dp-1", "8c16g", "m5.2xlarge", "owner-1", int64(4), now.Add(-time.Minute), now.Add(-time.Minute), now.Add(-time.Minute))
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_locks.*").WithArgs("dp-1", "8c16g", "m5.2xlarge", 1).WillReturnRows(rows)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_locks.*").WillReturnResult(sqlmock.NewResult(0, 1))
			},
			wantCode: codes.OK,
			validate: func(t *testing.T, lock *entity.DataplaneNodeInstanceGroupLock) {
				t.Helper()
				require.NotNil(t, lock)
				assert.Equal(t, "owner-2", lock.OwnerID)
				assert.Equal(t, int64(5), lock.FencingToken)
				assert.True(t, lock.ExpiresAt.After(now))
			},
		},
		{
			name:    "duplicate create triggers retry",
			dp:      "dp-1",
			spec:    "8C16G",
			primary: "m5.2xlarge",
			owner:   "owner-3",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_locks.*").WithArgs("dp-1", "8c16g", "m5.2xlarge", 1).WillReturnRows(sqlmock.NewRows([]string{"dataplane_name"}))
				mockDB.ExpectExec("INSERT .*dataplane_node_instance_group_locks.*").WillReturnError(&mysqlDriver.MySQLError{Number: 1062, Message: "Duplicate entry"})
				rows := sqlmock.NewRows([]string{
					"dataplane_name",
					"node_spec",
					"primary_instance",
					"owner_id",
					"fencing_token",
					"expires_at",
					"created_at",
					"updated_at",
				}).AddRow("dp-1", "8c16g", "m5.2xlarge", "owner-1", int64(1), now.Add(-time.Minute), now.Add(-time.Minute), now.Add(-time.Minute))
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_locks.*").WithArgs("dp-1", "8c16g", "m5.2xlarge", 1).WillReturnRows(rows)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_locks.*").WillReturnResult(sqlmock.NewResult(0, 1))
			},
			wantCode: codes.OK,
			validate: func(t *testing.T, lock *entity.DataplaneNodeInstanceGroupLock) {
				t.Helper()
				require.NotNil(t, lock)
				assert.Equal(t, "owner-3", lock.OwnerID)
				assert.Equal(t, int64(2), lock.FencingToken)
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := dbmocks.NewProvider(t)
			var (
				gdb    *gorm.DB
				mockDB sqlmock.Sqlmock
			)
			if tt.setupMock != nil {
				gdb, mockDB = newLockMockDB(t)
				tt.setupMock(gdb, mockDB, provider)
			}

			repo := NewDataplaneNodeInstanceGroupLockRepo(provider)
			lock, err := repo.TryAcquire(ctx, tt.dp, tt.spec, tt.primary, tt.owner, lease)
			if tt.wantCode == codes.OK {
				require.NoError(t, err)
				if tt.wantNil {
					assert.Nil(t, lock)
				} else if tt.validate != nil {
					tt.validate(t, lock)
				}
			} else {
				require.Error(t, err)
				assert.Equal(t, tt.wantCode, status.Code(err))
				assert.Nil(t, lock)
			}

			if mockDB != nil {
				require.NoError(t, mockDB.ExpectationsWereMet())
			}
		})
	}
}

func TestDataplaneNodeInstanceGroupLockRepo_Release(t *testing.T) {
	ctx := context.Background()
	tests := []struct {
		name      string
		dp        string
		spec      string
		primary   string
		owner     string
		setupMock func(*gorm.DB, sqlmock.Sqlmock, *dbmocks.Provider)
		wantCode  codes.Code
	}{
		{
			name:     "invalid args",
			dp:       "",
			spec:     "",
			primary:  "",
			owner:    "",
			wantCode: codes.InvalidArgument,
		},
		{
			name:    "not held",
			dp:      "dp-1",
			spec:    "8C16G",
			primary: "m5.2xlarge",
			owner:   "owner-1",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("DELETE .*dataplane_node_instance_group_locks.*").WillReturnResult(sqlmock.NewResult(0, 0))
			},
			wantCode: codes.FailedPrecondition,
		},
		{
			name:    "release error",
			dp:      "dp-1",
			spec:    "8C16G",
			primary: "m5.2xlarge",
			owner:   "owner-1",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("DELETE .*dataplane_node_instance_group_locks.*").WillReturnError(errors.New("boom"))
			},
			wantCode: codes.Internal,
		},
		{
			name:    "success",
			dp:      "dp-1",
			spec:    "8C16G",
			primary: "m5.2xlarge",
			owner:   "owner-1",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("DELETE .*dataplane_node_instance_group_locks.*").WillReturnResult(sqlmock.NewResult(0, 1))
			},
			wantCode: codes.OK,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := dbmocks.NewProvider(t)
			var (
				gdb    *gorm.DB
				mockDB sqlmock.Sqlmock
			)
			if tt.setupMock != nil {
				gdb, mockDB = newLockMockDB(t)
				tt.setupMock(gdb, mockDB, provider)
			}

			repo := NewDataplaneNodeInstanceGroupLockRepo(provider)
			err := repo.Release(ctx, tt.dp, tt.spec, tt.primary, tt.owner)
			if tt.wantCode == codes.OK {
				require.NoError(t, err)
			} else {
				require.Error(t, err)
				assert.Equal(t, tt.wantCode, status.Code(err))
			}

			if mockDB != nil {
				require.NoError(t, mockDB.ExpectationsWereMet())
			}
		})
	}
}

func TestDataplaneNodeInstanceGroupLockRepo_Load(t *testing.T) {
	ctx := context.Background()
	now := time.Now()
	tests := []struct {
		name      string
		dp        string
		spec      string
		primary   string
		setupMock func(*gorm.DB, sqlmock.Sqlmock, *dbmocks.Provider)
		wantCode  codes.Code
		wantNil   bool
	}{
		{
			name:     "invalid args",
			dp:       "",
			spec:     "",
			primary:  "",
			wantCode: codes.InvalidArgument,
		},
		{
			name:    "not found",
			dp:      "dp-1",
			spec:    "8C16G",
			primary: "m5.2xlarge",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_locks.*").WithArgs("dp-1", "8c16g", "m5.2xlarge", 1).WillReturnRows(sqlmock.NewRows([]string{"dataplane_name"}))
			},
			wantCode: codes.OK,
			wantNil:  true,
		},
		{
			name:    "query error",
			dp:      "dp-2",
			spec:    "4C8G",
			primary: "m5.xlarge",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_locks.*").WithArgs("dp-2", "4c8g", "m5.xlarge", 1).WillReturnError(errors.New("boom"))
			},
			wantCode: codes.Internal,
		},
		{
			name:    "success",
			dp:      "dp-3",
			spec:    "16C32G",
			primary: "m5.4xlarge",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				rows := sqlmock.NewRows([]string{
					"dataplane_name",
					"node_spec",
					"primary_instance",
					"owner_id",
					"fencing_token",
					"expires_at",
					"created_at",
					"updated_at",
				}).AddRow("dp-3", "16c32g", "m5.4xlarge", "owner-1", int64(2), now.Add(time.Minute), now, now)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_locks.*").WithArgs("dp-3", "16c32g", "m5.4xlarge", 1).WillReturnRows(rows)
			},
			wantCode: codes.OK,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := dbmocks.NewProvider(t)
			var (
				gdb    *gorm.DB
				mockDB sqlmock.Sqlmock
			)
			if tt.setupMock != nil {
				gdb, mockDB = newLockMockDB(t)
				tt.setupMock(gdb, mockDB, provider)
			}

			repo := NewDataplaneNodeInstanceGroupLockRepo(provider)
			lock, err := repo.Load(ctx, tt.dp, tt.spec, tt.primary)
			if tt.wantCode == codes.OK {
				require.NoError(t, err)
				if tt.wantNil {
					assert.Nil(t, lock)
				} else {
					require.NotNil(t, lock)
					assert.Equal(t, tt.dp, lock.DataplaneName)
					assert.Equal(t, strings.ToLower(tt.spec), lock.NodeSpec)
				}
			} else {
				require.Error(t, err)
				assert.Equal(t, tt.wantCode, status.Code(err))
				assert.Nil(t, lock)
			}

			if mockDB != nil {
				require.NoError(t, mockDB.ExpectationsWereMet())
			}
		})
	}
}

func TestDataplaneNodeInstanceGroupLockRepo_ForceRelease(t *testing.T) {
	ctx := context.Background()
	tests := []struct {
		name      string
		dp        string
		setupMock func(*gorm.DB, sqlmock.Sqlmock, *dbmocks.Provider)
		wantCode  codes.Code
	}{
		{
			name:     "invalid args",
			dp:       "",
			wantCode: codes.InvalidArgument,
		},
		{
			name: "delete error",
			dp:   "dp-1",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("DELETE .*dataplane_node_instance_group_locks.*").WillReturnError(errors.New("boom"))
			},
			wantCode: codes.Internal,
		},
		{
			name: "success - deletes all nodeSpec locks",
			dp:   "dp-2",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("DELETE .*dataplane_node_instance_group_locks.*").WillReturnResult(sqlmock.NewResult(0, 3))
			},
			wantCode: codes.OK,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := dbmocks.NewProvider(t)
			var (
				gdb    *gorm.DB
				mockDB sqlmock.Sqlmock
			)
			if tt.setupMock != nil {
				gdb, mockDB = newLockMockDB(t)
				tt.setupMock(gdb, mockDB, provider)
			}

			repo := NewDataplaneNodeInstanceGroupLockRepo(provider)
			err := repo.ForceRelease(ctx, tt.dp)
			if tt.wantCode == codes.OK {
				require.NoError(t, err)
			} else {
				require.Error(t, err)
				assert.Equal(t, tt.wantCode, status.Code(err))
			}

			if mockDB != nil {
				require.NoError(t, mockDB.ExpectationsWereMet())
			}
		})
	}
}

func TestDataplaneNodeInstanceGroupLockRepo_IsLockValid(t *testing.T) {
	ctx := context.Background()
	now := time.Now()
	tests := []struct {
		name      string
		dp        string
		spec      string
		primary   string
		fencing   int64
		setupMock func(*gorm.DB, sqlmock.Sqlmock, *dbmocks.Provider)
		wantCode  codes.Code
		wantValid bool
	}{
		{
			name:     "invalid args",
			dp:       "",
			spec:     "",
			primary:  "",
			wantCode: codes.InvalidArgument,
		},
		{
			name:    "not found",
			dp:      "dp-1",
			spec:    "8C16G",
			primary: "m5.2xlarge",
			fencing: 1,
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_locks.*").WithArgs("dp-1", "8c16g", "m5.2xlarge", 1).WillReturnRows(sqlmock.NewRows([]string{"dataplane_name"}))
			},
			wantCode:  codes.OK,
			wantValid: false,
		},
		{
			name:    "query error",
			dp:      "dp-2",
			spec:    "4C8G",
			primary: "m5.xlarge",
			fencing: 1,
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_locks.*").WithArgs("dp-2", "4c8g", "m5.xlarge", 1).WillReturnError(errors.New("boom"))
			},
			wantCode: codes.Internal,
		},
		{
			name:    "valid lock",
			dp:      "dp-3",
			spec:    "8C16G",
			primary: "m5.2xlarge",
			fencing: 2,
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				rows := sqlmock.NewRows([]string{
					"dataplane_name",
					"node_spec",
					"primary_instance",
					"owner_id",
					"fencing_token",
					"expires_at",
					"created_at",
					"updated_at",
				}).AddRow("dp-3", "8c16g", "m5.2xlarge", "owner-1", int64(2), now.Add(time.Minute), now, now)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_locks.*").WithArgs("dp-3", "8c16g", "m5.2xlarge", 1).WillReturnRows(rows)
			},
			wantCode:  codes.OK,
			wantValid: true,
		},
		{
			name:    "expired or token mismatch",
			dp:      "dp-4",
			spec:    "16C32G",
			primary: "m5.4xlarge",
			fencing: 2,
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				rows := sqlmock.NewRows([]string{
					"dataplane_name",
					"node_spec",
					"primary_instance",
					"owner_id",
					"fencing_token",
					"expires_at",
					"created_at",
					"updated_at",
				}).AddRow("dp-4", "16c32g", "m5.4xlarge", "owner-1", int64(3), now.Add(-time.Minute), now, now)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_locks.*").WithArgs("dp-4", "16c32g", "m5.4xlarge", 1).WillReturnRows(rows)
			},
			wantCode:  codes.OK,
			wantValid: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := dbmocks.NewProvider(t)
			var (
				gdb    *gorm.DB
				mockDB sqlmock.Sqlmock
			)
			if tt.setupMock != nil {
				gdb, mockDB = newLockMockDB(t)
				tt.setupMock(gdb, mockDB, provider)
			}

			repo := NewDataplaneNodeInstanceGroupLockRepo(provider)
			valid, err := repo.IsLockValid(ctx, tt.dp, tt.spec, tt.primary, tt.fencing)
			if tt.wantCode == codes.OK {
				require.NoError(t, err)
				assert.Equal(t, tt.wantValid, valid)
			} else {
				require.Error(t, err)
				assert.Equal(t, tt.wantCode, status.Code(err))
				assert.False(t, valid)
			}

			if mockDB != nil {
				require.NoError(t, mockDB.ExpectationsWereMet())
			}
		})
	}
}