package migration

import (
	"encoding/json"
	"fmt"
	"time"

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

	globaldedicatedv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/dedicated/v1"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/entity"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/infra/repoimpl/schema"
)

// Add following V2 SPUs and SKUs:
//   - TiDB: 8C64G, 16C128G
//   - TiKV: 64C256G
var _ = AddDBMigrateRecord(
	"release_20251202",
	func(db *gorm.DB) error {
		timestamp := time.Now()

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

		// Add instance types.
		instanceTypes := []*schema.InstanceType{
			{
				InstanceType: "r8g.2xlarge",
				ProviderId:   uint64(provider.ID),
				CPUNum:       8,
				MemorySizeGi: 64,
				CpuArch:      "arm64",
			},
			{
				InstanceType: "r8g.4xlarge",
				ProviderId:   uint64(provider.ID),
				CPUNum:       16,
				MemorySizeGi: 128,
				CpuArch:      "arm64",
			},
			{
				InstanceType: "m8g.16xlarge",
				ProviderId:   uint64(provider.ID),
				CPUNum:       64,
				MemorySizeGi: 256,
				CpuArch:      "arm64",
			},
		}
		for _, instanceType := range instanceTypes {
			instanceType.CreatedAt = timestamp
			instanceType.UpdatedAt = timestamp
			if err := db.Clauses(clause.OnConflict{
				Columns:   []clause.Column{{Name: "instance_type"}},
				UpdateAll: true,
			}).Create(instanceType).Error; err != nil {
				return fmt.Errorf("failed to create instance type %s: %w", instanceType.InstanceType, err)
			}
		}

		// Re-query instance types to get their ids.
		var createdInstanceTypes []*schema.InstanceType
		if err := db.Where("instance_type IN ?", lo.Map(instanceTypes, func(i *schema.InstanceType, _ int) string {
			return i.InstanceType
		})).Find(&createdInstanceTypes).Error; err != nil {
			return err
		}
		instanceTypeMap := lo.SliceToMap(
			createdInstanceTypes,
			func(i *schema.InstanceType) (string, *schema.InstanceType) { return i.InstanceType, i },
		)

		// Add volume types.
		volumeTypes := []*schema.StorageTemplate{
			{
				ProviderId: uint64(provider.ID),
				VolumeType: "io2",
			},
			{
				ProviderId: uint64(provider.ID),
				VolumeType: "gp3",
			},
		}
		for _, volumeType := range volumeTypes {
			volumeType.CreatedAt = timestamp
			volumeType.UpdatedAt = timestamp
			if err := db.Clauses(clause.OnConflict{
				Columns:   []clause.Column{{Name: "volume_type"}},
				UpdateAll: true,
			}).Create(volumeType).Error; err != nil {
				return fmt.Errorf("failed to create volume type %s: %w", volumeType.VolumeType, err)
			}
		}

		// Re-query volume types to get their ids.
		var createdVolumeTypes []*schema.StorageTemplate
		if err := db.Where("volume_type IN ?", lo.Map(volumeTypes, func(v *schema.StorageTemplate, _ int) string {
			return v.VolumeType
		})).Find(&createdVolumeTypes).Error; err != nil {
			return err
		}
		volumeTypeMap := lo.SliceToMap(
			createdVolumeTypes,
			func(v *schema.StorageTemplate) (string, *schema.StorageTemplate) { return v.VolumeType, v },
		)

		// Create new V2 spus.
		v2Spus := []*schema.ComponentTier{
			{ComponentType: "tidb", Name: "8C64G-V2", Version: newSpuVersion, CPUNum: 8, MemorySizeGi: 64},
			{ComponentType: "tidb", Name: "16C128G-V2", Version: newSpuVersion, CPUNum: 16, MemorySizeGi: 128},
			{ComponentType: "tikv", Name: "64C256G-V2", Version: newSpuVersion, CPUNum: 64, MemorySizeGi: 256},
		}
		v2Spus = lo.Map(v2Spus, func(spu *schema.ComponentTier, _ int) *schema.ComponentTier {
			var ct globaldedicatedv1.ComponentType
			switch spu.ComponentType {
			case "tidb":
				ct = globaldedicatedv1.ComponentType_TIDB
			case "tikv":
				ct = globaldedicatedv1.ComponentType_TIKV
			case "tiflash":
				ct = globaldedicatedv1.ComponentType_TIFLASH
			case "pd":
				ct = globaldedicatedv1.ComponentType_PD
			case "tiproxy":
				ct = globaldedicatedv1.ComponentType_TIPROXY
			default:
				// Unknown component type, return as-is.
				return spu
			}
			spu.Key = entity.BuildSpuKey(ct, spu.Name, spu.Version)
			spu.DisplayName = fmt.Sprintf("%d vCPU, %d GiB (%s)", spu.CPUNum, spu.MemorySizeGi, newSpuVersion)
			return spu
		})
		if err := db.Clauses(clause.OnConflict{
			UpdateAll: true,
		}).CreateInBatches(v2Spus, 20).Error; err != nil {
			return err
		}

		// Re-query v2 spus to get their ids.
		var createdV2Spus []*schema.ComponentTier
		if err := db.Where("version = ?", newSpuVersion).Find(&createdV2Spus).Error; err != nil {
			return err
		}
		v2SpuMap := lo.KeyBy(createdV2Spus, func(spu *schema.ComponentTier) string {
			return spu.Name
		})

		// Create corresponding SKUs.
		skuTidb8C64G := &schema.ComponentTierUnit{
			Key:             "aws:tidb:r8g.2xlarge:v2",
			ComponentType:   v2SpuMap["8C64G-V2"].ComponentType,
			ProviderId:      uint64(provider.ID),
			ComponentTierId: uint64(v2SpuMap["8C64G-V2"].ID),
			InstanceTypeId:  uint64(instanceTypeMap["r8g.2xlarge"].ID),
			DefaultNodeNum:  2,
		}
		skuTidb16C128G := &schema.ComponentTierUnit{
			Key:             "aws:tidb:r8g.4xlarge:v2",
			ComponentType:   v2SpuMap["16C128G-V2"].ComponentType,
			ProviderId:      uint64(provider.ID),
			ComponentTierId: uint64(v2SpuMap["16C128G-V2"].ID),
			InstanceTypeId:  uint64(instanceTypeMap["r8g.4xlarge"].ID),
			DefaultNodeNum:  2,
		}

		skuTikv64C256GBasic := &schema.ComponentTierUnit{
			Key:               "aws:tikv:m8g.16xlarge:Basic:v2",
			ComponentType:     v2SpuMap["64C256G-V2"].ComponentType,
			ProviderId:        uint64(provider.ID),
			ComponentTierId:   uint64(v2SpuMap["64C256G-V2"].ID),
			InstanceTypeId:    uint64(instanceTypeMap["m8g.16xlarge"].ID),
			DefaultNodeNum:    3,
			StorageTemplateId: lo.ToPtr(uint64(volumeTypeMap["gp3"].ID)),
			DefaultStorageGi:  500,
			MinStorageGi:      200,
			MaxStorageGi:      4096,
			DefaultIops:       16000,
			DefaultThroughPut: lo.ToPtr(uint32(1000)),
		}
		skuTikv64C256GPlus := &schema.ComponentTierUnit{
			Key:               "aws:tikv:m8g.16xlarge:Plus:v2",
			ComponentType:     v2SpuMap["64C256G-V2"].ComponentType,
			ProviderId:        uint64(provider.ID),
			ComponentTierId:   uint64(v2SpuMap["64C256G-V2"].ID),
			InstanceTypeId:    uint64(instanceTypeMap["m8g.16xlarge"].ID),
			DefaultNodeNum:    3,
			StorageTemplateId: lo.ToPtr(uint64(volumeTypeMap["io2"].ID)),
			DefaultStorageGi:  500,
			MinStorageGi:      200,
			MaxStorageGi:      4096,
			DefaultIops:       8000,
		}
		skuTikv64C256GStandard := &schema.ComponentTierUnit{
			Key:                    "aws:tikv:m8g.16xlarge:Standard:v2",
			ComponentType:          v2SpuMap["64C256G-V2"].ComponentType,
			ProviderId:             uint64(provider.ID),
			ComponentTierId:        uint64(v2SpuMap["64C256G-V2"].ID),
			InstanceTypeId:         uint64(instanceTypeMap["m8g.16xlarge"].ID),
			DefaultNodeNum:         3,
			StorageTemplateId:      lo.ToPtr(uint64(volumeTypeMap["gp3"].ID)),
			DefaultStorageGi:       500,
			MinStorageGi:           200,
			MaxStorageGi:           4096,
			MaxIops:                16000,
			MinIops:                3000,
			IopsPerGi:              4,
			MaxThroughPut:          1000,
			MinThroughPut:          125,
			ThroughputMibpsPerGi:   0.5,
			ThroughputMibpsPerIops: 0.046875,
			RaftStorageTemplateId:  lo.ToPtr(uint64(volumeTypeMap["gp3"].ID)),
			RaftDefaultStorageGi:   40,
			RaftDefaultIops:        3000,
			RaftDefaultThroughPut:  125,
		}
		skuTikv64C256GPerformance := &schema.ComponentTierUnit{
			Key:                    "aws:tikv:m8g.16xlarge:Performance:v2",
			ComponentType:          v2SpuMap["64C256G-V2"].ComponentType,
			ProviderId:             uint64(provider.ID),
			ComponentTierId:        uint64(v2SpuMap["64C256G-V2"].ID),
			InstanceTypeId:         uint64(instanceTypeMap["m8g.16xlarge"].ID),
			DefaultNodeNum:         3,
			StorageTemplateId:      lo.ToPtr(uint64(volumeTypeMap["gp3"].ID)),
			DefaultStorageGi:       500,
			MinStorageGi:           200,
			MaxStorageGi:           4096,
			MaxIops:                16000,
			MinIops:                3000,
			IopsPerGi:              4,
			MaxThroughPut:          1000,
			MinThroughPut:          125,
			ThroughputMibpsPerGi:   0.5,
			ThroughputMibpsPerIops: 0.046875,
			RaftStorageTemplateId:  lo.ToPtr(uint64(volumeTypeMap["io2"].ID)),
			RaftDefaultStorageGi:   40,
			RaftDefaultIops:        4000,
		}
		skus := []*schema.ComponentTierUnit{
			skuTidb8C64G,
			skuTidb16C128G,
			skuTikv64C256GBasic,
			skuTikv64C256GPlus,
			skuTikv64C256GStandard,
			skuTikv64C256GPerformance,
		}
		for _, sku := range skus {
			sku.CreatedAt = timestamp
			sku.UpdatedAt = timestamp
			if err := db.Clauses(clause.OnConflict{
				Columns:   []clause.Column{{Name: "key"}},
				UpdateAll: true,
			}).Create(sku).Error; err != nil {
				return fmt.Errorf("failed to create sku %s: %w", sku.Key, err)
			}
		}

		// Create SPU-SKU mappings.
		mappingTidb8C64G := &schema.TierTierUnit{
			ComponentTierId: uint64(v2SpuMap["8C64G-V2"].ID),
			TierUnitMapping: schema.TierUnitMapping{
				Items: []*schema.TierUnitMappingItem{
					{
						ProviderName: provider.Name,
						RegionNames:  []string{"NAME_UNSPECIFIED"},
						TenantIDs:    []string{}, // default to no tenant. modify it manually.
						SkuKey:       skuTidb8C64G.Key,
					},
				},
			},
		}
		mappingTidb16C128G := &schema.TierTierUnit{
			ComponentTierId: uint64(v2SpuMap["16C128G-V2"].ID),
			TierUnitMapping: schema.TierUnitMapping{
				Items: []*schema.TierUnitMappingItem{
					{
						ProviderName: provider.Name,
						RegionNames:  []string{"NAME_UNSPECIFIED"},
						TenantIDs:    []string{}, // default to no tenant. modify it manually.
						SkuKey:       skuTidb16C128G.Key,
					},
				},
			},
		}
		mappingTikv64C256G := &schema.TierTierUnit{
			ComponentTierId: uint64(v2SpuMap["64C256G-V2"].ID),
			TierUnitMapping: schema.TierUnitMapping{
				Items: []*schema.TierUnitMappingItem{
					{
						ProviderName: provider.Name,
						RegionNames:  []string{"NAME_UNSPECIFIED"},
						TenantIDs:    []string{}, // default to no tenant. modify it manually.
						SkuKey:       skuTikv64C256GBasic.Key,
					},
					{
						ProviderName: provider.Name,
						RegionNames:  []string{"NAME_UNSPECIFIED"},
						TenantIDs:    []string{}, // default to no tenant. modify it manually.
						SkuKey:       skuTikv64C256GPlus.Key,
					},
					{
						ProviderName: provider.Name,
						RegionNames:  []string{"NAME_UNSPECIFIED"},
						TenantIDs:    []string{}, // default to no tenant. modify it manually.
						SkuKey:       skuTikv64C256GStandard.Key,
					},
					{
						ProviderName: provider.Name,
						RegionNames:  []string{"NAME_UNSPECIFIED"},
						TenantIDs:    []string{}, // default to no tenant. modify it manually.
						SkuKey:       skuTikv64C256GPerformance.Key,
					},
				},
			},
		}
		mappings := []*schema.TierTierUnit{
			mappingTidb8C64G,
			mappingTidb16C128G,
			mappingTikv64C256G,
		}
		for _, mapping := range mappings {
			if err := db.Clauses(clause.OnConflict{
				Columns:   []clause.Column{{Name: "component_tier_id"}},
				UpdateAll: true,
			}).Create(mapping).Error; err != nil {
				return fmt.Errorf("failed to create spu-sku mapping for spu id %d: %w", mapping.ComponentTierId, err)
			}
		}

		// Bind TiKV SKUs to the PD SKU.
		// Query "aws:pd:c8g.8xlarge:gp3:3000:125:v2" PD.
		var pdSku *schema.ComponentTierUnit
		if err := db.Where("`key` = ?", "aws:pd:c8g.8xlarge:gp3:3000:125:v2").First(&pdSku).Error; err != nil {
			return fmt.Errorf("failed to query pd sku aws:pd:c8g.8xlarge:gp3:3000:125:v2: %w", err)
		}
		// Query bindings.
		var pdSkuBinding *schema.TierUnitCondition
		if err := db.Where("component_tier_unit_id = ?", pdSku.ID).First(&pdSkuBinding).Error; err != nil {
			return fmt.Errorf("failed to query pd sku binding for sku id %d: %w", pdSku.ID, err)
		}
		// Update bindings to include the new TiKV SKUs.
		var existingValues []string
		if err := json.Unmarshal([]byte(pdSkuBinding.Values), &existingValues); err != nil {
			return fmt.Errorf("failed to unmarshal existing pd sku binding values: %w", err)
		}
		newValues := lo.Uniq(
			append(existingValues, skuTikv64C256GBasic.Key, skuTikv64C256GPlus.Key, skuTikv64C256GStandard.Key, skuTikv64C256GPerformance.Key),
		)
		newValuesBytes, err := json.Marshal(newValues)
		if err != nil {
			return fmt.Errorf("failed to marshal new pd sku binding values: %w", err)
		}
		pdSkuBinding.Values = string(newValuesBytes)
		if err := db.Save(pdSkuBinding).Error; err != nil {
			return fmt.Errorf("failed to save pd sku binding for sku id %d: %w", pdSku.ID, err)
		}

		// Migration completed.

		return nil
	},
)