package core

import (
	"context"
	"sync"
	"time"

	"github.com/pingcap/log"
	"github.com/tidbcloud/tiflash-autoscale/manager/pkg/config"
	"github.com/tidbcloud/tiflash-autoscale/manager/pkg/k8s"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/kubernetes"
)

type Pools struct {
	conf  *config.Config
	pools sync.Map // key: pool name, value: []*Pods

	k8sClient *kubernetes.Clientset
}

func NewPools(conf *config.Config) *Pools {
	return &Pools{
		conf:      conf,
		pools:     sync.Map{},
		k8sClient: k8s.SingleKubernetesClient(),
	}
}

func (p *Pools) GetPool(ctx context.Context, poolName string) ([]*corev1.Pod, bool) {
	if poolName == "" {
		return nil, false
	}

	pods, ok := p.pools.Load(poolName)
	if !ok {
		log.Warn("Pool not found", zap.String("pool", poolName))
		return nil, false
	}

	return pods.([]*corev1.Pod), true
}

func (p *Pools) GetPools(ctx context.Context) map[string][]*corev1.Pod {
	pools := make(map[string][]*corev1.Pod)
	p.pools.Range(func(key, value interface{}) bool {
		poolName := key.(string)
		pods := value.([]*corev1.Pod)
		pools[poolName] = pods
		return true
	})
	return pools
}

func (p *Pools) Setup(ctx context.Context) *Pools {
	grp, grpCtx := errgroup.WithContext(ctx)
	grp.Go(func() error {
		p.watchPods(grpCtx, p.conf.Tiflash.Main.PoolName, p.conf.Tiflash.Main.StatefulSetName, p.conf.Tiflash.Namespace)
		return nil
	})
	grp.Go(func() error {
		p.watchPods(grpCtx, p.conf.Tiflash.Branch.PoolName, p.conf.Tiflash.Branch.StatefulSetName, p.conf.Tiflash.Namespace)
		return nil
	})

	return p
}

func (p *Pools) watchPods(ctx context.Context, poolName, statefulSetName, namespace string) {
	log.Info("Watching Pods for StatefulSet", zap.String("statefulSet", statefulSetName), zap.String("namespace", namespace), zap.String("poolName", poolName))

	selector, err := k8s.GetStatefulSetSelector(ctx, p.k8sClient, namespace, statefulSetName)
	if selector == nil || err != nil {
		log.Fatal("Failed to get StatefulSet selector:", zap.String("statefulSet", statefulSetName), zap.String("namespace", namespace), zap.String("error", err.Error()))
	}
	watcher, err := p.getPodsWatcher(ctx, namespace, selector.String())
	if err != nil {
		log.Fatal("Failed to create Pod watcher:", zap.String("error", err.Error()))
	}
	defer watcher.Stop()

	// init
	pods, err := p.k8sClient.CoreV1().Pods(namespace).List(ctx, metav1.ListOptions{
		LabelSelector: selector.String(),
	})
	if err != nil {
		log.Error("Failed to list initial pods: ", zap.String("error", err.Error()))
	} else {
		log.Info("Initial Pod status for StatefulSet ", zap.String("statefulSet", statefulSetName))
		var runningPods []*corev1.Pod
		for _, pod := range pods.Items {
			if pod.Status.Phase == corev1.PodRunning {
				runningPods = append(runningPods, &pod)
			}
		}
		p.pools.Store(poolName, runningPods)
	}

	for {
		// watch
		select {
		case event, ok := <-watcher.ResultChan():
			if !ok {
				log.Warn("Pod watcher channel closed, restarting watcher", zap.String("statefulSet", statefulSetName), zap.String("namespace", namespace), zap.String("poolName", poolName))
				time.Sleep(time.Second)
				watcher, _ = p.getPodsWatcher(ctx, namespace, selector.String())
				continue
			}

			pod, ok := event.Object.(*corev1.Pod)
			if !ok {
				continue
			}

			switch event.Type {
			case watch.Added:
				log.Info("Pod added: ", zap.String("namespace", pod.Namespace), zap.String("podName", pod.Name))
				if pod.Status.Phase == corev1.PodRunning {
					pods, _ := p.pools.LoadOrStore(poolName, []*corev1.Pod{})
					p.pools.Store(poolName, addIfNotExists(pods, pod))
				}
			case watch.Modified:
				log.Info("Pod modified: ", zap.String("namespace", pod.Namespace), zap.String("podName", pod.Name))
				if pod.Status.Phase == corev1.PodRunning {
					pods, _ := p.pools.LoadOrStore(poolName, []*corev1.Pod{})
					p.pools.Store(poolName, addIfNotExists(pods, pod))
					continue
				}

				log.Info("Pod is not running, removing from pool", zap.String("namespace", pod.Namespace), zap.String("podName", pod.Name))
				pods, _ := p.pools.Load(poolName)
				var updatedPods []*corev1.Pod
				for _, existingPod := range pods.([]*corev1.Pod) {
					if existingPod.Name != pod.Name {
						updatedPods = append(updatedPods, existingPod)
					}
				}
				p.pools.Store(poolName, updatedPods)
			case watch.Deleted:
				log.Info("Pod deleted: ", zap.String("namespace", pod.Namespace), zap.String("podName", pod.Name))
				pods, _ := p.pools.Load(poolName)
				var updatedPods []*corev1.Pod
				for _, existingPod := range pods.([]*corev1.Pod) {
					if existingPod.Name != pod.Name {
						updatedPods = append(updatedPods, existingPod)
					}
				}
				p.pools.Store(poolName, updatedPods)
			case watch.Error:
				log.Error("Error watching Pod: ", zap.String("namespace", pod.Namespace), zap.String("podName", pod.Name), zap.String("error", event.Object.(error).Error()))
			default:
				log.Warn("Unknown event type:", zap.String("namespace", pod.Namespace), zap.String("podName", pod.Name))
			}

		case <-ctx.Done():
			log.Info("Stopping Pod watcher for StatefulSet", zap.String("statefulSet", statefulSetName), zap.String("namespace", namespace), zap.String("poolName", poolName))
			return
		}

	}
}

func (p *Pools) getPodsWatcher(ctx context.Context, namespace, labelSelector string) (watch.Interface, error) {
	watcher, err := p.k8sClient.CoreV1().Pods(namespace).Watch(ctx, metav1.ListOptions{
		LabelSelector: labelSelector,
	})
	if err != nil {
		log.Error("Failed to watch Pods:", zap.String("error", err.Error()))
		return nil, err
	}
	return watcher, nil
}

func addIfNotExists(pds any, pod *corev1.Pod) []*corev1.Pod {
	if pds == nil {
		return []*corev1.Pod{pod}
	}
	pods, ok := pds.([]*corev1.Pod)
	if !ok {
		log.Error("Failed to cast pods to []*corev1.Pod", zap.Any("pods", pds))
		return []*corev1.Pod{pod}
	}
	for i, existingPod := range pods {
		if existingPod.Name == pod.Name {
			pods[i] = pod
			return pods
		}
	}
	return append(pods, pod)
}