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/util/yaml"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/kubernetes"
)

const specializeKeyInConfigMap = "tiflash-manager-config.yaml"

type TenantsPool struct {
	conf      *config.Config
	pools     sync.Map // key: tenant ID or token, value: pool name
	k8sClient *kubernetes.Clientset
}

type tenantsPoolMap struct {
	Specialize map[string]any `yaml:"specialize"` // tenant ID or token to pool name
}

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

func (t *TenantsPool) Setup(ctx context.Context) *TenantsPool {
	grp, grpCtx := errgroup.WithContext(ctx)
	grp.Go(func() error {
		t.initAndWatchConfigMap(grpCtx)
		return nil
	})

	return t
}

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

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

	return pool.(string)
}

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))

	// init
	configMap, err := t.k8sClient.CoreV1().ConfigMaps(t.conf.Tiflash.Namespace).Get(ctx, t.conf.Tiflash.ConfigMap, metav1.GetOptions{})
	if err != nil {
		log.Fatal("Failed to Get configmap:", zap.String("error", err.Error()))
	}

	poolMap, err := parseTenantsPool(configMap.Data)
	log.Debug("Parsed tenants pool from config map", zap.Any("poolMap", poolMap))
	if err != nil {
		log.Error("Failed to parse tenants pool from config map:", zap.String("error", err.Error()))
	} else {
		storeTenantsPool(poolMap, &t.pools)
	}

	// Watch the config map for changes
	watcher, err := t.getConfigMapWatcher(ctx, t.conf.Tiflash.Namespace, t.conf.Tiflash.ConfigMap)
	if err != nil {
		log.Fatal("Failed to watch configmap:", zap.String("error", err.Error()))
	}

	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.getConfigMapWatcher(ctx, t.conf.Tiflash.Namespace, t.conf.Tiflash.ConfigMap)
				continue
			}

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

			switch event.Type {
			case watch.Added, watch.Modified:
				poolMap, err := parseTenantsPool(configMap.Data)
				if err != nil {
					log.Error("Failed to parse tenants pool from config map:", zap.String("error", err.Error()))
				} else {
					storeTenantsPool(poolMap, &t.pools)
				}
			}
		}
	}
}

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 configmap:", zap.String("error", err.Error()))
		return nil, err
	}

	return watcher, nil
}

func parseTenantsPool(data map[string]string) (*tenantsPoolMap, error) {
	log.Debug("Parsing tenants pool from config map data", zap.Any("data", data))
	var poolMap tenantsPoolMap
	if err := yaml.Unmarshal([]byte(data[specializeKeyInConfigMap]), &poolMap); err != nil {
		return nil, err
	}
	return &poolMap, nil
}

func storeTenantsPool(poolMap *tenantsPoolMap, pools *sync.Map) {
	for tenant, pool := range poolMap.Specialize {
		if pool == "" {
			log.Warn("Tenant pool is empty", zap.String("tenant", tenant))
			continue
		}

		log.Info("Storing tenant in pool", zap.String("tenant", tenant), zap.String("pool", pool.(string)))
		pools.Store(tenant, pool)
	}

	pools.Range(func(key, value interface{}) bool {
		tenant := key.(string)
		pool := value.(string)

		if _, ok := poolMap.Specialize[tenant]; !ok {
			log.Info("Removing tenant from pool", zap.String("tenant", tenant), zap.String("pool", pool))
			pools.Delete(tenant)
		}
		return true
	})
}