//go:generate ${GOBIN}/mockgen -destination repo_mock.go -package=repo ${GO_MODULE}/pkg/service/premiumupgrade/repo Repo
package repo

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"time"

	apprent "github.com/tidbcloud/devops/entity/approval"
	ccent "github.com/tidbcloud/devops/entity/cloudconfig"
	prem "github.com/tidbcloud/devops/entity/premiumupgrade"
	"github.com/tidbcloud/devops/models/base"
	premmod "github.com/tidbcloud/devops/models/premiumupgrade"
	"github.com/tidbcloud/devops/pkg/infra/utils/repoopt"
	"github.com/tidbcloud/pkgv2/db"
	"gorm.io/gorm"
)

type ListFilter struct {
	Applicant string
	TenantID  *uint64
	Statuses  []prem.PlanStatus
	Page      int
	PerPage   int
}

type RetryReset struct {
	PlanID             uint64
	TargetID           uint64
	StepID             uint64
	LaterStepIDs       []uint64
	DependentTargetIDs []uint64
	DependentStepIDs   []uint64
}

type Repo interface {
	Create(ctx context.Context, plan *prem.UpgradePlan) (*prem.UpgradePlan, error)
	Get(ctx context.Context, id uint64) (*prem.UpgradePlan, error)
	GetByApprovalID(ctx context.Context, approvalID uint64) (*prem.UpgradePlan, error)
	GetByIdempotencyKey(ctx context.Context, applicant, key string) (*prem.UpgradePlan, error)
	FindActiveByTargetLocators(ctx context.Context, locators []ccent.ClusterConfigurationLocator) ([]*prem.UpgradePlan, error)
	List(ctx context.Context, filter ListFilter) ([]*prem.UpgradePlan, int64, error)
	ListByStatus(ctx context.Context, statuses ...prem.PlanStatus) ([]*prem.UpgradePlan, error)
	PatchPlan(ctx context.Context, id uint64, values map[string]any) error
	PatchTarget(ctx context.Context, id uint64, values map[string]any) error
	ClaimPlanStatus(ctx context.Context, id uint64, from, to prem.PlanStatus, version uint64) (bool, error)
	ClaimStep(ctx context.Context, id uint64, from []prem.StepStatus, version uint64) (bool, error)
	ClaimStepForExecution(ctx context.Context, id uint64, from []prem.StepStatus, version uint64) (bool, error)
	ClaimStepForManualApply(ctx context.Context, id, version uint64, startedAt, deadline time.Time) (bool, error)
	PatchStep(ctx context.Context, id uint64, values map[string]any) error
	ResetFailedStepForRetry(ctx context.Context, reset RetryReset) error
}

type repo struct{ dbp db.TransProvider }

func NewRepo(dbp db.TransProvider) Repo { return &repo{dbp: dbp} }

func (r *repo) Create(ctx context.Context, plan *prem.UpgradePlan) (*prem.UpgradePlan, error) {
	model := planToModel(plan)
	err := r.dbp.Transaction(ctx, func(txCtx context.Context) error {
		db := r.dbp.WriteDB(txCtx)
		if err := db.Omit("Targets").Create(model).Error; err != nil {
			return err
		}
		for _, target := range model.Targets {
			target.PlanID = model.ID.Uint64()
			if err := db.Omit("Steps").Create(target).Error; err != nil {
				return err
			}
			for _, step := range target.Steps {
				step.PlanID = model.ID.Uint64()
				step.TargetID = target.ID.Uint64()
				if err := db.Create(step).Error; err != nil {
					return err
				}
			}
		}
		return nil
	})
	if err != nil {
		return nil, fmt.Errorf("create premium upgrade plan: %w", err)
	}
	return r.Get(ctx, model.ID.Uint64())
}

