package reconciler

import (
	"context"
	"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"
	providertypes "github.com/tidbcloud/aws-shared-provider/pkg/types"
	"go.uber.org/zap"
	v1 "k8s.io/api/core/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

type reconcilerImpl struct {
	manager    rm.PoolManager
	namespace  string
	clusterCli client.Client
	cfg        *rm.ReconcilerConfig
}

func NewReconciler(manager rm.PoolManager, cfg *rm.ReconcilerConfig) rm.Reconciler {
	return &reconcilerImpl{
		manager:    manager,
		namespace:  manager.GetNamespace(),
		clusterCli: manager.GetClusterClient(),
		cfg:        cfg,
	}
}

func (r *reconcilerImpl) Start(ctx context.Context) error {
	ticker := time.NewTicker(r.cfg.ReconcileInterval)
	defer ticker.Stop()
	log.Info("start reconciling worker pools")

	for {
		select {
		case <-ctx.Done():
			log.Info("stop reconciling worker pools, context done")
			return nil
		case <-ticker.C:
			if err := r.reconcile(ctx); err != nil {
				log.Error("failed to reconcile worker pools, will retry later", zap.Error(err))
			}
		}
	}
}

func (r *reconcilerImpl) reconcile(ctx context.Context) error {
	poolsConfig, err := r.manager.GetPools()
	if err != nil {
		log.Error("failed to get pools config", zap.Error(err))
		return err
	}
	// If config-watcher not initialized, return directly.
	if len(poolsConfig) == 0 {
		log.Info("no pool config found, waiting for config watcher")
		return nil
	}

	// Remove delete pool.
	if err = r.reconcileExtraPool(ctx, poolsConfig); err != nil {
		log.Error("failed to remove extra pools", zap.Error(err))
		return err
	}

	// Reconcile each pool.
	for pool, poolConfig := range poolsConfig {
		if err = r.reconcilePool(ctx, pool, poolConfig); err != nil {
			log.Error("failed to reconcile pool", zap.Any("pool", pool), zap.Error(err))
			return err
		}
	}
	return nil
}

func (r *reconcilerImpl) reconcilePool(ctx context.Context, pool rm.WorkerPool, poolConfig *rm.PoolConfig) error {
	if err := r.reconcileRevision(ctx, pool, poolConfig); err != nil {
		log.Error("failed to remove outdated workers", zap.Any("pool", pool), zap.Error(err))
		return err
	}
	if err := r.reconcileIndex(ctx, pool, poolConfig); err != nil {
		log.Error("failed to remove excessive workers", zap.Any("pool", pool), zap.Error(err))
		return err
	}

	if err := r.reconcileMissingWorkers(ctx, pool, poolConfig); err != nil {
		log.Error("failed to create missing workers", zap.Any("pool", pool), zap.Error(err))
		return err
	}

	if err := r.reconcileMissingService(ctx, pool, poolConfig); err != nil {
		log.Error("failed to create missing service", zap.Any("pool", pool), zap.Error(err))
		return err
	}

	return nil
}

// reconcileExtraPool removes worker with invalid pool.
func (r *reconcilerImpl) reconcileExtraPool(ctx context.Context, poolsConfig map[rm.WorkerPool]*rm.PoolConfig) error {
	invalidPoolSelector, err := invalidPool(poolsConfig)
	if err != nil {
		return err
	}
	if err = r.clusterCli.DeleteAllOf(ctx, &v1.Service{},
		client.InNamespace(r.namespace),
		client.MatchingLabelsSelector{Selector: invalidPoolSelector},
	); err != nil {
		return err
	}
	return r.clusterCli.DeleteAllOf(ctx, &v1.Pod{},
		client.InNamespace(r.namespace),
		client.MatchingLabelsSelector{Selector: invalidPoolSelector},
	)
}

// reconcileRevision removes worker with outdated revision.
func (r *reconcilerImpl) reconcileRevision(ctx context.Context, pool rm.WorkerPool, poolConfig *rm.PoolConfig) error {
	outdatedRevisionSelector, err := outdatedRevision(pool, poolConfig.Revision)
	if err != nil {
		return err
	}
	return r.clusterCli.DeleteAllOf(ctx, &v1.Pod{},
		client.InNamespace(r.namespace),
		client.MatchingLabelsSelector{Selector: outdatedRevisionSelector},
	)
}

// reconcileIndex removes worker and service with invalid index.
func (r *reconcilerImpl) reconcileIndex(ctx context.Context, pool rm.WorkerPool, poolConfig *rm.PoolConfig) error {
	invalidIndexSelector, err := invalidIndex(pool, poolConfig.Size)
	if err != nil {
		return err
	}
	// Delete worker.
	if err := r.clusterCli.DeleteAllOf(ctx, &v1.Pod{},
		client.InNamespace(r.namespace),
		client.MatchingLabelsSelector{Selector: invalidIndexSelector},
	); err != nil {
		return err
	}
	// Delete service.
	return r.clusterCli.DeleteAllOf(ctx, &v1.Service{},
		client.InNamespace(r.namespace),
		client.MatchingLabelsSelector{Selector: invalidIndexSelector},
	)
}

func (r *reconcilerImpl) reconcileMissingWorkers(
	ctx context.Context,
	pool rm.WorkerPool,
	poolConfig *rm.PoolConfig,
) error {
	missingWorkers := make(map[rm.Worker]struct{})
	for i := 0; i < poolConfig.Size; i++ {
		missingWorkers[factory.WorkerHandle(pool, i)] = struct{}{}
	}
	labelSet := factory.WorkerLabels()
	labelSet[providertypes.WorkerPoolLabelKey] = string(pool)
	currentWorkers := &v1.PodList{}
	if err := r.clusterCli.List(ctx, currentWorkers,
		client.InNamespace(r.namespace),
		client.MatchingLabels(labelSet),
	); err != nil {
		return err
	}
	for _, worker := range currentWorkers.Items {
		delete(missingWorkers, rm.Worker(worker.Labels[providertypes.WorkerLabelKey]))
	}
	for worker := range missingWorkers {
		workerPod := factory.WorkerPod(
			worker,
			pool,
			r.namespace,
			poolConfig.Revision,
			poolConfig.WorkerPodSpec,
		)
		if err := r.clusterCli.Create(ctx, workerPod); err != nil {
			return err
		}
	}
	return nil
}

func (r *reconcilerImpl) reconcileMissingService(
	ctx context.Context,
	pool rm.WorkerPool,
	poolConfig *rm.PoolConfig,
) error {
	missingServices := make(map[rm.Worker]struct{})
	for i := 0; i < poolConfig.Size; i++ {
		missingServices[factory.WorkerHandle(pool, i)] = struct{}{}
	}
	labelSet := factory.WorkerLabels()
	labelSet[providertypes.WorkerPoolLabelKey] = string(pool)
	currentServices := &v1.ServiceList{}
	if err := r.clusterCli.List(ctx, currentServices,
		client.InNamespace(r.namespace),
		client.MatchingLabels(labelSet),
	); err != nil {
		return err
	}
	for _, service := range currentServices.Items {
		delete(missingServices, rm.Worker(service.Labels[providertypes.WorkerLabelKey]))
	}
	for worker := range missingServices {
		workerSvc := factory.WorkerSvc(worker, pool, r.namespace)
		if err := r.clusterCli.Create(ctx, workerSvc); err != nil {
			return err
		}
	}
	return nil
}