package gitops

import (
	"fmt"
	"math"

	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
)

const (
	storageNodeImage = "centos-tidb-v20210420"
	NodeImage        = "centos-v20210420"

	X86ImageVersion     = "7.8.2003"
	Aarch64ImageVersion = "7.8.2003-aarch64"

	defaultAWSVolumeType = "gp2"
	defaultGCPVolumeType = "pd-standard"
	defaultVolumeSize    = "50Gi"

	megaBytes = 1024 * 1024
	gigaBytes = 1024 * 1024 * 1024

	// reservedMemory for system components and daemon pods
	reservedMemory = 1536 * megaBytes

	localStorageClassName = "local-storage"
)

type HardcodeConfigProvider struct{}

var _ ConfigProvider = &HardcodeConfigProvider{}

func NewHardcode() *HardcodeConfigProvider {
	return &HardcodeConfigProvider{}
}

func (m *HardcodeConfigProvider) GetMachineSpec(
	provider Provider,
	c cloudv1.ClusterComponent,
) (*cloudv1.MachineSpec, error) {

	// TODO(https://jira.tidbcloud.com/browse/DM-1576): machineImage should be controlled by GitOPS
	machineImage := NodeImage
	if c.GetComponentType() == cloudv1.ComponentTiKV || c.GetComponentType() == cloudv1.ComponentTiFlash {
		machineImage = storageNodeImage
	}
	// TODO(https://jira.tidbcloud.com/browse/DM-1576): the cloud spec of an instance type should be controlled by GitOPS
	machineType := c.GetInstanceType()
	machineArch, err := m.getMachineArch(provider, machineType)
	if err != nil {
		return nil, fmt.Errorf("failed to get machine arch from provider for component %s/%s: %s", c.GetComponentType(), machineType, err.Error())
	}
	// Check the arch of the given instance type to determine which image version to use.
	machineImageVersion := m.getMachineImageVersion(machineArch)

	var systemVolumeType string
	switch provider.GetProviderType() {
	case consts.ProviderTypeGCP, consts.ProviderTypeGKE:
		systemVolumeType = defaultGCPVolumeType
	case consts.ProviderTypeAWS, consts.ProviderTypeEKS:
		systemVolumeType = defaultAWSVolumeType
	default:
		return nil, fmt.Errorf("unknown provider type %s", provider.GetProviderType())
	}
	return &cloudv1.MachineSpec{
		MachineType: machineType,
		MachineArch: machineArch,
		// The following two fields are gardener specific and can be inferred from MachineArch.
		// However this function is used by the general nodegroup package.
		// So we'd better refactor this to exclude gardener specific logic
		MachineImage:        machineImage,
		MachineImageVersion: machineImageVersion,
		SystemVolumeType:    systemVolumeType,
		SystemVolumeSize:    defaultVolumeSize,
	}, nil
}

func (m *HardcodeConfigProvider) getMachineArch(provider Provider, machineType string) (string, error) {
	providertype := provider.GetProviderType()
	// TODO: support ARM architecture on GCP in the future.
	// It's possible that the machineType passed is empty, for example, the TiCDC component.
	if (providertype != consts.ProviderTypeAWS && providertype != consts.ProviderTypeEKS) || machineType == "" {
		return consts.AMD64ArchName, nil
	}
	// Try to find the arch of the instance type.
	instanceTypes, err := provider.GetInstanceTypes()
	if err != nil {
		return "", fmt.Errorf("can't load instance types")
	}
	for _, instanceType := range instanceTypes {
		if machineType == instanceType.Name {
			switch instanceType.Arch {
			case consts.ARM64ArchName:
				return consts.ARM64ArchName, nil
			default:
				// In order to be compatible with the old version,
				// return x86 image version by default if the arch is not given.
				return consts.AMD64ArchName, nil
			}
		}
	}
	return "", fmt.Errorf("unknown instance type")
}

func (m *HardcodeConfigProvider) getMachineImageVersion(machineArch string) string {
	switch machineArch {
	case consts.ARM64ArchName:
		return Aarch64ImageVersion
	case consts.AMD64ArchName:
		return X86ImageVersion
	}
	return ""
}

