package migration

import (
	"time"

	"github.com/samber/lo"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"

	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/infra/repoimpl/schema"
)

// Batch add 1:8 TiDB SKUs
var _ = AddDBMigrateRecord(
	"202506171200",
	func(db *gorm.DB) error {
		timestamp := time.Now()

		return db.Transaction(func(tx *gorm.DB) error {
			// Fetch the AWS provider.
			var provider *schema.Provider
			if err := tx.Where("name = ?", "aws").First(&provider).Error; err != nil {
				return err
			}

			// Add instance types.
			instanceTypes := []*schema.InstanceType{
				{
					InstanceType: "r6g.2xlarge",
					ProviderId:   uint64(provider.ID),
					CPUNum:       8,
					MemorySizeGi: 64,
					CpuArch:      "arm64",
				},
				{
					InstanceType: "r6g.4xlarge",
					ProviderId:   uint64(provider.ID),
					CPUNum:       16,
					MemorySizeGi: 128,
					CpuArch:      "arm64",
				},
			}
			for _, instanceType := range instanceTypes {
				instanceType.CreatedAt = timestamp
				instanceType.UpdatedAt = timestamp
				if err := tx.Clauses(clause.OnConflict{
					Columns:   []clause.Column{{Name: "instance_type"}},
					UpdateAll: true,
				}).Create(instanceType).Error; err != nil { // The ID will be populated in the instanceType object after successful insertion
					return err
				}
			}
			instanceTypeMap := lo.SliceToMap(
				instanceTypes,
				func(i *schema.InstanceType) (string, *schema.InstanceType) {
					return i.InstanceType, i
				},
			)

			// Add tidb skus.
			{
				tidbSkus := make([]*schema.ComponentTierUnit, 0, 4)

				tidbSpu8c64g := &schema.ComponentTier{
					Key:           "8C64G:DEPLOY_TYPE_DEDICATED:COMPONENT_TYPE_TIDB",
					ComponentType: "tidb",
					DisplayName:   "8 vCPU, 64 GiB",
					CPUNum:        8,
					MemorySizeGi:  64,
				}
				if err := tx.Clauses(clause.OnConflict{
					Columns:   []clause.Column{{Name: "key"}},
					UpdateAll: true,
				}).Create(tidbSpu8c64g).Error; err != nil {
					return err
				}
				tidbSkus = append(tidbSkus, &schema.ComponentTierUnit{
					Key:             "aws:tidb:r6g.2xlarge",
					Priority:        999,
					ComponentType:   "tidb",
					ProviderId:      uint64(provider.ID),
					ComponentTierId: uint64(tidbSpu8c64g.ID),
					InstanceTypeId:  uint64(instanceTypeMap["r6g.2xlarge"].ID),
					DefaultNodeNum:  2,
				})

				tidbSpu16c128g := &schema.ComponentTier{
					Key:           "16C128G:DEPLOY_TYPE_DEDICATED:COMPONENT_TYPE_TIDB",
					ComponentType: "tidb",
					DisplayName:   "16 vCPU, 128 GiB",
					CPUNum:        16,
					MemorySizeGi:  128,
				}
				if err := tx.Clauses(clause.OnConflict{
					Columns:   []clause.Column{{Name: "key"}},
					UpdateAll: true,
				}).Create(tidbSpu16c128g).Error; err != nil {
					return err
				}
				tidbSkus = append(tidbSkus, &schema.ComponentTierUnit{
					Key:             "aws:tidb:r6g.4xlarge",
					Priority:        999,
					ComponentType:   "tidb",
					ProviderId:      uint64(provider.ID),
					ComponentTierId: uint64(tidbSpu16c128g.ID),
					InstanceTypeId:  uint64(instanceTypeMap["r6g.4xlarge"].ID),
					DefaultNodeNum:  2,
				})

				for _, tidbSku := range tidbSkus {
					tidbSku.CreatedAt = timestamp
					tidbSku.UpdatedAt = timestamp
					if err := tx.Clauses(clause.OnConflict{
						Columns:   []clause.Column{{Name: "key"}},
						UpdateAll: true,
					}).Create(tidbSku).Error; err != nil {
						return err
					}
				}
			}

			return nil
		})
	},
)