package core

import (
	"context"
	"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"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/kubernetes"
)

type TenantsPool struct {
	conf       *config.Config
	pools      sync.Map // key: keyspace ID or token, value: pool name
	k8sClient  kubernetes.Interface
	poolCache  *Pools
	reconciler *PoolReconciler
}

func NewTenantsPool(conf *config.Config, poolCache *Pools, reconciler *PoolReconciler) *TenantsPool {
	return &TenantsPool{
		conf:       conf,
		pools:      sync.Map{},
		k8sClient:  k8s.SingleKubernetesClient(),
		poolCache:  poolCache,
		reconciler: reconciler,
	}
}

func (t *TenantsPool) Setup(ctx context.Context) *TenantsPool {
	go t.initAndWatchConfigMap(ctx)
	return t
}

func (t *TenantsPool) GetKeyspacePool(keyspaceID string) string {
	if keyspaceID == "" {
		return t.conf.Tiflash.DefaultPool
	}

	pool, ok := t.pools.Load(keyspaceID)
	if !ok {
		return t.conf.Tiflash.DefaultPool
	}

	return pool.(string)
}

func (t *TenantsPool) GetTenantPool(tenant string) string {
	return t.GetKeyspacePool(tenant)
}

func (t *TenantsPool) initAndWatchConfigMap(ctx context.Context) {
	log.Info("Initializing tenants pool from config map", zap.String("configMapName", t.conf.Tiflash.ConfigMap), zap.String("namespace", t.conf.Tiflash.Namespace))

	// Start from a safe default so the manager can serve the main pool even before
	// the ConfigMap exists.
	t.applyManagerConfig(ctx, &ManagerConfig{
		Keyspaces:  map[string]string{},
		Specialize: map[string]string{},
		Pools:      map[string]PoolSpec{},
	})

	configMap, err := t.k8sClient.CoreV1().ConfigMaps(t.conf.Tiflash.Namespace).Get(ctx, t.conf.Tiflash.ConfigMap, metav1.GetOptions{})
	switch {
	case err == nil:
		cfg, err := parseManagerConfig(configMap.Data)
		if err != nil {
			log.Error("Failed to parse tenants pool from config map", zap.Error(err))
		} else {
			t.applyManagerConfig(ctx, cfg)
		}
	case apierrors.IsNotFound(err):
		log.Warn("Config map not found, continuing with default pool only", zap.String("configMapName", t.conf.Tiflash.ConfigMap))
	case err != nil:
		log.Error("Failed to get config map, continuing with default pool only", zap.Error(err))
	}

	watcher := t.waitForConfigMapWatcher(ctx, t.conf.Tiflash.Namespace, t.conf.Tiflash.ConfigMap)
	if watcher == nil {
		return
	}
	defer watcher.Stop()

	for {
		select {
		case <-ctx.Done():
			log.Info("Stopping config map watcher", zap.String("configMapName", t.conf.Tiflash.ConfigMap), zap.String("namespace", t.conf.Tiflash.Namespace))
			return
		case event, ok := <-watcher.ResultChan():
			if !ok {
				log.Warn("Config map watcher channel closed, restarting watcher", zap.String("configMapName", t.conf.Tiflash.ConfigMap), zap.String("namespace", t.conf.Tiflash.Namespace))
				time.Sleep(time.Second)
				watcher = t.waitForConfigMapWatcher(ctx, t.conf.Tiflash.Namespace, t.conf.Tiflash.ConfigMap)
				if watcher == nil {
					return
				}
				continue
			}

			switch event.Type {
			case watch.Added, watch.Modified:
				configMap, ok := event.Object.(*corev1.ConfigMap)
				if !ok {
					log.Warn("Received non-ConfigMap event", zap.Any("event", event))
					continue
				}

				cfg, err := parseManagerConfig(configMap.Data)
				if err != nil {
					log.Error("Failed to parse keyspace pool config from config map", zap.Error(err))
					continue
				}
				t.applyManagerConfig(ctx, cfg)
			case watch.Deleted:
				t.applyManagerConfig(ctx, &ManagerConfig{
					Keyspaces:  map[string]string{},
					Specialize: map[string]string{},
					Pools:      map[string]PoolSpec{},
				})
			case watch.Error:
				log.Warn("Received ConfigMap watch error event", zap.Any("event", event.Object))
			}
		}
	}
}

func (t *TenantsPool) applyManagerConfig(ctx context.Context, cfg *ManagerConfig) {
	storeTenantsPool(cfg, &t.pools)

	desiredPools := cfg.DesiredPools(t.conf.Tiflash.DefaultPool)
	if err := t.reconciler.Reconcile(ctx, desiredPools); err != nil {
		log.Error("Failed to reconcile desired pools", zap.Error(err))
	}

	t.poolCache.SyncDesiredPools(ctx, desiredPools)
}

func (t *TenantsPool) getConfigMapWatcher(ctx context.Context, namespace, name string) (watch.Interface, error) {
	watcher, err := t.k8sClient.CoreV1().ConfigMaps(namespace).Watch(ctx, metav1.ListOptions{
		FieldSelector: "metadata.name=" + name,
	})
	if err != nil {
		log.Error("Failed to watch config map", zap.Error(err))
		return nil, err
	}

	return watcher, nil
}

func (t *TenantsPool) waitForConfigMapWatcher(ctx context.Context, namespace, name string) watch.Interface {
	for {
		watcher, err := t.getConfigMapWatcher(ctx, namespace, name)
		if err == nil {
			return watcher
		}
		if ctx.Err() != nil {
			return nil
		}

		log.Error("Failed to create config map watcher, retrying", zap.Error(err))
		time.Sleep(time.Second)
	}
}