package configwatcher

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

	"github.com/pingcap/log"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/poolmanager/factory"
	rm "github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/types"
	"go.uber.org/zap"
	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/util/yaml"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

const (
	poolsConfigKey = "pools"
	workerImageKey = "tidb-image"
)

type configWatcherImpl struct {
	manager    rm.PoolManager
	namespace  string
	clusterCli client.Client
	cfg        *rm.ConfigWatcherConfig

	sync.RWMutex
	configs map[rm.WorkerPool]*rm.PoolConfig
}

// NewConfigWatcher creates a new ConfigWatcher.
func NewConfigWatcher(manager rm.PoolManager, cfg *rm.ConfigWatcherConfig) rm.ConfigWatcher {
	return &configWatcherImpl{
		manager:    manager,
		namespace:  manager.GetNamespace(),
		clusterCli: manager.GetClusterClient(),
		cfg:        cfg,
		configs:    make(map[rm.WorkerPool]*rm.PoolConfig),
	}
}

func (c *configWatcherImpl) GetPools() (map[rm.WorkerPool]*rm.PoolConfig, error) {
	c.RLock()
	defer c.RUnlock()
	return c.configs, nil
}

func (c *configWatcherImpl) Start(ctx context.Context) error {
	ticker := time.NewTicker(c.cfg.PollInterval)
	defer ticker.Stop()
	log.Info("start watching pool configurations")
	for {
		select {
		case <-ctx.Done():
			log.Info("stop watching pool configurations, context done")
			return nil
		case <-ticker.C:
			if err := c.updatePoolConfigs(ctx); err != nil {
				log.Error("failed to update pool configurations, will retry later", zap.Error(err))
			}
		}
	}
}

func (c *configWatcherImpl) updatePoolConfigs(ctx context.Context) error {
	workerImage, err := c.getWorkerImage(ctx)
	if err != nil {
		return err
	}
	poolConfigs, err := c.parsePoolConfig(ctx)
	if err != nil {
		return err
	}
	configs, err := c.buildPoolConfigs(ctx, workerImage, poolConfigs)
	if err != nil {
		return err
	}
	c.Lock()
	defer c.Unlock()
	c.configs = configs
	return nil
}

type poolConfig struct {
	Size      int                      `yaml:"size"`
	Resources *v1.ResourceRequirements `yaml:"resources"`
}

func (c *configWatcherImpl) parsePoolConfig(ctx context.Context) (map[rm.WorkerPool]*poolConfig, error) {
	cm := v1.ConfigMap{}
	if err := c.clusterCli.Get(ctx, client.ObjectKey{
		Namespace: c.namespace,
		Name:      c.cfg.PoolsConfigMap,
	}, &cm); err != nil {
		return nil, err
	}

	configs := make(map[rm.WorkerPool]*poolConfig)
	if err := yaml.Unmarshal([]byte(cm.Data[poolsConfigKey]), &configs); err != nil {
		log.Error("failed to unmarshal pools config", zap.String("cm", cm.String()), zap.Error(err))
		return nil, err
	}
	return configs, nil
}

func (c *configWatcherImpl) getWorkerImage(ctx context.Context) (string, error) {
	cm := v1.ConfigMap{}
	if err := c.clusterCli.Get(ctx, client.ObjectKey{
		Namespace: c.namespace,
		Name:      c.cfg.WorkerConfigMap,
	}, &cm); err != nil {
		return "", err
	}

	data, ok := cm.Data[workerImageKey]
	if !ok {
		return "", fmt.Errorf("worker image not found, cm: %s, key: %s", c.cfg.WorkerConfigMap, workerImageKey)
	}
	return data, nil
}

func (c *configWatcherImpl) buildPoolConfigs(
	ctx context.Context,
	workerImage string,
	poolConfigs map[rm.WorkerPool]*poolConfig,
) (map[rm.WorkerPool]*rm.PoolConfig, error) {
	configs := make(map[rm.WorkerPool]*rm.PoolConfig)
	for pool, pc := range poolConfigs {
		podSpec := factory.WorkerSpec(workerImage, pc.Resources, c.cfg.WorkerConfigMap, c.cfg.WorkerSecret)
		revision, err := factory.GenRevision(ctx, c.namespace, c.clusterCli, podSpec, c.cfg.WorkerConfigMap)
		if err != nil {
			log.Error("failed to generate revision", zap.Error(err))
			return nil, err
		}
		configs[pool] = &rm.PoolConfig{
			WorkerPodSpec: podSpec,
			Revision:      revision,
			Size:          pc.Size,
		}
	}
	return configs, nil
}