package migration

import (
	"gorm.io/gorm"

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

// Set iops/throughput mapping parameters for **data** storages, for tikv:
// - Standard and Performance (volume type of data disk is gp3).
//   - ops= 4/ GB, 3000-16000
//   - throughput=max(0.046875 MB per IOPS, 0.5MB per GB Size), 125-1000
//
// - Basic, only set for Azure.
var _ = AddDBMigrateRecord(
	"202503241200",
	func(db *gorm.DB) error {
		return db.Transaction(setIopsThroughputMapping)
	},
)

func setIopsThroughputMapping(tx *gorm.DB) error {
	// STEP 1: Fetch original TiKV components
	var originalUnits []schema.ComponentTierUnit

	subQuery := tx.Model(&schema.ComponentTierUnit{}).
		Select("component_tier_units.id").
		Joins("JOIN instance_types ON component_tier_units.instance_type_id = instance_types.id AND instance_types.cpu_arch = 'arm64'")

	if err := tx.Model(&schema.ComponentTierUnit{}).
		Where("id IN (?)", subQuery).
		Where("component_type LIKE ?", "%tikv%").
		Where(
			tx.Where("`key` LIKE ?", "%:Standard").
				Or("`key` LIKE ?", "%:Performance"),
		).
		Find(&originalUnits).Error; err != nil {
		return err
	}

	var originalUnitsAzure []schema.ComponentTierUnit
	if err := tx.Model(&schema.ComponentTierUnit{}).Where(
		tx.Where("component_type LIKE ?", "%tikv%").
			Or("component_type LIKE ?", "%tiflash%"),
	).
		Where("`key` LIKE ?", "%premiumv2%").
		Find(&originalUnitsAzure).Error; err != nil {
		return err
	}

	originalUnits = append(originalUnits, originalUnitsAzure...)

	// SETP 2: Update iops/throughput mapping parameters.
	for i := range originalUnits {
		if err := tx.Model(&schema.ComponentTierUnit{}).
			Where("id = ?", originalUnits[i].ID).
			Updates(map[string]any{
				"default_iops":              0,
				"default_through_put":       0,
				"max_iops":                  16000,
				"min_iops":                  3000,
				"iops_per_gi":               4,
				"max_through_put":           1000,
				"min_through_put":           125,
				"throughput_mibps_per_gi":   0.5,
				"throughput_mibps_per_iops": 0.046875,
			}).Error; err != nil {
			return err
		}
	}

	return nil
}