package vpa

import (
	"context"
	"errors"
	"fmt"
	"os"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/shirou/gopsutil/v3/common"
	"github.com/shirou/gopsutil/v3/mem"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	k8slabels "k8s.io/apimachinery/pkg/labels"
	"k8s.io/client-go/informers"
	v1 "k8s.io/client-go/informers/core/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/cache"
	ctrl "sigs.k8s.io/controller-runtime"
)

type podState struct {
	name       string
	namespace  string
	limit      int64
	minLimit   int64
	maxLimit   int64
	used       atomic.Int64
	reportTime atomic.Int64
}

func podKey(namespace, name string) string {
	return namespace + "/" + name
}

func (p *podState) isVIP() bool {
	// export-tidb-pool will be treated as VIP.
	return strings.Contains(p.namespace, "vip") || strings.Contains(p.namespace, "export")
}

// MemoryAllocator is a memory allocator for vpa.
type MemoryAllocator struct {
	capacity                int64
	nodeName                string
	enableKillPod           bool
	measureRealtimeMemUsage bool
	clientset               *kubernetes.Clientset
	podInformer             v1.PodInformer
	nodeInformer            v1.NodeInformer
	stop                    chan struct{}
	ready                   atomic.Bool

	sync.RWMutex
	allocated      int64 // sum of all pods' limits. When it is greater than capacity, we need to delete some pods.
	schedulePaused bool
	pods           map[string]*podState
}

// NewMemoryAllocator creates a new memory allocator.
func NewMemoryAllocator(
	capacity int64,
	enableKillPod bool,
	measureRealtimeMemUsage bool,
	clientset *kubernetes.Clientset,
	nodeInformer v1.NodeInformer,
	nodeName string,
) (*MemoryAllocator, error) {
	podFilter := func(opts *metav1.ListOptions) {
		opts.LabelSelector = fmt.Sprintf("%s=serverless-tidb,%s=true", labels.ComponentLabelKey, types.UseVPALabelKey)
		opts.FieldSelector = fmt.Sprintf("spec.nodeName=%s", nodeName)
	}
	podFactory := informers.NewSharedInformerFactoryWithOptions(
		clientset, 10*time.Minute, informers.WithTweakListOptions(podFilter))
	podInformer := podFactory.Core().V1().Pods()
	podSharedIndexInformer := podInformer.Informer()

	m := &MemoryAllocator{
		capacity:                capacity,
		enableKillPod:           enableKillPod,
		measureRealtimeMemUsage: measureRealtimeMemUsage,
		nodeName:                nodeName,
		clientset:               clientset,
		podInformer:             podInformer,
		nodeInformer:            nodeInformer,
		stop:                    make(chan struct{}),
		pods:                    make(map[string]*podState),
	}
	//nolint:errcheck // not need
	podSharedIndexInformer.AddEventHandler(&cache.ResourceEventHandlerFuncs{
		DeleteFunc: func(obj interface{}) { m.onPodDelete(obj.(*corev1.Pod)) },
	})
	podFactory.Start(m.stop)
	podFactory.WaitForCacheSync(m.stop)
	go m.checkReady()
	if enableKillPod {
		go m.runPodDelete()
	}
	return m, nil
}

func (m *MemoryAllocator) onPodDelete(pod *corev1.Pod) {
	ctrl.Log.Info("pod deleted", "namespace", pod.Namespace, "name", pod.Name)

	m.Lock()
	defer m.Unlock()
	key := podKey(pod.Namespace, pod.Name)
	if p, ok := m.pods[key]; ok {
		m.allocated -= p.limit
		delete(m.pods, key)
		m.onAllocatedUpdated()
	}
}

func (m *MemoryAllocator) checkReady() {
	start := time.Now()
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	for range ticker.C {
		if m.checkReadyOnce() {
			ctrl.Log.Info("memory allocator is ready")
			WaitReadySeconds.Observe(time.Since(start).Seconds())
			m.ready.Store(true)
			return
		}
		if time.Since(start) > 30*time.Second {
			ctrl.Log.Info("memory allocator is not ready, exit process", "waited", time.Since(start))
			os.Exit(1) //nolint:gocritic // ignore ticker.Stop() warning
		}
	}
}