func (r *repo) Get(ctx context.Context, id uint64) (*prem.UpgradePlan, error) {
	var model premmod.UpgradePlan
	err := r.dbp.ReadDB(ctx).
		Preload("Targets", func(db *gorm.DB) *gorm.DB { return db.Order("id") }).
		Preload("Targets.Steps", func(db *gorm.DB) *gorm.DB { return db.Order("sequence, id") }).
		First(&model, id).Error
	if err != nil {
		return nil, err
	}
	return planToEntity(&model), nil
}

func (r *repo) GetByApprovalID(ctx context.Context, approvalID uint64) (*prem.UpgradePlan, error) {
	var model premmod.UpgradePlan
	err := r.dbp.ReadDB(ctx).Where("approval_id = ?", approvalID).First(&model).Error
	if errors.Is(err, gorm.ErrRecordNotFound) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	return r.Get(ctx, model.ID.Uint64())
}

func (r *repo) GetByIdempotencyKey(ctx context.Context, applicant, key string) (*prem.UpgradePlan, error) {
	var model premmod.UpgradePlan
	err := r.dbp.ReadDB(ctx).
		Where("applicant = ? AND idempotency_key = ?", applicant, key).
		First(&model).Error
	if errors.Is(err, gorm.ErrRecordNotFound) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	return r.Get(ctx, model.ID.Uint64())
}

func (r *repo) FindActiveByTargetLocators(ctx context.Context, locators []ccent.ClusterConfigurationLocator) ([]*prem.UpgradePlan, error) {
	if len(locators) == 0 {
		return nil, nil
	}
	nonTerminal := make([]string, 0, len(prem.PlanStatusValues))
	for _, status := range prem.PlanStatusValues {
		if !prem.IsPlanTerminal(status) {
			nonTerminal = append(nonTerminal, string(status))
		}
	}
	query := r.dbp.ReadDB(ctx).
		Model(&premmod.UpgradePlan{}).
		Preload("Targets", func(db *gorm.DB) *gorm.DB { return db.Order("id") }).
		Joins("JOIN premium_upgrade_targets ON premium_upgrade_targets.plan_id = premium_upgrade_plans.id").
		Where("premium_upgrade_plans.status IN ?", nonTerminal)
	locatorQuery := r.dbp.ReadDB(ctx).Where("1 = 0")
	for _, locator := range locators {
		locatorQuery = locatorQuery.Or(
			"(premium_upgrade_targets.cluster_type = ? AND premium_upgrade_targets.cluster_id = ?)",
			string(locator.ClusterType), locator.ClusterID,
		)
	}
	var models []*premmod.UpgradePlan
	if err := query.Where(locatorQuery).Distinct("premium_upgrade_plans.*").Find(&models).Error; err != nil {
		return nil, err
	}
	items := make([]*prem.UpgradePlan, 0, len(models))
	for _, model := range models {
		items = append(items, planToEntity(model))
	}
	return items, nil
}

