package migration

import (
	"context"
	"fmt"
	"regexp"

	"github.com/samber/lo"
	"google.golang.org/protobuf/proto"
	"gorm.io/gorm"

	globalcommonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	globaldedicatedv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/dedicated/v1"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/entity"
	spurepo "github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/repository/spu"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/infra/repoimpl/db"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/infra/repoimpl/schema"
	spurepoimpl "github.com/tidbcloud/tidb-management-service/dedicated-global-service/infra/repoimpl/spu"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
)

var keyRegex = regexp.MustCompile(`^(aws:tikv):([^:]+):([^:]+):(\d+):(\d+)(.*)$`)

// Ensure there is a Basic storage type SKU for all Plus storage type SKUs.
var _ = AddDBMigrateRecord(
	"release_20251104",
	func(tx *gorm.DB) error {
		ctx := context.Background()

		_, p, _ := db.NewDBProvider(tx)
		skuRepo := spurepoimpl.NewSkuRepo(p)
		spuRepo := spurepoimpl.NewSpuRepo(p)

		// Fetch storage for gp3 volumes.
		var gp3Storage *schema.StorageTemplate
		if err := tx.Where("volume_type = ?", "gp3").First(&gp3Storage).Error; err != nil {
			return err
		}

		// Query all PD SKUs.
		pdSkus, err := skuRepo.Query(ctx, &spurepo.SkuFilter{
			Component: globaldedicatedv1.ComponentType_PD,
			Provider:  globalcommonv1.CloudProvider_AWS,
		})
		if err != nil {
			return fmt.Errorf("failed to query PD SKUs: %w", err)
		}

		findPd := func(tikvSku *entity.Sku) *entity.Sku {
			return lo.FindOrElse(pdSkus, nil, func(pd *entity.Sku) bool {
				return lo.ContainsBy(pd.GetConditions(), func(c *globaldedicatedv1.Sku_Condition) bool {
					return c.GetKey() == globaldedicatedv1.Sku_Condition_ATTACH_TIERUNIT &&
						c.GetOperator() == globaldedicatedv1.Sku_Condition_IN &&
						lo.Contains(c.GetValues(), tikvSku.GetKey())
				})
			})
		}

		// Query all Plus storage type SKUs.
		tikvSkus, err := skuRepo.Query(ctx, &spurepo.SkuFilter{
			Component: globaldedicatedv1.ComponentType_TIKV,
			Provider:  globalcommonv1.CloudProvider_AWS,
		})
		if err != nil {
			return fmt.Errorf("failed to query TiKV SKUs: %w", err)
		}
		plusSkus := lo.Filter(tikvSkus, func(k *entity.Sku, _ int) bool {
			return k.GetDataStorage().GetVolumeType() == globaldedicatedv1.StorageMu_IO2 &&
				k.GetRaftlogStorage().GetVolumeType() == globaldedicatedv1.StorageMu_VOLUME_TYPE_UNSPECIFIED
		})

		// Filter out those which already have a Basic storage type SKU.
		plusSkus = lo.Filter(plusSkus, func(plusSku *entity.Sku, _ int) bool {
			return !lo.ContainsBy(tikvSkus, func(k *entity.Sku) bool {
				// Basic storage type.
				return k.GetDataStorage().GetVolumeType() == globaldedicatedv1.StorageMu_GP3 &&
					k.GetRaftlogStorage().GetVolumeType() == globaldedicatedv1.StorageMu_VOLUME_TYPE_UNSPECIFIED &&
					// Same SPU ID, instance type, and mapped PD.
					k.GetSpuId() == plusSku.GetSpuId() &&
					k.GetInstanceType() == plusSku.GetInstanceType() &&
					findPd(k).GetId() == findPd(plusSku).GetId()
			})
		})

		setIopsThroughput := func(sku *entity.Sku) {
			switch {
			case sku.GetCpuCore() >= 32:
				sku.DataStorage.Iops = &globaldedicatedv1.Sku_Storage_IopsFixed_{
					IopsFixed: &globaldedicatedv1.Sku_Storage_IopsFixed{
						Iops: 16000,
					},
				}
				sku.DataStorage.Throughput = &globaldedicatedv1.Sku_Storage_ThroughputFixed_{
					ThroughputFixed: &globaldedicatedv1.Sku_Storage_ThroughputFixed{
						ThroughputMibps: 1000,
					},
				}
			case sku.GetCpuCore() >= 16:
				sku.DataStorage.Iops = &globaldedicatedv1.Sku_Storage_IopsFixed_{
					IopsFixed: &globaldedicatedv1.Sku_Storage_IopsFixed{
						Iops: 8000,
					},
				}
				sku.DataStorage.Throughput = &globaldedicatedv1.Sku_Storage_ThroughputFixed_{
					ThroughputFixed: &globaldedicatedv1.Sku_Storage_ThroughputFixed{
						ThroughputMibps: 594,
					},
				}
			case sku.GetCpuCore() >= 8:
				sku.DataStorage.Iops = &globaldedicatedv1.Sku_Storage_IopsFixed_{
					IopsFixed: &globaldedicatedv1.Sku_Storage_IopsFixed{
						Iops: 4000,
					},
				}
				sku.DataStorage.Throughput = &globaldedicatedv1.Sku_Storage_ThroughputFixed_{
					ThroughputFixed: &globaldedicatedv1.Sku_Storage_ThroughputFixed{
						ThroughputMibps: 297,
					},
				}
			default:
				sku.DataStorage.Iops = &globaldedicatedv1.Sku_Storage_IopsFixed_{
					IopsFixed: &globaldedicatedv1.Sku_Storage_IopsFixed{
						Iops: 3000,
					},
				}
				sku.DataStorage.Throughput = &globaldedicatedv1.Sku_Storage_ThroughputFixed_{
					ThroughputFixed: &globaldedicatedv1.Sku_Storage_ThroughputFixed{
						ThroughputMibps: 149,
					},
				}
			}
		}

		for _, plusSku := range plusSkus {
			// Create Basic storage type SKUs for the remaining Plus SKUs.
			basicSku := proto.CloneOf(plusSku)
			basicSku.Id = 0
			basicSku.GetDataStorage().StorageMuId = uint64(gp3Storage.ID)
			setIopsThroughput(basicSku)
			basicSku.Key = keyRegex.ReplaceAllString(
				plusSku.GetKey(),
				fmt.Sprintf(
					"$1:$2:gp3:%d:%d$6",
					basicSku.GetDataStorage().GetIopsFixed().GetIops(),
					basicSku.GetDataStorage().GetThroughputFixed().GetThroughputMibps(),
				),
			)
			if _, err := skuRepo.Save(ctx, basicSku); err != nil && err.ErrCode() != bizerr.ErrCodeRecordExist {
				return fmt.Errorf("failed to create Basic storage type SKU for Plus SKU %s: %w", plusSku.GetKey(), err)
			}

			// Bind new Basic SKU to PD.
			pdSku := findPd(plusSku)
			cdt := lo.FindOrElse(pdSku.GetConditions(), nil, func(c *globaldedicatedv1.Sku_Condition) bool {
				return c.GetKey() == globaldedicatedv1.Sku_Condition_ATTACH_TIERUNIT &&
					c.GetOperator() == globaldedicatedv1.Sku_Condition_IN
			})
			if cdt == nil {
				return fmt.Errorf("PD SKU %s missing ATTACH_TIERUNIT condition", pdSku.GetKey())
			}
			cdt.Values = lo.Uniq(append(cdt.Values, basicSku.GetKey()))
			if _, err := skuRepo.Save(ctx, pdSku); err != nil {
				return fmt.Errorf("failed to update PD SKU %s to bind new Basic SKU %s: %w", pdSku.GetKey(), basicSku.GetKey(), err)
			}

			// Add the new SKU to mapping.
			mappings, err := skuRepo.QuerySpuMappings(ctx, []uint64{plusSku.GetSpuId()})
			if err != nil {
				return fmt.Errorf("failed to query SPU mapping for SKU %s: %w", plusSku.GetKey(), err)
			}
			mapping := mappings[plusSku.GetSpuId()]
			plusMapping, ok := lo.Find(mapping, func(m *entity.SpuMapping) bool {
				return m.SkuKey == plusSku.GetKey()
			})
			if ok {
				mapping = append(mapping, &entity.SpuMapping{
					SpuId:     plusMapping.SpuId,
					SkuKey:    basicSku.GetKey(),
					Provider:  plusMapping.Provider,
					Regions:   plusMapping.Regions,
					TenantIDs: plusMapping.TenantIDs,
				})
				if _, err := spuRepo.SaveSpuMappings(ctx, mapping); err != nil {
					return fmt.Errorf("failed to update SPU mapping to add new Basic SKU %s: %w", basicSku.GetKey(), err)
				}
			}
		}

		return nil
	},
)