//go:build unit

package common

import (
	"errors"
	"testing"

	opv1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/controller-runtime/pkg/log/zap"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
)

func TestCountRestartPhaseMigration(t *testing.T) {
	tidbPhase := RestartPhase{
		Name:       "TiDB",
		Components: []cloudv1.ComponentType{cloudv1.ComponentTiDB},
		Reason:     ReasonRestartingTiDB,
	}
	tiproxyPhase := RestartPhase{
		Name:       "TiProxy",
		Components: []cloudv1.ComponentType{cloudv1.ComponentTiProxy},
		Reason:     ReasonRestartingTiProxy,
	}

	tcWithTiProxy := &opv1alpha1.TidbCluster{
		ObjectMeta: metav1.ObjectMeta{Name: "tc-1"},
		Spec: opv1alpha1.TidbClusterSpec{
			TiProxy: &opv1alpha1.TiProxySpec{Replicas: 1},
		},
	}
	tcWithZeroTiProxy := &opv1alpha1.TidbCluster{
		ObjectMeta: metav1.ObjectMeta{Name: "tc-2"},
		Spec: opv1alpha1.TidbClusterSpec{
			TiProxy: &opv1alpha1.TiProxySpec{Replicas: 0},
		},
	}

	logger := zap.New(zap.UseDevMode(true))

	tests := []struct {
		name            string
		allTCs          []*opv1alpha1.TidbCluster
		ph              RestartPhase
		counter         ComponentMigrationCounter
		wantDone        int
		wantTotal       int
		wantMigratable  bool
		wantErr         bool
	}{
		{
			name:           "tidb-only phase skips migration",
			allTCs:         []*opv1alpha1.TidbCluster{tcWithTiProxy},
			ph:             tidbPhase,
			counter:        func(_, _ string, _ *opv1alpha1.TidbCluster) (int, int, error) { return 0, 0, nil },
			wantDone:       0,
			wantTotal:      0,
			wantMigratable: false,
		},
		{
			name:           "zero-replica component skips migration",
			allTCs:         []*opv1alpha1.TidbCluster{tcWithZeroTiProxy},
			ph:             tiproxyPhase,
			counter:        func(_, _ string, _ *opv1alpha1.TidbCluster) (int, int, error) { return 0, 0, nil },
			wantDone:       0,
			wantTotal:      0,
			wantMigratable: false,
		},
		{
			name:   "migratable component with no pods found",
			allTCs: []*opv1alpha1.TidbCluster{tcWithTiProxy},
			ph:     tiproxyPhase,
			counter: func(_, _ string, _ *opv1alpha1.TidbCluster) (int, int, error) {
				return 0, 0, nil
			},
			wantDone:       0,
			wantTotal:      0,
			wantMigratable: true,
		},
		{
			name:   "partial migration in progress",
			allTCs: []*opv1alpha1.TidbCluster{tcWithTiProxy},
			ph:     tiproxyPhase,
			counter: func(_, _ string, _ *opv1alpha1.TidbCluster) (int, int, error) {
				return 1, 2, nil
			},
			wantDone:       1,
			wantTotal:      2,
			wantMigratable: true,
		},
		{
			name:   "migration complete",
			allTCs: []*opv1alpha1.TidbCluster{tcWithTiProxy},
			ph:     tiproxyPhase,
			counter: func(_, _ string, _ *opv1alpha1.TidbCluster) (int, int, error) {
				return 2, 2, nil
			},
			wantDone:       2,
			wantTotal:      2,
			wantMigratable: true,
		},
		{
			name:   "counter error propagates",
			allTCs: []*opv1alpha1.TidbCluster{tcWithTiProxy},
			ph:     tiproxyPhase,
			counter: func(_, _ string, _ *opv1alpha1.TidbCluster) (int, int, error) {
				return 0, 0, errors.New("list pods failed")
			},
			wantMigratable: true,
			wantErr:        true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			gotDone, gotTotal, gotMigratable, err := CountRestartPhaseMigration(logger, "ns", tt.allTCs, tt.ph, tt.counter)
			if (err != nil) != tt.wantErr {
				t.Fatalf("CountRestartPhaseMigration() error = %v, wantErr %v", err, tt.wantErr)
			}
			if gotDone != tt.wantDone {
				t.Fatalf("done = %d, want %d", gotDone, tt.wantDone)
			}
			if gotTotal != tt.wantTotal {
				t.Fatalf("total = %d, want %d", gotTotal, tt.wantTotal)
			}
			if gotMigratable != tt.wantMigratable {
				t.Fatalf("migratable = %v, want %v", gotMigratable, tt.wantMigratable)
			}
		})
	}
}