// Package capacity helps us estimate capacity requirements.
package capacity

import (
	"database/sql/driver"
	"math"
	"strings"
)

//revive:disable:exported

type DiskType string

// Disk types
const (
	SSDNvme    DiskType = "ssd_nvme"
	SSDScsi    DiskType = "ssd_scsi"
	SSDNetwork DiskType = "ssd_network"
	HDD        DiskType = "HDD"
	HddNetwork DiskType = "hdd_network"
	GP3        DiskType = "gp3"
)

func (dt DiskType) IsNetworked() bool {
	switch dt {
	case SSDNvme, SSDScsi, HDD:
		return false
	case SSDNetwork, HddNetwork, GP3:
		return true
	}
	panic("disk type not known " + string(dt))
}

func (dt DiskType) IsGP3() bool {
	return dt == GP3
}

type MachineType string

func (u MachineType) String() string {
	return string(u)
}

func (u *MachineType) IsN2Type() bool {
	return strings.Contains(u.String(), "n2")
}

// Capacity models Pod capacity requests
type Capacity struct {
	MachineType *MachineType
	Replicas    uint32
	Compute     *Compute
	Disk        *Disk
	Node        *Node
}

type Disk struct {
	SizeGi uint32
	Type   DiskType
}

type Compute struct {
	MemoryMi uint32
	CPU      float32
}

type Node struct {
	NodeProfileID uint64
	Name          string
}

type ComputeRecommendation struct {
	Query   Compute
	Storage Compute
}

func RecommendCompute(TiKVGi uint32) ComputeRecommendation {
	return ComputeRecommendation{
		Query:   MysqlComputeCapacity(TiKVGi),
		Storage: TikvComputeCapacity(TiKVGi),
	}
}

// Assume we store 2TB on a TiKV node when using the recommended capacity
// https://github.com/pingcap/docs/blob/master/op-guide/recommendation.md
// Scale based on the requested storage.
// But establish some minimum constraints for production deployments.
// For our minimum reservations we use 1.7 as CPU instead of 2 to allow for k8s daemons to run.
const fullTiKVStorage = 2 * 1024

var (
	fullTiKVMemoryMi = uint32(math.Trunc(32 * 1024))
	minTiKVMemoryMi  = uint32(math.Trunc(7 * 1024))
)

const (
	fullTiKVCPU = 16
	minTiKVCPU  = 2.3
)

func TikvComputeCapacity(TiKVGi uint32) Compute {
	sizeFactor := float32(TiKVGi) / fullTiKVStorage
	return scaledCompute(sizeFactor, minTiKVCPU, fullTiKVCPU, minTiKVMemoryMi, fullTiKVMemoryMi)
}

var two uint32 = 2

const (
	fullPDStorage = 16
	minPDStorage  = 5
)

var (
	minPDMemoryMi  = uint32(math.Trunc(1 * 1024))
	fullPDMemoryMi = uint32(math.Trunc(8 * 1024))
)

const (
	fullPDCPU = 4.0
	minPDCPU  = 1.4
)

func PdComputeCapacity(TiKVGi uint32) Compute {
	sizeFactor := float32(TiKVGi) / fullTiKVStorage
	return scaledCompute(sizeFactor, minPDCPU, fullPDCPU, minPDMemoryMi, fullPDMemoryMi)
}

func PdStorageCapacity(TiKVGi uint32) uint32 {
	sizeFactor := float32(TiKVGi) / fullTiKVStorage
	storage := uint32(sizeFactor * fullPDStorage)
	var minStorage uint32 = minPDStorage
	if storage < minStorage {
		return minStorage
	}
	return storage
}

var (
	fullMySQLMemoryMi = uint32(math.Trunc(32 * 1024))
	minMySQLMemoryMi  = uint32(math.Trunc(3.7 * 1024))
)

const (
	fullMySQLCPU = 16.0
	minMySQLCPU  = 2.3
)

func scaledCompute(sizeFactor float32,
	minCPU float32, fullCPU float32, minMemoryMi uint32, fullMemoryMi uint32,
) Compute {
	cpu := sizeFactor * fullCPU
	if cpu < minCPU {
		cpu = minCPU
	}
	if cpu > fullCPU {
		cpu = fullCPU
	}
	memoryMb := sizeFactor * float32(fullMemoryMi)
	if memoryMb < float32(minMemoryMi) {
		memoryMb = float32(minMemoryMi)
	}
	if memoryMb > float32(fullMemoryMi) {
		memoryMb = float32(fullMemoryMi)
	}

	return Compute{
		MemoryMi: uint32(memoryMb),
		CPU:      cpu,
	}
}

func MysqlComputeCapacity(TiKVGi uint32) Compute {
	sizeFactor := float32(TiKVGi) / fullTiKVStorage
	return scaledCompute(sizeFactor, minMySQLCPU, fullMySQLCPU, minMySQLMemoryMi, fullMySQLMemoryMi)
}

// nolint: dupl
func MysqlCapacity(compute *Compute, count uint32) Capacity {
	replicas := count
	// Currently for TiDB service, we don't allow tidb less than 2.
	// But in future when we support TiKV service, we can remove this restriction
	if count <= two {
		replicas = two
	}
	return Capacity{
		Replicas: replicas,
		Compute:  compute,
		Disk:     nil,
	}
}

// Scan is for gorm
func (u *MachineType) Scan(value interface{}) error {
	*u = MachineType(value.([]byte))
	return nil
}

// Value is for gorm
func (u MachineType) Value() (driver.Value, error) { return string(u), nil }