func (r *repo) List(ctx context.Context, filter ListFilter) ([]*prem.UpgradePlan, int64, error) {
	query := r.dbp.ReadDB(ctx).Model(&premmod.UpgradePlan{})
	if filter.TenantID != nil {
		query = query.Joins("JOIN premium_upgrade_targets ON premium_upgrade_targets.plan_id = premium_upgrade_plans.id")
		tenantID := fmt.Sprintf("%d", *filter.TenantID)
		if query.Dialector.Name() == "mysql" {
			query = query.Where("JSON_CONTAINS(premium_upgrade_targets.tenant_ids, ?)", strconv.Quote(tenantID))
		} else {
			// SQLite repository tests use a JSON string array and do not provide
			// MySQL JSON_CONTAINS. Quoting prevents 12 from matching 112.
			query = query.Where("premium_upgrade_targets.tenant_ids LIKE ?", "%\""+tenantID+"\"%")
		}
		query = query.Distinct("premium_upgrade_plans.*")
	}
	if filter.Applicant != "" {
		query = query.Where("premium_upgrade_plans.applicant = ?", filter.Applicant)
	}
	if len(filter.Statuses) > 0 {
		values := make([]string, 0, len(filter.Statuses))
		for _, status := range filter.Statuses {
			values = append(values, string(status))
		}
		query = query.Where("premium_upgrade_plans.status IN ?", values)
	}
	var total int64
	if err := query.Distinct("premium_upgrade_plans.id").Count(&total).Error; err != nil {
		return nil, 0, err
	}
	page, perPage := filter.Page, filter.PerPage
	if page <= 0 {
		page = 1
	}
	if perPage <= 0 || perPage > 100 {
		perPage = 20
	}
	query = repoopt.Aggregate(
		&repoopt.Pager{Page: page, PerPage: perPage},
		&repoopt.Order{Column: "premium_upgrade_plans.created_at", Desc: true},
	).Apply(query)
	query = query.Select("premium_upgrade_plans.*")
	var models []*premmod.UpgradePlan
	if err := query.Find(&models).Error; err != nil {
		return nil, 0, err
	}
	items := make([]*prem.UpgradePlan, 0, len(models))
	for _, model := range models {
		items = append(items, planToEntity(model))
	}
	return items, total, nil
}

func (r *repo) ListByStatus(ctx context.Context, statuses ...prem.PlanStatus) ([]*prem.UpgradePlan, error) {
	values := make([]string, 0, len(statuses))
	for _, status := range statuses {
		values = append(values, string(status))
	}
	var models []*premmod.UpgradePlan
	query := r.dbp.ReadDB(ctx).Model(&premmod.UpgradePlan{})
	if len(values) > 0 {
		query = query.Where("status IN ?", values)
	}
	if err := query.Find(&models).Error; err != nil {
		return nil, err
	}
	items := make([]*prem.UpgradePlan, 0, len(models))
	for _, model := range models {
		items = append(items, planToEntity(model))
	}
	return items, nil
}

func (r *repo) PatchPlan(ctx context.Context, id uint64, values map[string]any) error {
	return r.dbp.WriteDB(ctx).Model(&premmod.UpgradePlan{}).Where("id = ?", id).Updates(values).Error
}

func (r *repo) PatchTarget(ctx context.Context, id uint64, values map[string]any) error {
	return r.dbp.WriteDB(ctx).Model(&premmod.UpgradeTarget{}).Where("id = ?", id).Updates(values).Error
}

func (r *repo) ClaimPlanStatus(ctx context.Context, id uint64, from, to prem.PlanStatus, version uint64) (bool, error) {
	result := r.dbp.WriteDB(ctx).Model(&premmod.UpgradePlan{}).
		Where("id = ? AND status = ? AND version = ?", id, from, version).
		Updates(map[string]any{"status": string(to), "version": gorm.Expr("version + 1")})
	return result.RowsAffected == 1, result.Error
}

func (r *repo) ClaimStep(ctx context.Context, id uint64, from []prem.StepStatus, version uint64) (bool, error) {
	return r.claimStep(ctx, id, from, version, false)
}

func (r *repo) ClaimStepForExecution(ctx context.Context, id uint64, from []prem.StepStatus, version uint64) (bool, error) {
	return r.claimStep(ctx, id, from, version, true)
}

func (r *repo) ClaimStepForManualApply(ctx context.Context, id, version uint64, startedAt, deadline time.Time) (bool, error) {
	result := r.dbp.WriteDB(ctx).Model(&premmod.UpgradeStep{}).
		Where("id = ? AND status = ? AND version = ?", id, string(prem.StepStatusWaitingManualApply), version).
		Where(
			"plan_id IN (?)",
			r.dbp.WriteDB(ctx).Model(&premmod.UpgradePlan{}).Select("id").Where("status = ?", string(prem.PlanStatusRunning)),
		).
		Updates(map[string]any{
			"status":            string(prem.StepStatusApplying),
			"version":           gorm.Expr("version + 1"),
			"started_at":        startedAt,
			"apply_deadline_at": deadline,
			"next_retry_at":     nil,
			"status_reason":     "",
		})
	return result.RowsAffected == 1, result.Error
}

