package migration

import (
	"errors"
	"fmt"
	"log"

	"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"
)

const (
	existingSpuVersion = "V1"
	newSpuVersion      = "V2"
)

var _ = AddDBMigrateRecord(
	"202509021200",
	func(db *gorm.DB) error {
		if err := db.AutoMigrate(&schema.ComponentTier{}); err != nil {
			return err
		}

		// Query all ComponentTier.
		var spus []*schema.ComponentTier
		if err := db.Find(&spus).Error; err != nil {
			return err
		}

		// Update Name and ComponentType for each spu.
		// Set version of all existing spus to "v1".
		for _, spu := range spus {
			name, comp, err := entity.ParseSpuKey(spu.Key)
			if err != nil {
				return err
			}
			spu.Name = name
			spu.ComponentType = comp.String()
			spu.Version = existingSpuVersion
		}
		if len(spus) > 0 {
			if err := db.Clauses(clause.OnConflict{
				UpdateAll: true,
			}).CreateInBatches(spus, 20).Error; err != nil {
				return err
			}
		}

		// Create new V2 spus.
		v2Spus := []*schema.ComponentTier{
			{ComponentType: "tidb", Name: "4C16G-V2", Version: newSpuVersion, CPUNum: 4, MemorySizeGi: 16},
			{ComponentType: "tidb", Name: "8C16G-V2", Version: newSpuVersion, CPUNum: 8, MemorySizeGi: 16},
			{ComponentType: "tidb", Name: "8C32G-V2", Version: newSpuVersion, CPUNum: 8, MemorySizeGi: 32},
			{ComponentType: "tidb", Name: "16C32G-V2", Version: newSpuVersion, CPUNum: 16, MemorySizeGi: 32},
			{ComponentType: "tidb", Name: "16C64G-V2", Version: newSpuVersion, CPUNum: 16, MemorySizeGi: 64},
			{ComponentType: "tidb", Name: "32C64G-V2", Version: newSpuVersion, CPUNum: 32, MemorySizeGi: 64},
			{ComponentType: "tidb", Name: "32C128G-V2", Version: newSpuVersion, CPUNum: 32, MemorySizeGi: 128},
			{ComponentType: "tikv", Name: "4C16G-V2", Version: newSpuVersion, CPUNum: 4, MemorySizeGi: 16},
			{ComponentType: "tikv", Name: "8C32G-V2", Version: newSpuVersion, CPUNum: 8, MemorySizeGi: 32},
			{ComponentType: "tikv", Name: "8C64G-V2", Version: newSpuVersion, CPUNum: 8, MemorySizeGi: 64},
			{ComponentType: "tikv", Name: "16C64G-V2", Version: newSpuVersion, CPUNum: 16, MemorySizeGi: 64},
			{ComponentType: "tikv", Name: "32C128G-V2", Version: newSpuVersion, CPUNum: 32, MemorySizeGi: 128},
			{ComponentType: "tiflash", Name: "8C64G-V2", Version: newSpuVersion, CPUNum: 8, MemorySizeGi: 64},
			{ComponentType: "tiflash", Name: "16C128G-V2", Version: newSpuVersion, CPUNum: 16, MemorySizeGi: 128},
			{ComponentType: "tiflash", Name: "32C128G-V2", Version: newSpuVersion, CPUNum: 32, MemorySizeGi: 128},
			{ComponentType: "tiflash", Name: "32C256G-V2", Version: newSpuVersion, CPUNum: 32, MemorySizeGi: 256},
			{ComponentType: "pd", Name: "2C8G-V2", Version: newSpuVersion, CPUNum: 2, MemorySizeGi: 8},
			{ComponentType: "pd", Name: "4C8G-V2", Version: newSpuVersion, CPUNum: 4, MemorySizeGi: 8},
			{ComponentType: "pd", Name: "16C32G-V2", Version: newSpuVersion, CPUNum: 16, MemorySizeGi: 32},
			{ComponentType: "pd", Name: "32C64G-V2", Version: newSpuVersion, CPUNum: 32, MemorySizeGi: 64},
		}
		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 len(v2Spus) > 0 {
			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.Key
		})

		// migrate aws g4 skus to new v2 spus.
		for _, v2Spu := range v2Spus {
			createdV2Spu, ok := v2SpuMap[v2Spu.Key]
			if !ok {
				log.Printf("v2 spu %s not found in database, skip", v2Spu.Key)
				continue
			}
			// query v1 spu with same component type and cpu/mem.
			var v1Spu schema.ComponentTier
			if err := db.Where("`component_type` = ? AND `cpu_num` = ? AND `memory_size_gi` = ? AND `version` = ? AND `name` = ?",
				v2Spu.ComponentType, v2Spu.CPUNum, v2Spu.MemorySizeGi, existingSpuVersion, fmt.Sprintf("%dC%dG", v2Spu.CPUNum, v2Spu.MemorySizeGi)).
				First(&v1Spu).Error; err != nil {
				if errors.Is(err, gorm.ErrRecordNotFound) {
					// no v1 spu found, skip to next v2 spu.
					log.Printf("no v1 spu found for v2 spu %s, skip", v2Spu.Key)
					continue
				}
				return err
			}

			// query g4 skus with v1 spu id.
			var g4Skus []*schema.ComponentTierUnit
			if err := db.Where("`component_tier_id` = ? AND `key` LIKE ?", v1Spu.ID, "aws%8g.%").
				Find(&g4Skus).Error; err != nil {
				return err
			}

			// duplicate skus for v2 spu.
			// - set id to 0 to insert new record.
			// - set component_tier_id to v2 spu id.
			// - append :v2 to key to make it unique.
			for _, sku := range g4Skus {
				sku.ComponentTierId = uint64(createdV2Spu.ID)
				sku.Key = fmt.Sprintf("%s:v2", sku.Key) // append :v2 to key to make it unique.
				sku.ID = 0
			}

			// insert duplicated skus.
			if len(g4Skus) > 0 {
				if err := db.Clauses(clause.OnConflict{
					UpdateAll: true,
				}).CreateInBatches(g4Skus, 20).Error; err != nil {
					return err
				}
			}

			// Override the spu mapping.
			var mappings *schema.TierTierUnit
			if err := db.Where("`component_tier_id` = ?", createdV2Spu.ID).FirstOrInit(&mappings).Error; err != nil {
				return err
			}
			mappings.ComponentTierId = uint64(createdV2Spu.ID)
			mappings.TierUnitMapping.Items = lo.Map(g4Skus, func(sku *schema.ComponentTierUnit, _ int) *schema.TierUnitMappingItem {
				return &schema.TierUnitMappingItem{
					ProviderName: "aws",
					RegionNames:  []string{"NAME_UNSPECIFIED"},
					TenantIDs:    []string{}, // default to no tenant. modify it manually.
					SkuKey:       sku.Key,
				}
			})
			if err := db.Save(mappings).Error; err != nil {
				return err
			}
		}

		return nil
	})