package scheduler

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

	"github.com/pingcap/log"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/job"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/scheduler/nodemanager"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/scheduler/resourcelimiter"
	rm "github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/types"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
	"gorm.io/gorm"
	v1 "k8s.io/api/core/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

const (
	poolInterval = 30 * time.Second

	schedulerConfigMap     = "scheduler-config"
	poolKey                = "pool"
	totalMemoryLimitKey    = "memory-limit"
	keyspaceMemoryLimitKey = "keyspace.memory-limit"
	totalCPULimitKey       = "cpu-limit"
	keyspaceCPULimitKey    = "keyspace.cpu-limit"
)

type Scheduler struct {
	nodeManager     nodemanager.Manager
	resourceLimiter resourcelimiter.ResourceLimiter
	poolManager     rm.PoolManager
	namespace       string
	clusterCli      client.Client
	config          *struct {
		sync.RWMutex
		version string
		data    map[string]string
	}
}

func NewScheduler(db *gorm.DB, poolManager rm.PoolManager) (*Scheduler, error) {
	nodeManager := nodemanager.NewNodeManager(db, poolManager)

	scheduler := Scheduler{
		nodeManager:     nodeManager,
		resourceLimiter: resourcelimiter.NewResourceLimiter(poolManager),
		poolManager:     poolManager,
		namespace:       poolManager.GetNamespace(),
		clusterCli:      poolManager.GetClusterClient(),
		config: &struct {
			sync.RWMutex
			version string
			data    map[string]string
		}{
			data: make(map[string]string),
		},
	}
	return &scheduler, nil
}

func (s *Scheduler) AllocateResource(ctx context.Context, req *nodemanager.Request) error {
	log.Info("begin allocate resource ",
		zap.String("keyspace-name", req.ActivateRequest.KeyspaceName),
		zap.String("task-id", req.TaskID))
	var err error
	req.PoolName, err = s.getPool(job.Type(req.ActivateRequest.WorkerRole))
	if err != nil {
		log.Error("failed to get pool",
			zap.String("keyspace-name", req.ActivateRequest.KeyspaceName),
			zap.String("task-id", req.TaskID),
			zap.Error(err))
		return err
	}
	// TODO: batch task
	if err = s.resourceLimiter.Acquire(req.ActivateRequest.KeyspaceName, []string{req.TaskID}, req.PoolName); err != nil {
		log.Error("failed to acquire resource",
			zap.String("keyspace-name", req.ActivateRequest.KeyspaceName),
			zap.String("task-id", req.TaskID),
			zap.Error(err))
		return err
	}
	tasks := []*nodemanager.Request{req}
	err = s.nodeManager.Assign(ctx, req.ActivateRequest.KeyspaceName, tasks)
	if err != nil {
		log.Error("failed to assign worker node",
			zap.String("keyspace-name", req.ActivateRequest.KeyspaceName),
			zap.String("task-id", req.TaskID),
			zap.Error(err))
		if releaseErr := s.resourceLimiter.Release(req.ActivateRequest.KeyspaceName, req.TaskID); releaseErr != nil {
			log.Error("failed to release resource",
				zap.String("keyspace-name", req.ActivateRequest.KeyspaceName),
				zap.String("task-id", req.TaskID),
				zap.Error(releaseErr))
		}
		return err
	}
	return nil
}

func (s *Scheduler) RecycleResource(ctx context.Context, keyspaceName string, taskID string) error {
	err := s.nodeManager.Recycle(ctx, keyspaceName, taskID)
	if err != nil {
		log.Error("recycle worker node failed",
			zap.String("keyspace-name", keyspaceName),
			zap.String("task-id", taskID),
			zap.Error(err))
		return err
	}
	if err = s.resourceLimiter.Release(keyspaceName, taskID); err != nil {
		log.Error("failed to release resource",
			zap.String("keyspace-name", keyspaceName),
			zap.String("task-id", taskID),
			zap.Error(err))
		// Should not return error here, since actual pods has already been released.
	}
	return nil
}

// getPool return the corresponding pool for the job according to job type.
func (s *Scheduler) getPool(taskType job.Type) (rm.WorkerPool, error) {
	s.config.RLock()
	defer s.config.RUnlock()
	pool, ok := s.config.data[strings.Join([]string{string(taskType), poolKey}, ".")]
	if !ok {
		pool = s.config.data[poolKey]
	}
	if pool == "" {
		return "", fmt.Errorf("no pool assigned for task type %s, assignment map %v", taskType, s.config.data)
	}
	return rm.WorkerPool(pool), nil
}

func (s *Scheduler) updateConfig(ctx context.Context) error {
	s.config.Lock()
	defer s.config.Unlock()
	cm := v1.ConfigMap{}
	if err := s.clusterCli.Get(ctx, client.ObjectKey{
		Namespace: s.namespace,
		Name:      schedulerConfigMap,
	}, &cm); err != nil {
		return err
	}
	if cm.ResourceVersion != s.config.version {
		s.config.data = cm.Data
		s.config.version = cm.ResourceVersion
		log.Info("scheduler config updated",
			zap.String("version", s.config.version),
			zap.Any("data", s.config.data),
		)
		totalMem := cm.Data[totalMemoryLimitKey]
		totalCPU := cm.Data[totalCPULimitKey]
		keyspaceMem := cm.Data[keyspaceMemoryLimitKey]
		keyspaceCPU := cm.Data[keyspaceCPULimitKey]
		if err := s.resourceLimiter.UpdateLimit(
			totalMem,
			totalCPU,
			keyspaceMem,
			keyspaceCPU,
		); err != nil {
			log.Error("failed to update resource limit", zap.Error(err))
			return err
		}
	}
	return nil
}

func (s *Scheduler) updateConfigLoop(ctx context.Context) error {
	ticker := time.NewTicker(poolInterval)
	defer ticker.Stop()
	log.Info("start watching scheduler config")
	for {
		select {
		case <-ctx.Done():
			log.Info("stop watching scheduler config, context done")
			return nil
		case <-ticker.C:
			if err := s.updateConfig(ctx); err != nil {
				log.Error("failed to update scheduler configs, will retry later", zap.Error(err))
			}
		}
	}
}

func (s *Scheduler) Start(ctx context.Context) error {
	log.Info("start resource scheduler")
	g, ctx := errgroup.WithContext(ctx)
	// Start pool manager.
	g.Go(func() error {
		return s.poolManager.Start(ctx)
	})
	// Start node manager.
	g.Go(func() error {
		return s.nodeManager.Start(ctx)
	})
	// Start watching scheduler config.
	g.Go(func() error {
		return s.updateConfigLoop(ctx)
	})
	if err := s.resourceLimiter.Initialize(ctx, s.nodeManager); err != nil {
		log.Error("failed to initialize resource limiter", zap.Error(err))
	}
	return g.Wait()
}