func (m *MemoryAllocator) checkReadyOnce() bool {
	pods, err := m.podInformer.Lister().List(k8slabels.Everything())
	if err != nil {
		ctrl.Log.Error(err, "failed to list pods")
		return false
	}

	m.RLock()
	defer m.RUnlock()

	var reported int
	for _, p := range pods {
		if _, ok := m.pods[podKey(p.Namespace, p.Name)]; ok {
			reported++
		}
	}
	ctrl.Log.Info("check ready", "reported", reported, "total", len(pods))
	return reported == len(pods)
}

// Ready returns true if the memory allocator is ready to serve scale requests.
func (m *MemoryAllocator) Ready() bool {
	return m.ready.Load()
}

func (m *MemoryAllocator) report(pod Pod) {
	m.RLock()
	if p, ok := m.pods[pod.key()]; ok && p.limit == pod.Limit {
		p.used.Store(pod.Used)
		p.reportTime.Store(time.Now().Unix())
		m.RUnlock()
		return
	}
	m.RUnlock()

	m.Lock()
	defer m.Unlock()
	p, ok := m.pods[pod.key()]
	if !ok {
		p = &podState{
			name:      pod.Name,
			namespace: pod.Namespace,
			limit:     pod.Limit,
			minLimit:  pod.MinLimit,
			maxLimit:  pod.MaxLimit,
		}
		m.allocated += p.limit
		m.onAllocatedUpdated()
		m.pods[pod.key()] = p
	}
	if p.limit != pod.Limit {
		m.allocated += pod.Limit - p.limit
		m.onAllocatedUpdated()
		p.limit = pod.Limit
	}
	p.used.Store(pod.Used)
	p.reportTime.Store(time.Now().Unix())
}

func (m *MemoryAllocator) scale(namespace, name string, limit int64) error {
	m.Lock()
	defer m.Unlock()
	p, ok := m.pods[podKey(namespace, name)]
	if !ok {
		return errors.New("wait for report")
	}
	newAllocated := m.allocated - p.limit + limit
	if newAllocated > m.capacity && !m.enableKillPod {
		return fmt.Errorf("allocated memory %d exceeds capacity %d", newAllocated, m.capacity)
	}
	ScalePodHistogram.WithLabelValues(namespace).Observe(float64(limit))
	m.allocated, p.limit = newAllocated, limit
	m.onAllocatedUpdated()
	return nil
}

func (m *MemoryAllocator) runPodDelete() {
	ticker := time.NewTicker(10 * time.Second) //nolint:mnd // no need
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			if !m.Ready() {
				continue
			}
			m.checkMemoryUsage()
			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) //nolint:mnd // no need
			if err := m.ensureUnschedulable(ctx); err != nil {
				ctrl.Log.Error(err, "failed to ensure store taint")
			}
			cancel()
			podsToDelete := m.selectPodsToDelete()
			if len(podsToDelete) > 0 {
				DeletePodHistogram.Observe(float64(len(podsToDelete)))
			}
			for _, p := range podsToDelete {
				ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
				ctrl.Log.Info("[pod lifecycle] start to delete pod", "namespace", p.namespace, "name", p.name)
				err := m.clientset.CoreV1().Pods(p.namespace).Delete(ctx, p.name, metav1.DeleteOptions{})
				if err != nil {
					ctrl.Log.Error(err, "failed to delete pod", "namespace", p.namespace, "name", p.name)
				}
				cancel()
			}
		case <-m.stop:
			return
		}
	}
}

// call it with lock held
func (m *MemoryAllocator) onAllocatedUpdated() {
	counts := make(map[string]int)
	for _, p := range m.pods {
		counts[p.namespace]++
	}
	for ns, count := range counts {
		PodCountGuage.WithLabelValues(ns).Set(float64(count))
	}
	AllocatedGauge.Set(float64(m.allocated))

	m.updateSchedulePaused()
}

