package migrate

import (
	"fmt"

	strutil "github.com/tidbcloud/aisupport/packages/sdk-go/infra/utils/container/str"

	"gorm.io/gorm"
)

var _ = addRecord("202607201100", migrateDispatchPolicyFamilyID)

type dispatchPolicyWithNullableFamilyID struct {
	ID               uint64  `gorm:"column:id;primaryKey"`
	FamilyID         *string `gorm:"column:family_id;type:varchar(20);index:idx_dispatch_policy_family_id"`
	ReplacesPolicyID *uint64 `gorm:"column:replaces_policy_id"`
}

func (dispatchPolicyWithNullableFamilyID) TableName() string {
	return "dispatch_policies"
}

type dispatchPolicyWithFamilyID struct {
	FamilyID string `gorm:"column:family_id;type:varchar(20);not null;index:idx_dispatch_policy_family_id"`
}

func (dispatchPolicyWithFamilyID) TableName() string {
	return "dispatch_policies"
}

func migrateDispatchPolicyFamilyID(db *gorm.DB) error {
	legacy := &dispatchPolicyWithNullableFamilyID{}
	if err := db.AutoMigrate(legacy); err != nil {
		return err
	}

	var policies []*dispatchPolicyWithNullableFamilyID
	if err := db.Order("id ASC").Find(&policies).Error; err != nil {
		return err
	}
	byID := make(map[uint64]*dispatchPolicyWithNullableFamilyID, len(policies))
	for _, policy := range policies {
		byID[policy.ID] = policy
	}

	familyIDs := make(map[uint64]string, len(policies))
	resolving := make(map[uint64]bool, len(policies))
	usedFamilyIDs := make(map[string]struct{}, len(policies))
	var resolveFamilyID func(uint64) (string, error)
	resolveFamilyID = func(policyID uint64) (string, error) {
		if familyID, ok := familyIDs[policyID]; ok {
			return familyID, nil
		}
		if resolving[policyID] {
			return "", fmt.Errorf("dispatch policy replacement chain contains a cycle at policy %d", policyID)
		}
		policy, ok := byID[policyID]
		if !ok {
			return "", fmt.Errorf("dispatch policy replacement chain references missing policy %d", policyID)
		}

		resolving[policyID] = true
		defer delete(resolving, policyID)

		var familyID string
		if policy.ReplacesPolicyID != nil {
			var err error
			familyID, err = resolveFamilyID(*policy.ReplacesPolicyID)
			if err != nil {
				return "", err
			}
		} else if policy.FamilyID != nil && *policy.FamilyID != "" {
			familyID = *policy.FamilyID
		} else {
			for {
				familyID = "pol_" + strutil.RandomBase62String(16)
				if _, exists := usedFamilyIDs[familyID]; !exists {
					break
				}
			}
		}
		familyIDs[policyID] = familyID
		usedFamilyIDs[familyID] = struct{}{}
		return familyID, nil
	}

	for _, policy := range policies {
		familyID, err := resolveFamilyID(policy.ID)
		if err != nil {
			return err
		}
		if err := db.Exec("UPDATE dispatch_policies SET family_id = ? WHERE id = ?", familyID, policy.ID).Error; err != nil {
			return err
		}
	}

	if db.Dialector.Name() == "sqlite" {
		return nil
	}
	return db.Migrator().AlterColumn(&dispatchPolicyWithFamilyID{}, "FamilyID")
}
