package nodemanager

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/pingcap/errors"
	"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"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/tidb"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
	"gorm.io/gorm"
)

const (
	updatePodListLoopInterval               = 5 * time.Second
	reactivateRestartNodeLoopInterval       = 5 * time.Second
	checkAllPoolNeedScalingNodeLoopInterval = 5 * time.Second
	defaultPoolSize                         = 10
	scalingRate                             = 0.8
	defaultNodeCount                        = 1
)

type Request struct {
	TaskID          string
	ActivateRequest tidb.ActivateRequest
	PoolName        rm.WorkerPool
}

func (req *Request) GetTaskType() string {
	return req.ActivateRequest.WorkerRole
}

type manager struct {
	nodeBinder       Binder // It's used to bind worker pod and task.
	PoolManager      rm.PoolManager
	db               *gorm.DB
	isAboveThreshold bool
	lastAboveTime    time.Time
}

type Manager interface {
	Assign(ctx context.Context, keyspaceName string, tasks []*Request) error
	Recycle(ctx context.Context, keyspaceName string, taskID string) error
	GetActiveNodes(ctx context.Context) ([]*Node, error)
	Start(ctx context.Context) error
}

func NewNodeManager(db *gorm.DB, poolManager rm.PoolManager) Manager {
	nodeBinder := NewBinder(db)
	nodeManager := manager{
		nodeBinder:  nodeBinder,
		db:          db,
		PoolManager: retryablePoolManager{poolManager},
	}
	return &nodeManager
}

func NewMockNodeManager(
	mockManager rm.PoolManager,
	db *gorm.DB,
) *manager {
	nodeManager := manager{
		nodeBinder:  NewBinder(db),
		db:          db,
		PoolManager: mockManager,
	}
	return &nodeManager
}

func (m *manager) dbAutoMigrate(ctx context.Context) error {
	return m.db.WithContext(ctx).AutoMigrate(&Node{})
}

func (m *manager) Start(ctx context.Context) error {
	log.Info("start node manager")
	err := m.dbAutoMigrate(ctx)
	if err != nil {
		return err
	}

	if err := m.db.WithContext(ctx).AutoMigrate(&Node{}); err != nil {
		log.Error("failed to auto migrate node table", zap.Error(err))
		return errors.Trace(err)
	}
	g, ctx := errgroup.WithContext(ctx)

	// updateNodeListLoop periodically retrieve the list of nodes for all pools.
	// If there are new nodes, insert them into the node list.
	g.Go(func() error {
		return m.updateNodeListLoop(ctx)
	})
	// reactivateRestartNodeLoop periodically retrieve standby nodes. If a node is bound to a task in the metadata,
	// use the metadata to reactivate this node.
	g.Go(func() error {
		return m.reactivateRestartNodeLoop(ctx)
	})
	// checkAllPoolNeedScalingNodeLoop periodically checks if the current capacity needs to be expanded.
	g.Go(func() error {
		return m.checkAllPoolNeedScalingNodeLoop(ctx)
	})
	return g.Wait()
}

func (m *manager) Assign(ctx context.Context, keyspaceName string, requests []*Request) error {
	log.Debug("before assign multi worker",
		zap.Int("worker-count", defaultNodeCount), zap.Any("requests", requests))
	nodes, err := m.nodeBinder.Bind(ctx, defaultNodeCount, requests)
	if err != nil {
		log.Error("bind worker err", zap.Int("workerCount", defaultNodeCount), zap.Any("requests", requests),
			zap.Error(err))
		return err
	}

	if len(nodes) != len(requests) {
		err = fmt.Errorf("failed to bind nodes, expected %d nodes but got %d", len(requests), len(nodes))
		log.Error("failed to bind nodes", zap.Error(err))
		return errors.Trace(err)
	}

	// To activate node
	for i, node := range nodes {
		request := requests[i]
		err := m.PoolManager.ActivateWorker(ctx, request.PoolName, rm.Worker(node.NodeID), request.ActivateRequest)
		if err != nil {
			log.Error("activate worker node err",
				zap.Int("workerCount", defaultNodeCount), zap.Any("request", request),
				zap.Error(err))
			// Do not return error if only failed at this stage, it will be automatically retried.
		}
	}

	log.Info("assign worker node success", zap.Any("requests", len(requests)), zap.Any("nodes", nodes))

	scaleNodeCount := defaultNodeCount - len(nodes)
	if scaleNodeCount > 0 {
		// todo(ystaticy):scale worker node
	}
	return nil
}

// Recycle if task done, should recycle worker node.
func (m *manager) Recycle(ctx context.Context, keyspaceName string, taskID string) error {
	log.Debug("recycle worker node",
		zap.String("keyspace-name", keyspaceName),
		zap.String("task-id", taskID))

	nodeLock, err := m.nodeBinder.GetNodeByTaskID(ctx, taskID)
	if err != nil {
		return err
	}
	log.Debug("recycle worker node",
		zap.String("keyspace-name", keyspaceName),
		zap.String("task-id", taskID), zap.Any("nodeLock", nodeLock))

	err = m.PoolManager.RecycleWorker(ctx, rm.WorkerPool(nodeLock.PoolName),
		rm.Worker(nodeLock.NodeID), keyspaceName)
	if err != nil {
		return err
	}

	_, err = m.nodeBinder.Unbind(ctx, []string{taskID})
	if err != nil {
		return err
	}

	return nil
}