// call it with lock held
func (m *MemoryAllocator) updateSchedulePaused() {
	usedMemory := m.allocated
	if m.measureRealtimeMemUsage {
		used, err := getNodeMemoryUsage()
		if err != nil {
			ctrl.Log.Error(err, "fallback to allocated memory")
		} else {
			usedMemory = int64(used)
		}
	}

	if m.schedulePaused && usedMemory <= m.capacity*8/10 {
		m.schedulePaused = false
	} else if !m.schedulePaused && usedMemory > m.capacity {
		m.schedulePaused = true
	}
}

func (m *MemoryAllocator) checkMemoryUsage() {
	if !m.measureRealtimeMemUsage {
		return
	}
	m.Lock()
	defer m.Unlock()
	m.updateSchedulePaused()
}

func (m *MemoryAllocator) getSchedulePaused() bool {
	m.RLock()
	defer m.RUnlock()
	return m.schedulePaused
}

const (
	unschedulableReasonKey   = "tidbcloud.com/unschedulable-reason"
	unschedulableReasonValue = "vpa"
)

func (m *MemoryAllocator) ensureUnschedulable(ctx context.Context) error {
	start := time.Now()
	defer func() {
		EnsureStoreUnschedulableHistogram.Observe(time.Since(start).Seconds())
	}()
	node, err := m.nodeInformer.Lister().Get(m.nodeName)
	if err != nil {
		return err
	}
	if m.getSchedulePaused() {
		if !node.Spec.Unschedulable {
			node = node.DeepCopy()
			node.Annotations[unschedulableReasonKey] = unschedulableReasonValue
			node.Spec.Unschedulable = true
			_, err = m.clientset.CoreV1().Nodes().Update(ctx, node, metav1.UpdateOptions{})
		}
	} else {
		if node.Spec.Unschedulable && node.Annotations[unschedulableReasonKey] == unschedulableReasonValue {
			node = node.DeepCopy()
			delete(node.Annotations, unschedulableReasonKey)
			node.Spec.Unschedulable = false
			_, err = m.clientset.CoreV1().Nodes().Update(ctx, node, metav1.UpdateOptions{})
		}
	}
	return err
}

func (m *MemoryAllocator) selectPodsToDelete() []*podState {
	m.RLock()
	defer m.RUnlock()

	needMemory := m.allocated - m.capacity
	if needMemory <= 0 {
		return nil
	}

	// try to delete only from tidb-pool
	// select pods allocate most memory.
	var tidbPods []*podState
	for _, p := range m.pods {
		if !p.isVIP() {
			tidbPods = append(tidbPods, p)
		}
	}
	sort.Slice(tidbPods, func(i, j int) bool {
		return tidbPods[i].used.Load()/tidbPods[i].limit < tidbPods[j].used.Load()/tidbPods[j].limit
	})
	var total int64
	for i, pod := range tidbPods {
		total += pod.limit
		if total >= needMemory {
			return tidbPods[:i+1]
		}
	}

	// consider delete some pods from vip pool
	// select pods used least memory
	var vipPods []*podState
	for _, p := range m.pods {
		if p.isVIP() {
			vipPods = append(vipPods, p)
		}
	}
	sort.Slice(vipPods, func(i, j int) bool {
		return vipPods[i].used.Load() < vipPods[j].used.Load()
	})
	for i, pod := range vipPods {
		total += pod.limit
		if total >= needMemory {
			pods := append(tidbPods, vipPods[:i+1]...) //nolint: gocritic // on purpose
			sort.Slice(pods, func(i, j int) bool {
				return pods[i].limit > pods[j].limit
			})
			return pods
		}
	}

	ctrl.Log.Error(errors.New("failed to select pods to delete"),
		"failed to select pods to delete",
		"allocated", m.allocated,
		"needMemory", needMemory)
	return nil
}

func getNodeMemoryUsage() (uint64, error) {
	ctx := context.WithValue(context.Background(),
		common.EnvKey, common.EnvMap{common.HostProcEnvKey: "/host_proc"})
	v, err := mem.VirtualMemoryWithContext(ctx)
	if err != nil {
		ctrl.Log.Error(err, "failed to get node memory usage")
		return 0, err
	}
	UsedGauge.Set(float64(v.Used))
	return v.Used, nil
}