package core

import (
	"context"
	"fmt"
	"slices"
	"sync"
	"time"

	"github.com/pingcap/log"
	"github.com/tidbcloud/tiflash-manager/pkg/config"
	"github.com/tidbcloud/tiflash-manager/pkg/k8s"
	"go.uber.org/zap"
	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

	mu       sync.RWMutex
	pools    map[string][]*corev1.Pod
	watchers map[string]context.CancelFunc

	k8sClient kubernetes.Interface
}

func NewPools(conf *config.Config) *Pools {
	return &Pools{
		conf:      conf,
		pools:     make(map[string][]*corev1.Pod),
		watchers:  make(map[string]context.CancelFunc),
		k8sClient: k8s.SingleKubernetesClient(),
	}
}

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

	p.mu.RLock()
	defer p.mu.RUnlock()

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

	return slices.Clone(pods), true
}

func (p *Pools) GetPools(ctx context.Context) map[string][]*corev1.Pod {
	p.mu.RLock()
	defer p.mu.RUnlock()

	pools := make(map[string][]*corev1.Pod, len(p.pools))
	for poolName, pods := range p.pools {
		pools[poolName] = slices.Clone(pods)
	}
	return pools
}

func (p *Pools) Setup(ctx context.Context) *Pools {
	p.SyncDesiredPools(ctx, map[string]PoolSpec{
		p.conf.Tiflash.DefaultPool: {},
	})
	return p
}

// SyncDesiredPools keeps pod watchers aligned with the current desired pool set.
func (p *Pools) SyncDesiredPools(ctx context.Context, desired map[string]PoolSpec) {
	toStart := make([]string, 0, len(desired))
	toStop := make([]context.CancelFunc, 0)

	p.mu.Lock()
	for poolName := range p.watchers {
		if _, ok := desired[poolName]; ok {
			continue
		}
		toStop = append(toStop, p.watchers[poolName])
		delete(p.watchers, poolName)
		delete(p.pools, poolName)
	}

	for poolName := range desired {
		if poolName == "" {
			continue
		}
		if _, ok := p.watchers[poolName]; ok {
			continue
		}
		watchCtx, cancel := context.WithCancel(ctx)
		p.watchers[poolName] = cancel
		toStart = append(toStart, poolName)
		go p.watchPool(watchCtx, poolName)
	}
	p.mu.Unlock()

	for _, cancel := range toStop {
		cancel()
	}

	for _, poolName := range toStart {
		log.Info("Started pool watcher", zap.String("poolName", poolName))
	}
}

func (p *Pools) watchPool(ctx context.Context, poolName string) {
	statefulSetName := p.conf.StatefulSetName(poolName)
	namespace := p.conf.Tiflash.Namespace
	log.Info("Watching Pods for pool", zap.String("poolName", poolName), zap.String("statefulSet", statefulSetName), zap.String("namespace", namespace))

	for {
		if err := p.runWatchLoop(ctx, poolName, statefulSetName, namespace); err != nil {
			if ctx.Err() != nil {
				break
			}
			log.Warn("Pool watch loop exited, retrying", zap.String("poolName", poolName), zap.String("statefulSet", statefulSetName), zap.Error(err))
			time.Sleep(time.Second)
			continue
		}
		break
	}

	log.Info("Stopping Pod watcher for pool", zap.String("poolName", poolName), zap.String("statefulSet", statefulSetName), zap.String("namespace", namespace))
}

func (p *Pools) runWatchLoop(ctx context.Context, poolName, statefulSetName, namespace string) error {
	selector, err := k8s.GetStatefulSetSelector(ctx, p.k8sClient, namespace, statefulSetName)
	if err != nil {
		return fmt.Errorf("get StatefulSet selector: %w", err)
	}

	if err := p.refreshPool(ctx, poolName, namespace, selector.String()); err != nil {
		return fmt.Errorf("list initial pods: %w", err)
	}

	watcher, err := p.getPodsWatcher(ctx, namespace, selector.String())
	if err != nil {
		return fmt.Errorf("create pod watcher: %w", err)
	}
	defer watcher.Stop()

	for {
		select {
		case <-ctx.Done():
			return nil
		case event, ok := <-watcher.ResultChan():
			if !ok {
				return fmt.Errorf("pod watcher channel closed")
			}
			if err := p.handlePodEvent(ctx, poolName, event); err != nil {
				log.Warn("Failed to handle Pod event", zap.String("poolName", poolName), zap.Error(err))
			}
		}
	}
}