func (r *repo) claimStep(ctx context.Context, id uint64, from []prem.StepStatus, version uint64, requireRunnablePlan bool) (bool, error) {
	values := make([]string, 0, len(from))
	for _, status := range from {
		values = append(values, string(status))
	}
	query := r.dbp.WriteDB(ctx).Model(&premmod.UpgradeStep{}).
		Where("id = ? AND status IN ? AND version = ?", id, values, version)
	if requireRunnablePlan {
		query = query.Where(
			"plan_id IN (?)",
			r.dbp.WriteDB(ctx).Model(&premmod.UpgradePlan{}).Select("id").Where("status = ?", string(prem.PlanStatusRunning)),
		)
	}
	result := query.
		Updates(map[string]any{
			"status":        string(prem.StepStatusPreparing),
			"version":       gorm.Expr("version + 1"),
			"attempt_count": gorm.Expr("attempt_count + 1"),
		})
	return result.RowsAffected == 1, result.Error
}

func (r *repo) PatchStep(ctx context.Context, id uint64, values map[string]any) error {
	return r.dbp.WriteDB(ctx).Model(&premmod.UpgradeStep{}).Where("id = ?", id).Updates(values).Error
}

func (r *repo) ResetFailedStepForRetry(ctx context.Context, reset RetryReset) error {
	return r.dbp.Transaction(ctx, func(txCtx context.Context) error {
		db := r.dbp.WriteDB(txCtx)
		result := db.Model(&premmod.UpgradeStep{}).
			Where("id = ? AND plan_id = ? AND target_id = ? AND status = ? AND configuration_change_id IS NULL",
				reset.StepID, reset.PlanID, reset.TargetID, string(prem.StepStatusFailed)).
			Updates(map[string]any{
				"status": string(prem.StepStatusPending), "status_reason": "manual retry requested",
				"next_retry_at": nil, "apply_deadline_at": nil, "finished_at": nil,
				"last_error_code": "", "last_error_message": "",
			})
		if result.Error != nil {
			return result.Error
		}
		if result.RowsAffected != 1 {
			return fmt.Errorf("failed step %d is no longer retryable", reset.StepID)
		}
		resetSkipped := func(ids []uint64) error {
			if len(ids) == 0 {
				return nil
			}
			return db.Model(&premmod.UpgradeStep{}).
				Where("id IN ? AND plan_id = ? AND status = ?", ids, reset.PlanID, string(prem.StepStatusSkipped)).
				Updates(map[string]any{
					"status": string(prem.StepStatusPending), "status_reason": "",
					"finished_at": nil, "next_retry_at": nil,
				}).Error
		}
		if err := resetSkipped(reset.LaterStepIDs); err != nil {
			return err
		}
		if err := resetSkipped(reset.DependentStepIDs); err != nil {
			return err
		}
		if err := db.Model(&premmod.UpgradeTarget{}).Where("id = ? AND plan_id = ?", reset.TargetID, reset.PlanID).
			Updates(map[string]any{
				"status": string(prem.TargetStatusPending), "last_error": "", "finished_at": nil,
			}).Error; err != nil {
			return err
		}
		if len(reset.DependentTargetIDs) > 0 {
			if err := db.Model(&premmod.UpgradeTarget{}).
				Where("id IN ? AND plan_id = ?", reset.DependentTargetIDs, reset.PlanID).
				Updates(map[string]any{
					"status": string(prem.TargetStatusPending), "last_error": "", "finished_at": nil,
				}).Error; err != nil {
				return err
			}
		}
		result = db.Model(&premmod.UpgradePlan{}).
			Where("id = ? AND status IN ?", reset.PlanID, []string{
				string(prem.PlanStatusRunning), string(prem.PlanStatusFailed), string(prem.PlanStatusPartialFailed),
			}).
			Updates(map[string]any{
				"status": string(prem.PlanStatusRunning), "last_error": "", "finished_at": nil,
			})
		if result.Error != nil {
			return result.Error
		}
		if result.RowsAffected != 1 {
			return fmt.Errorf("plan %d is no longer retryable", reset.PlanID)
		}
		return nil
	})
}

