package repo

import (
	"context"
	"testing"
	"time"

	ccent "github.com/tidbcloud/devops/entity/cloudconfig"
	prem "github.com/tidbcloud/devops/entity/premiumupgrade"
	premmod "github.com/tidbcloud/devops/models/premiumupgrade"
	testutil "github.com/tidbcloud/devops/pkg/infra/utils/test"

	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
	"go.uber.org/mock/gomock"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

func TestRepo(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Premium Upgrade Repository Suite")
}

var _ = Describe("Premium upgrade repository", func() {
	var (
		ctx      context.Context
		gormDB   *gorm.DB
		repo     Repo
		provider *testutil.MockTransProvider
	)

	BeforeEach(func() {
		ctx = context.Background()
		var err error
		gormDB, err = gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
		Expect(err).NotTo(HaveOccurred())
		Expect(gormDB.AutoMigrate(&premmod.UpgradePlan{}, &premmod.UpgradeTarget{}, &premmod.UpgradeStep{})).To(Succeed())

		ctrl := gomock.NewController(GinkgoT())
		provider = testutil.NewMockTransProvider(ctrl)
		provider.EXPECT().ReadDB(gomock.Any()).Return(gormDB).AnyTimes()
		provider.EXPECT().WriteDB(gomock.Any()).Return(gormDB).AnyTimes()
		provider.EXPECT().Transaction(gomock.Any(), gomock.Any()).DoAndReturn(
			func(ctx context.Context, callback func(context.Context) error) error { return callback(ctx) },
		).AnyTimes()
		repo = NewRepo(provider)
	})

	It("creates a plan with targets and steps in one repository operation", func() {
		created, err := repo.Create(ctx, &prem.UpgradePlan{
			Name: "rollout", Reason: "align versions", Applicant: "owner@example.com",
			OperatorEntity: "user:owner@example.com", DesiredVersion: "v8.5.3",
			ExecutionMode: prem.ExecutionModeStepByStep,
			ScopeMode:     prem.ScopeModeSelected, RequestedClusterIDs: []string{"10"},
			RequestedComponents: []string{"LOGICAL_CLUSTER/tidb"}, ApprovalManifest: []byte(`{"targets":[]}`),
			ManifestHash: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
			Status:       prem.PlanStatusApprovalPendingCreate, MaxParallelTargets: 2,
			ComponentOrderRevision: "v1", IdempotencyKey: "key-1",
			Targets: []*prem.UpgradeTarget{{
				ClusterType: ccent.ClusterTypeLogicalCluster, ClusterID: "10",
				SourceClusterIDs: []string{"10"}, ClusterName: "cluster-10", TenantIDs: []string{"100"},
				Status: prem.TargetStatusPending,
				Steps: []*prem.UpgradeStep{{
					Sequence: 1, ApplyGroup: "tidb",
					Components: []prem.UpgradeStepComponent{{
						Component: "tidb", FromVersion: "v8.5.2", ToVersion: "v8.5.3",
						FromImageTag: "tidb-v8.5.2-build.1", ToImageTag: "tidb-v8.5.3-build.2",
					}},
					ApplyStrategy: "RollingRestart", ApplyMode: []byte("null"), Status: prem.StepStatusPending,
				}},
			}},
		})
		Expect(err).NotTo(HaveOccurred())
		Expect(created.ID).NotTo(BeZero())
		Expect(created.Targets).To(HaveLen(1))
		Expect(created.Targets[0].PlanID).To(Equal(created.ID))
		Expect(created.Targets[0].Steps).To(HaveLen(1))
		Expect(created.Targets[0].Steps[0].PlanID).To(Equal(created.ID))
		Expect(created.Targets[0].Steps[0].TargetID).To(Equal(created.Targets[0].ID))
		Expect(created.Targets[0].Steps[0].Components[0].FromImageTag).To(Equal("tidb-v8.5.2-build.1"))
		Expect(created.Targets[0].Steps[0].Components[0].ToImageTag).To(Equal("tidb-v8.5.3-build.2"))
		Expect(created.ExecutionMode).To(Equal(prem.ExecutionModeStepByStep))
	})

	It("claims only a waiting manual step whose plan is running", func() {
		created, err := repo.Create(ctx, &prem.UpgradePlan{
			Name: "manual", Reason: "test", Applicant: "owner@example.com",
			OperatorEntity: "user:owner@example.com", DesiredVersion: "v8.5.3",
			ExecutionMode: prem.ExecutionModeStepByStep, ScopeMode: prem.ScopeModeSelected,
			RequestedClusterIDs: []string{"10"}, RequestedComponents: []string{"LOGICAL_CLUSTER/tidb"},
			ApprovalManifest: []byte(`{}`), ManifestHash: "manual", Status: prem.PlanStatusRunning,
			MaxParallelTargets: 1, ComponentOrderRevision: "v1", IdempotencyKey: "manual",
			Targets: []*prem.UpgradeTarget{{
				ClusterType: ccent.ClusterTypeLogicalCluster, ClusterID: "10", SourceClusterIDs: []string{"10"},
				ClusterName: "cluster", TenantIDs: []string{"100"}, Status: prem.TargetStatusRunning,
				Steps: []*prem.UpgradeStep{{
					Sequence: 1, ApplyGroup: "tidb",
					Components:    []prem.UpgradeStepComponent{{Component: "tidb", ToVersion: "v8.5.3"}},
					ApplyStrategy: "RollingRestart", ApplyMode: []byte("null"),
					Status: prem.StepStatusWaitingManualApply, ConfigurationChangeID: pointerUint64(42),
				}},
			}},
		})
		Expect(err).NotTo(HaveOccurred())
		step := created.Targets[0].Steps[0]
		now := time.Now()

		claimed, err := repo.ClaimStepForManualApply(ctx, step.ID, step.Version, now, now.Add(time.Hour))
		Expect(err).NotTo(HaveOccurred())
		Expect(claimed).To(BeTrue())

		refreshed, err := repo.Get(ctx, created.ID)
		Expect(err).NotTo(HaveOccurred())
		Expect(refreshed.Targets[0].Steps[0].Status).To(Equal(prem.StepStatusApplying))
		Expect(refreshed.Targets[0].Steps[0].StartedAt).NotTo(BeNil())
		Expect(refreshed.Targets[0].Steps[0].ApplyDeadlineAt).NotTo(BeNil())
	})

	It("enforces one configuration change per step", func() {
		step := &premmod.UpgradeStep{
			PlanID: 1, TargetID: 1, Sequence: 1, ApplyGroup: "tidb",
			Components:    []premmod.UpgradeStepComponent{{Component: "tidb", FromVersion: "v8.5.2", ToVersion: "v8.5.3"}},
			ApplyStrategy: "RollingRestart",
			ApplyMode:     []byte("null"), Status: string(prem.StepStatusApplying), ConfigurationChangeID: pointerUint64(42),
		}
		Expect(gormDB.Create(step).Error).To(Succeed())
		duplicate := *step
		duplicate.ID = 0
		duplicate.TargetID = 2
		Expect(gormDB.Create(&duplicate).Error).To(HaveOccurred())
	})

	It("allows multiple apply groups in one target sequence and rejects a duplicate group", func() {
		first := &premmod.UpgradeStep{
			PlanID: 1, TargetID: 1, Sequence: 9, ApplyGroup: "shared_pool_tiproxy",
			Components:    []premmod.UpgradeStepComponent{{Component: "tiproxy", ToVersion: "v8.5.3"}},
			ApplyStrategy: "RollingRestart", ApplyMode: []byte("null"), Status: string(prem.StepStatusPending),
		}
		Expect(gormDB.Create(first).Error).To(Succeed())
		second := &premmod.UpgradeStep{
			PlanID: 1, TargetID: 1, Sequence: 9, ApplyGroup: "data_migration",
			Components:    []premmod.UpgradeStepComponent{{Component: "data_migration", ToVersion: "v8.5.3"}},
			ApplyStrategy: "RollingRestart", ApplyMode: []byte("null"), Status: string(prem.StepStatusPending),
		}
		Expect(gormDB.Create(second).Error).To(Succeed())
		duplicateGroup := &premmod.UpgradeStep{
			PlanID: 1, TargetID: 1, Sequence: 10, ApplyGroup: "shared_pool_tiproxy",
			Components:    []premmod.UpgradeStepComponent{{Component: "tiproxy", ToVersion: "v8.5.4"}},
			ApplyStrategy: "RollingRestart", ApplyMode: []byte("null"), Status: string(prem.StepStatusPending),
		}
		Expect(gormDB.Create(duplicateGroup).Error).To(HaveOccurred())
	})

	It("finds active plans that overlap an actual locator", func() {
		created, err := repo.Create(ctx, &prem.UpgradePlan{
			Name: "rollout", Reason: "align versions", Applicant: "owner@example.com",
			OperatorEntity: "user:owner@example.com", DesiredVersion: "v8.5.3",
			ScopeMode: prem.ScopeModeSelected, RequestedClusterIDs: []string{"10"}, RequestedComponents: []string{"SHARED_POOL/tikv"},
			ApprovalManifest: []byte(`{}`), ManifestHash: "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc",
			Status: prem.PlanStatusRunning, MaxParallelTargets: 1, ComponentOrderRevision: "v1", IdempotencyKey: "key-overlap",
			Targets: []*prem.UpgradeTarget{{
				ClusterType: ccent.ClusterTypeSharedPool, ClusterID: "200", SourceClusterIDs: []string{"10"},
				ClusterName: "pool", TenantIDs: []string{"100"}, Status: prem.TargetStatusRunning,
			}},
		})
		Expect(err).NotTo(HaveOccurred())

		plans, err := repo.FindActiveByTargetLocators(ctx, []ccent.ClusterConfigurationLocator{{ClusterType: ccent.ClusterTypeSharedPool, ClusterID: "200"}})
		Expect(err).NotTo(HaveOccurred())
		Expect(plans).To(HaveLen(1))
		Expect(plans[0].ID).To(Equal(created.ID))
		Expect(plans[0].Targets).To(HaveLen(1))
		Expect(plans[0].Targets[0].ClusterID).To(Equal("200"))

		Expect(repo.PatchPlan(ctx, created.ID, map[string]any{"status": string(prem.PlanStatusCompleted)})).To(Succeed())
		plans, err = repo.FindActiveByTargetLocators(ctx, []ccent.ClusterConfigurationLocator{{ClusterType: ccent.ClusterTypeSharedPool, ClusterID: "200"}})
		Expect(err).NotTo(HaveOccurred())
		Expect(plans).To(BeEmpty())
	})

	It("filters plans by an exact tenant id", func() {
		createForTenant := func(key, clusterID, tenantID string) *prem.UpgradePlan {
			created, err := repo.Create(ctx, &prem.UpgradePlan{
				Name: key, Reason: "test", Applicant: "owner@example.com", OperatorEntity: "user:owner@example.com",
				DesiredVersion: "v8.5.3", ScopeMode: prem.ScopeModeSelected,
				RequestedClusterIDs: []string{clusterID}, RequestedComponents: []string{"LOGICAL_CLUSTER/tidb"},
				ApprovalManifest: []byte(`{}`), ManifestHash: key, Status: prem.PlanStatusRunning,
				MaxParallelTargets: 1, ComponentOrderRevision: "v1", IdempotencyKey: key,
				Targets: []*prem.UpgradeTarget{{
					ClusterType: ccent.ClusterTypeLogicalCluster, ClusterID: clusterID, SourceClusterIDs: []string{clusterID},
					ClusterName: clusterID, TenantIDs: []string{tenantID}, Status: prem.TargetStatusPending,
				}},
			})
			Expect(err).NotTo(HaveOccurred())
			return created
		}
		tenant12 := createForTenant("tenant-12", "12", "12")
		_ = createForTenant("tenant-112", "112", "112")
		id := uint64(12)

		plans, total, err := repo.List(ctx, ListFilter{TenantID: &id})
		Expect(err).NotTo(HaveOccurred())
		Expect(total).To(Equal(int64(1)))
		Expect(plans).To(HaveLen(1))
		Expect(plans[0].ID).To(Equal(tenant12.ID))
	})

	It("reopens a terminal plan and its failure-skipped steps for retry", func() {
		created, err := repo.Create(ctx, &prem.UpgradePlan{
			Name: "retry", Reason: "test", Applicant: "owner@example.com", OperatorEntity: "user:owner@example.com",
			DesiredVersion: "v8.5.3", ScopeMode: prem.ScopeModeSelected,
			RequestedClusterIDs: []string{"10"}, RequestedComponents: []string{"SHARED_POOL/tikv", "LOGICAL_CLUSTER/tidb"},
			ApprovalManifest: []byte(`{}`), ManifestHash: "retry", Status: prem.PlanStatusFailed,
			MaxParallelTargets: 1, ComponentOrderRevision: "v1", IdempotencyKey: "retry",
			Targets: []*prem.UpgradeTarget{
				{
					ClusterType: ccent.ClusterTypeSharedPool, ClusterID: "200", SourceClusterIDs: []string{"10"},
					ClusterName: "pool", TenantIDs: []string{"100"},
					Status: prem.TargetStatusFailed,
					Steps: []*prem.UpgradeStep{
						{Sequence: 1, ApplyGroup: "tikv", Components: []prem.UpgradeStepComponent{{Component: "tikv", FromVersion: "v8.5.2", ToVersion: "v8.5.3"}}, ApplyStrategy: "RollingRestart", ApplyMode: []byte("null"), Status: prem.StepStatusFailed},
						{Sequence: 2, ApplyGroup: "tiflash", Components: []prem.UpgradeStepComponent{{Component: "tiflash", FromVersion: "v8.5.2", ToVersion: "v8.5.3"}}, ApplyStrategy: "RollingRestart", ApplyMode: []byte("null"), Status: prem.StepStatusSkipped, StatusReason: "stopped after apply group tikv failed"},
					},
				},
				{
					ClusterType: ccent.ClusterTypeLogicalCluster, ClusterID: "10", SourceClusterIDs: []string{"10"},
					ClusterName: "cluster", TenantIDs: []string{"100"},
					Status: prem.TargetStatusFailed,
					Steps:  []*prem.UpgradeStep{{Sequence: 3, ApplyGroup: "tidb", Components: []prem.UpgradeStepComponent{{Component: "tidb", FromVersion: "v8.5.2", ToVersion: "v8.5.3"}}, ApplyStrategy: "RollingRestart", ApplyMode: []byte("null"), Status: prem.StepStatusSkipped, StatusReason: "stopped after shared pool apply group tikv failed"}},
				},
			},
		})
		Expect(err).NotTo(HaveOccurred())
		shared, logical := created.Targets[0], created.Targets[1]

		Expect(repo.ResetFailedStepForRetry(ctx, RetryReset{
			PlanID: created.ID, TargetID: shared.ID, StepID: shared.Steps[0].ID,
			LaterStepIDs: []uint64{shared.Steps[1].ID}, DependentTargetIDs: []uint64{logical.ID},
			DependentStepIDs: []uint64{logical.Steps[0].ID},
		})).To(Succeed())

		refreshed, err := repo.Get(ctx, created.ID)
		Expect(err).NotTo(HaveOccurred())
		Expect(refreshed.Status).To(Equal(prem.PlanStatusRunning))
		Expect(refreshed.Targets[0].Status).To(Equal(prem.TargetStatusPending))
		Expect(refreshed.Targets[0].Steps[0].Status).To(Equal(prem.StepStatusPending))
		Expect(refreshed.Targets[0].Steps[1].Status).To(Equal(prem.StepStatusPending))
		Expect(refreshed.Targets[1].Status).To(Equal(prem.TargetStatusPending))
		Expect(refreshed.Targets[1].Steps[0].Status).To(Equal(prem.StepStatusPending))
	})
})

func pointerUint64(value uint64) *uint64 { return &value }