func (m *HardcodeConfigProvider) GetResourceSpec(provider Provider, c cloudv1.ClusterComponent) (*cloudv1.ResourceSpec, error) {
	instanceType := c.GetInstanceType()
	providerType := provider.GetProviderType()
	// TODO: storage class should be managed in gitops
	storageClass := consts.DefaultStorageClassName
	if c.GetUseLocalStorage() {
		storageClass = localStorageClassName
	} else if consts.ProviderTypeAWS == providerType || consts.ProviderTypeEKS == providerType ||
		consts.ProviderTypeGKE == providerType {
		switch c.GetComponentType() {
		case cloudv1.ComponentPD:
			storageClass = consts.PDGp3StorageClassName
		case cloudv1.ComponentTiKV:
			storageClass = consts.TiKVGp3StorageClassName
		case cloudv1.ComponentTiFlash:
			storageClass = consts.TiFlashGp3StorageClassName
		case cloudv1.ComponentTiCDC:
			storageClass = consts.TiCDCGp3StorageClassName
		default:
			storageClass = consts.DefaultStorageClassName
		}
	}
	requirements, err := m.GetCPUMemResourceRequirements(provider, instanceType)
	if err != nil {
		return nil, err
	}
	if size, ok := c.GetStorageSize(); ok {
		requirements.Requests[corev1.ResourceStorage] = *size
	}
	return &cloudv1.ResourceSpec{
		Resources:    *requirements,
		StorageClass: storageClass,
	}, nil
}

func (m *HardcodeConfigProvider) GetCPUMemResourceRequirements(provider InstanceTypeProvider, instanceType string) (*corev1.ResourceRequirements, error) {
	instanceSpec, err := provider.FindInstanceType(instanceType)
	if err != nil {
		return nil, err
	}
	// capacity/availableMemory = memory - unavailableMemory
	availableMemory := instanceSpec.Memory
	if !instanceSpec.UnavailableMemory.IsZero() {
		availableMemory.Sub(instanceSpec.UnavailableMemory)
	}
	requests := map[corev1.ResourceName]resource.Quantity{
		corev1.ResourceCPU:    computeCPU(instanceSpec.CPU),
		corev1.ResourceMemory: computeMemory(availableMemory),
	}
	limits := map[corev1.ResourceName]resource.Quantity{
		corev1.ResourceMemory: computeMemoryLimit(availableMemory),
	}
	return &corev1.ResourceRequirements{
		Requests: requests,
		Limits:   limits,
	}, nil
}

func computeCPU(total resource.Quantity) resource.Quantity {
	// reserve 1 core
	total.Sub(resource.MustParse("1"))
	if total.Value() <= 0 {
		total.Set(0)
	}
	return total
}

// ref: https://cloud.google.com/kubernetes-engine/docs/concepts/cluster-architecture#node_allocatable
// and looks like Azure and EKS share the same algorithm https://github.com/awslabs/amazon-eks-ami/issues/387#issuecomment-582419188
// https://docs.microsoft.com/en-us/azure/aks/concepts-clusters-workloads#resource-reservations
var memTaxTable = []struct {
	boundary int64
	rate     float64

	quickDeduction int64
}{
	{
		boundary:       1 * gigaBytes,
		rate:           0.0,
		quickDeduction: 255 * megaBytes,
	},
	{
		boundary:       4 * gigaBytes,
		rate:           0.25,
		quickDeduction: 0,
	},
	{
		boundary:       8 * gigaBytes,
		rate:           0.2,
		quickDeduction: 205 * megaBytes, // 0.05 * 4GB rounded
	},
	{
		boundary:       16 * gigaBytes,
		rate:           0.1,
		quickDeduction: 1 * gigaBytes, // 0.1 * 4GB + 0.15 * 4GB
	},
	{
		boundary:       128 * gigaBytes,
		rate:           0.06,
		quickDeduction: 1680 * megaBytes, // 0.19 * 4GB + 0.14 * 4GB + 0.04 * 8GB rounded
	},
	{
		boundary:       math.MaxInt64,
		rate:           0.02,
		quickDeduction: 6922 * megaBytes, // (0.25 - 0.02) * 4 + (0.2 - 0.02) * 4 + (0.1 - 0.02) * 8 + (0.06 - 0.02) * 112 rounded
	},
}

// computeMemory set a proper memory request for dedicated DB container
func computeMemory(total resource.Quantity) resource.Quantity {
	value := total.Value()
	for _, level := range memTaxTable {
		// ASC order of memTaxTable is guaranteed
		if value <= level.boundary {
			value = value - int64(math.Ceil(float64(value)*level.rate)) - level.quickDeduction
			break
		}
	}
	// subtract the reserved memory for daemon pods
	value -= reservedMemory
	if value <= 0 {
		value = 0
	}
	return *resource.NewQuantity(value, resource.BinarySI)
}

// computeMemoryLimit compute a proper memory limit for dedicated DB container
func computeMemoryLimit(total resource.Quantity) resource.Quantity {
	value := total.Value()
	for _, level := range memTaxTable {
		// ASC order of memTaxTable is guaranteed
		if value <= level.boundary {
			value = value - int64(math.Ceil(float64(value)*level.rate)) - level.quickDeduction
			break
		}
	}
	if value <= 0 {
		value = 0
	}
	return *resource.NewQuantity(value, resource.BinarySI)
}