func planToModel(entity *prem.UpgradePlan) *premmod.UpgradePlan {
	model := &premmod.UpgradePlan{
		Updatable: base.Updatable{Base: base.Base{ID: base.ID(entity.ID)}},
		Name:      entity.Name, Reason: entity.Reason, Applicant: entity.Applicant,
		OperatorEntity: entity.OperatorEntity, DesiredVersion: entity.DesiredVersion,
		ExecutionMode: string(entity.ExecutionMode),
		ScopeMode:     string(entity.ScopeMode), RequestedClusterIDs: entity.RequestedClusterIDs,
		RequestedComponents: entity.RequestedComponents, ApprovalManifest: entity.ApprovalManifest,
		ManifestHash: entity.ManifestHash, ApprovalID: entity.ApprovalID,
		Status: string(entity.Status), LastError: entity.LastError,
		MaxParallelTargets:     entity.MaxParallelTargets,
		ComponentOrderRevision: entity.ComponentOrderRevision,
		IdempotencyKey:         entity.IdempotencyKey, CancelRequestedAt: entity.CancelRequestedAt,
		StartedAt: entity.StartedAt, FinishedAt: entity.FinishedAt, Version: entity.Version,
	}
	if entity.ApprovalStatus != nil {
		status := string(*entity.ApprovalStatus)
		model.ApprovalStatus = &status
	}
	for _, target := range entity.Targets {
		targetModel := &premmod.UpgradeTarget{
			Updatable: base.Updatable{Base: base.Base{ID: base.ID(target.ID)}},
			PlanID:    target.PlanID, ClusterType: string(target.ClusterType), ClusterID: target.ClusterID,
			SourceClusterIDs: target.SourceClusterIDs, ClusterName: target.ClusterName,
			TenantIDs: target.TenantIDs, SharedPoolID: target.SharedPoolID,
			LogicalRegion: target.LogicalRegion, CloudProvider: target.CloudProvider,
			Status: string(target.Status), LastError: target.LastError,
			StartedAt: target.StartedAt, FinishedAt: target.FinishedAt,
		}
		for _, step := range target.Steps {
			stepModel := &premmod.UpgradeStep{
				Updatable: base.Updatable{Base: base.Base{ID: base.ID(step.ID)}},
				PlanID:    step.PlanID, TargetID: step.TargetID, Sequence: step.Sequence,
				ApplyGroup:    step.ApplyGroup,
				ApplyStrategy: step.ApplyStrategy, ApplyMode: step.ApplyMode,
				ConfigurationChangeID: step.ConfigurationChangeID, Status: string(step.Status),
				StatusReason: step.StatusReason, AttemptCount: step.AttemptCount,
				NextRetryAt: step.NextRetryAt, ApplyDeadlineAt: step.ApplyDeadlineAt,
				LastErrorCode: step.LastErrorCode, LastErrorMessage: step.LastErrorMessage,
				Version: step.Version, StartedAt: step.StartedAt, FinishedAt: step.FinishedAt,
			}
			for _, component := range step.Components {
				stepModel.Components = append(stepModel.Components, premmod.UpgradeStepComponent{
					Component: component.Component, FromVersion: component.FromVersion, ToVersion: component.ToVersion,
					FromImageTag: component.FromImageTag, ToImageTag: component.ToImageTag,
				})
			}
			targetModel.Steps = append(targetModel.Steps, stepModel)
		}
		model.Targets = append(model.Targets, targetModel)
	}
	return model
}