// GetActiveRequests returns the currently active requests.
func (m *manager) GetActiveNodes(ctx context.Context) ([]*Node, error) {
	log.Debug("get assigned worker node")
	allNodes, err := m.nodeBinder.GetAllNodes(ctx)
	if err != nil {
		return nil, err
	}
	boundNodes := make([]*Node, 0, len(allNodes))
	for _, node := range allNodes {
		if !node.IsBound {
			continue
		}
		boundNodes = append(boundNodes, node)
	}
	return boundNodes, nil
}

// updateNodeListLoop periodically retrieve the list of nodes for all pools.
// If there are new nodes, insert them into the node list.
func (m *manager) updateNodeListLoop(ctx context.Context) error {
	log.Debug("start update node list loop")

	ticker := time.NewTicker(updatePodListLoopInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
			err := m.UpdateNodeList(ctx)
			if err != nil {
				return err
			}
		}
	}
}

func (m *manager) UpdateNodeList(ctx context.Context) error {
	pools, err := m.PoolManager.GetPools()
	if err != nil {
		return err
	}
	for poolName, pool := range pools {
		poolSize := pool.Size
		for i := 0; i < poolSize; i++ {
			nodeID := factory.WorkerHandle(poolName, i)
			err := m.nodeBinder.InsertNewNode(ctx, string(nodeID), string(poolName))
			if err != nil {
				log.Error("failed to insert new node", zap.Error(err))
				continue
			}
		}
	}
	return nil
}

// reactivateRestartNodeLoop Periodically retrieve standby nodes.
// If a node is bound to a task in the metadata, use the metadata to reactivate this node.
func (m *manager) reactivateRestartNodeLoop(ctx context.Context) error {
	log.Debug("start retry activate restart node loop")

	ticker := time.NewTicker(reactivateRestartNodeLoopInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
			err := m.reactivateAllPoolsNodes(ctx)
			if err != nil {
				return err
			}
		}
	}
}

func (m *manager) reactivateAllPoolsNodes(ctx context.Context) error {
	pools, err := m.PoolManager.GetPools()
	if err != nil {
		return err
	}
	for poolName := range pools {
		err := m.reactivateRestartStandbyNode(ctx, poolName)
		if err != nil {
			log.Error("failed to activate standby node", zap.Error(err))
			return err
		}
	}

	return nil
}

// reactivateRestartStandbyNode reactivate the standby node.
func (m *manager) reactivateRestartStandbyNode(ctx context.Context, poolName rm.WorkerPool) error {
	workers, err := m.PoolManager.GetStandbyWorkers(poolName)
	if err != nil {
		return err
	}

	for _, workerName := range workers {
		node, err := m.nodeBinder.GetNodeByNodeID(ctx, string(workerName))
		if err != nil {
			log.Error("failed to get node", zap.Error(err))
			return err
		}
		if !node.IsBound {
			continue
		}
		var request tidb.ActivateRequest
		err = json.Unmarshal([]byte(node.ActivateRequest), &request)
		if err != nil {
			log.Error("failed to unmarshal activation request", zap.Error(err))
			return err
		}
		err = m.PoolManager.ActivateWorker(ctx, poolName, workerName, request)
		if err != nil {
			log.Error("failed to  reactivate node", zap.Error(err))
			continue
		}
	}
	return nil
}

// checkAllPoolNeedScalingNodeLoop Periodically checks if the current capacity needs to be expanded.
func (m *manager) checkAllPoolNeedScalingNodeLoop(ctx context.Context) error {
	log.Debug("start check for pool scale up")

	ticker := time.NewTicker(checkAllPoolNeedScalingNodeLoopInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
			err := m.checkAllPoolNeedScalingNode()
			if err != nil {
				return err
			}
		}
	}
}

func (m *manager) checkAllPoolNeedScalingNode() error {
	pools, err := m.PoolManager.GetPools()
	if err != nil {
		log.Error("get pools error", zap.Error(err))
		return err
	}
	for poolName, pool := range pools {
		err = m.tryScalingNode(poolName, pool.Size)
		if err != nil {
			log.Error("try scaling node error", zap.Error(err))
			return err
		}
	}
	return nil
}

func (m *manager) tryScalingNode(poolName rm.WorkerPool, poolSize int) error {
	nodes, err := m.PoolManager.GetStandbyWorkers(poolName)
	if err != nil {
		log.Error("get standby workers error", zap.Error(err))
		return err
	}
	currentNodes := poolSize - len(nodes)
	if currentNodes > defaultPoolSize*scalingRate {
		if !m.isAboveThreshold {
			// first time above threshold
			m.lastAboveTime = time.Now()
			m.isAboveThreshold = true
		} else if time.Since(m.lastAboveTime) > 10*time.Minute {
			m.ScalingNode()
		}
	} else {
		m.isAboveThreshold = false
	}
	return nil
}

// ScalingNode scale the node according to the required resource configuration
func (m *manager) ScalingNode() {
	// todo call pool manager
	// workerPool , scale count
}