func (p *Pools) refreshPool(ctx context.Context, poolName, namespace, labelSelector string) error {
	pods, err := p.k8sClient.CoreV1().Pods(namespace).List(ctx, metav1.ListOptions{
		LabelSelector: labelSelector,
	})
	if err != nil {
		return err
	}
	if ctx.Err() != nil {
		return ctx.Err()
	}

	runningPods := make([]*corev1.Pod, 0, len(pods.Items))
	for i := range pods.Items {
		pod := pods.Items[i]
		if pod.Status.Phase != corev1.PodRunning {
			continue
		}
		runningPods = append(runningPods, pod.DeepCopy())
	}

	p.mu.Lock()
	p.pools[poolName] = runningPods
	p.mu.Unlock()

	log.Info("Refreshed pool pods", zap.String("poolName", poolName), zap.Int("runningPods", len(runningPods)))
	return nil
}

func (p *Pools) handlePodEvent(ctx context.Context, poolName string, event watch.Event) error {
	if ctx.Err() != nil {
		return ctx.Err()
	}

	pod, ok := event.Object.(*corev1.Pod)
	if !ok {
		if event.Type == watch.Error {
			return fmt.Errorf("watch error event: %T", event.Object)
		}
		return nil
	}

	switch event.Type {
	case watch.Added:
		log.Info("Pod added", zap.String("namespace", pod.Namespace), zap.String("poolName", poolName), zap.String("podName", pod.Name))
		if pod.Status.Phase == corev1.PodRunning {
			p.upsertPod(poolName, pod.DeepCopy())
		}
	case watch.Modified:
		log.Info("Pod modified", zap.String("namespace", pod.Namespace), zap.String("poolName", poolName), zap.String("podName", pod.Name))
		if pod.Status.Phase == corev1.PodRunning {
			p.upsertPod(poolName, pod.DeepCopy())
			return nil
		}
		log.Info("Pod is not running, removing from pool", zap.String("namespace", pod.Namespace), zap.String("poolName", poolName), zap.String("podName", pod.Name))
		p.removePod(poolName, pod.Name)
	case watch.Deleted:
		log.Info("Pod deleted", zap.String("namespace", pod.Namespace), zap.String("poolName", poolName), zap.String("podName", pod.Name))
		p.removePod(poolName, pod.Name)
	case watch.Error:
		return fmt.Errorf("watch error event for pod %s", pod.Name)
	default:
		log.Warn("Unknown Pod event type", zap.String("poolName", poolName), zap.String("podName", pod.Name), zap.String("eventType", string(event.Type)))
	}
	return nil
}

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("labelSelector", labelSelector), zap.Error(err))
		return nil, err
	}
	return watcher, nil
}

func (p *Pools) upsertPod(poolName string, pod *corev1.Pod) {
	p.mu.Lock()
	defer p.mu.Unlock()

	current := slices.Clone(p.pools[poolName])
	for i, existingPod := range current {
		if existingPod.Name == pod.Name {
			current[i] = pod
			p.pools[poolName] = current
			return
		}
	}

	p.pools[poolName] = append(current, pod)
}

func (p *Pools) removePod(poolName, podName string) {
	p.mu.Lock()
	defer p.mu.Unlock()

	current := p.pools[poolName]
	if len(current) == 0 {
		return
	}

	updated := make([]*corev1.Pod, 0, len(current))
	for _, pod := range current {
		if pod.Name != podName {
			updated = append(updated, pod)
		}
	}
	p.pools[poolName] = updated
}