func planToEntity(model *premmod.UpgradePlan) *prem.UpgradePlan {
	entity := &prem.UpgradePlan{
		ID: model.ID.Uint64(), Name: model.Name, Reason: model.Reason, Applicant: model.Applicant,
		OperatorEntity: model.OperatorEntity, DesiredVersion: model.DesiredVersion,
		ExecutionMode: prem.ExecutionMode(model.ExecutionMode),
		ScopeMode:     prem.ScopeMode(model.ScopeMode), RequestedClusterIDs: model.RequestedClusterIDs,
		RequestedComponents: model.RequestedComponents, ApprovalManifest: model.ApprovalManifest,
		ManifestHash: model.ManifestHash, ApprovalID: model.ApprovalID,
		Status: prem.PlanStatus(model.Status), LastError: model.LastError,
		MaxParallelTargets: model.MaxParallelTargets, ComponentOrderRevision: model.ComponentOrderRevision,
		IdempotencyKey: model.IdempotencyKey, CancelRequestedAt: model.CancelRequestedAt,
		StartedAt: model.StartedAt, FinishedAt: model.FinishedAt, Version: model.Version,
		CreatedAt: model.CreatedAt, UpdatedAt: model.UpdatedAt,
	}
	if model.ApprovalStatus != nil {
		status := apprent.ApprovalStatus(*model.ApprovalStatus)
		entity.ApprovalStatus = &status
	}
	for _, targetModel := range model.Targets {
		target := &prem.UpgradeTarget{
			ID: targetModel.ID.Uint64(), PlanID: targetModel.PlanID,
			ClusterType: ccent.ClusterType(targetModel.ClusterType), ClusterID: targetModel.ClusterID,
			SourceClusterIDs: targetModel.SourceClusterIDs, ClusterName: targetModel.ClusterName,
			TenantIDs: targetModel.TenantIDs, SharedPoolID: targetModel.SharedPoolID,
			LogicalRegion: targetModel.LogicalRegion, CloudProvider: targetModel.CloudProvider,
			Status: prem.TargetStatus(targetModel.Status), LastError: targetModel.LastError,
			StartedAt: targetModel.StartedAt, FinishedAt: targetModel.FinishedAt,
			CreatedAt: targetModel.CreatedAt, UpdatedAt: targetModel.UpdatedAt,
		}
		for _, stepModel := range targetModel.Steps {
			step := &prem.UpgradeStep{
				ID: stepModel.ID.Uint64(), PlanID: stepModel.PlanID, TargetID: stepModel.TargetID,
				Sequence: stepModel.Sequence, ApplyGroup: stepModel.ApplyGroup,
				ApplyStrategy: stepModel.ApplyStrategy, ApplyMode: stepModel.ApplyMode,
				ConfigurationChangeID: stepModel.ConfigurationChangeID, Status: prem.StepStatus(stepModel.Status),
				StatusReason: stepModel.StatusReason, AttemptCount: stepModel.AttemptCount,
				NextRetryAt: stepModel.NextRetryAt, ApplyDeadlineAt: stepModel.ApplyDeadlineAt,
				LastErrorCode: stepModel.LastErrorCode, LastErrorMessage: stepModel.LastErrorMessage,
				Version: stepModel.Version, StartedAt: stepModel.StartedAt, FinishedAt: stepModel.FinishedAt,
				CreatedAt: stepModel.CreatedAt, UpdatedAt: stepModel.UpdatedAt,
			}
			for _, component := range stepModel.Components {
				step.Components = append(step.Components, prem.UpgradeStepComponent{
					Component: component.Component, FromVersion: component.FromVersion, ToVersion: component.ToVersion,
					FromImageTag: component.FromImageTag, ToImageTag: component.ToImageTag,
				})
			}
			target.Steps = append(target.Steps, step)
		}
		entity.Targets = append(entity.Targets, target)
	}
	return entity
}