package nodemanager

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

	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	rm "github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/types"
	"go.uber.org/zap"
	"gorm.io/gorm"
)

type binder struct {
	db *gorm.DB
	mu sync.Mutex
}

// Binder manages worker node and worker task bindings.
type Binder interface {
	InsertNewNode(ctx context.Context, nodeID, nodeType string) error
	Bind(ctx context.Context, expectCount int,
		request []*Request) ([]*Node, error)
	Unbind(ctx context.Context, taskID []string) (*Node, error)
	GetAllNodes(ctx context.Context) ([]*Node, error)
	GetNodeByNodeID(ctx context.Context, nodeID string) (*Node, error)
	GetNodeByTaskID(ctx context.Context, taskID string) (*Node, error)
}

type Node struct {
	NodeID          string    `gorm:"primaryKey"`
	TaskID          string    `gorm:"default:''"`
	KeyspaceName    string    `gorm:"default:''"`
	PoolName        string    `gorm:"default:''"`
	ActivateRequest string    `gorm:"comment:'the message that needs to be transmitted through the lock'"`
	IsBound         bool      `gorm:"default:false"`
	UpdatedAt       time.Time `gorm:"type:timestamp;default:CURRENT_TIMESTAMP"`
}

func (Node) TableName() string {
	return "worker_nodes"
}

func (n *Node) GetWorkerPool() rm.WorkerPool {
	return rm.WorkerPool(n.PoolName)
}

func NewBinder(db *gorm.DB) Binder {
	bindHandler := &binder{
		db: db,
	}
	return bindHandler
}

func (lock *binder) Unbind(ctx context.Context, taskIDs []string) (*Node, error) {
	log.Debug("unbind by task ids", zap.Any("task-ids", taskIDs))
	nodes := &Node{}

	err := lock.getDBWithContext(ctx).Transaction(func(tx *gorm.DB) error {
		for _, taskID := range taskIDs {
			err := tx.Where("task_id = ?", taskID).First(&nodes).Error
			if err != nil {
				return errors.Trace(err)
			}
			log.Info("unbind ", zap.Any("nodes", nodes))
			nodes.IsBound = false
			nodes.ActivateRequest = ""
			nodes.TaskID = ""
			nodes.KeyspaceName = ""
			err = tx.Save(&nodes).Error
			if err != nil {
				return errors.Trace(err)
			}
			return nil
		}
		return nil
	})

	if err != nil {
		return nil, errors.Trace(err)
	}
	return nodes, nil
}

func (lock *binder) Bind(
	ctx context.Context,
	expectCount int,
	requests []*Request,
) ([]*Node, error) {
	lock.mu.Lock()
	defer lock.mu.Unlock()
	log.Debug("try to bind task to node",
		zap.Any("requests", requests))
	var resNodes []*Node

	err := lock.getDBWithContext(ctx).Transaction(func(tx *gorm.DB) error {
		for _, request := range requests {
			nodes, err := lock.getUnboundNodes(tx, expectCount, string(request.PoolName))
			if err != nil {
				return errors.Trace(err)
			}
			log.Debug("bind task to node",
				zap.Any("nodes", nodes), zap.Any("requests", requests))
			for _, lock := range nodes {
				requestBytes, err := json.Marshal(request.ActivateRequest)
				if err != nil {
					log.Error("failed to connect to database", zap.Error(err))
				}
				requestStr := string(requestBytes)
				if err := tx.Model(&lock).
					Update("task_id", request.TaskID).
					Update("is_bound", true).
					Update("activate_request", requestStr).
					Update("keyspace_name", request.ActivateRequest.KeyspaceName).
					Error; err != nil {
					return errors.Trace(err)
				}
			}
			resNodes = append(resNodes, nodes...)
		}
		return nil
	})

	if err != nil {
		return nil, errors.Trace(err)
	}

	return resNodes, nil
}

func (lock *binder) GetNodeByNodeID(ctx context.Context, nodeID string) (*Node, error) {
	var activeNodes Node
	if err := lock.getDBWithContext(ctx).Where("node_id = ? and is_bound = ?",
		nodeID,
		true).
		Find(&activeNodes).Error; err != nil {
		return nil, errors.Trace(err)
	}
	log.Debug("bind get node by node id", zap.String("node-id", nodeID), zap.Any("active-nodes", activeNodes))
	return &activeNodes, nil
}

func (lock *binder) GetNodeByTaskID(ctx context.Context, taskID string) (*Node, error) {
	var activeNodes *Node
	if err := lock.getDBWithContext(ctx).Where("task_id = ? AND is_bound = ?",
		taskID,
		true).
		First(&activeNodes).Error; err != nil {
		if err == gorm.ErrRecordNotFound {
			log.Error("get lock by task id, no record found", zap.Error(err))
		}
		return nil, errors.Trace(err)
	}
	log.Debug("bind get node by task id", zap.String("task-id", taskID), zap.Any("active-nodes", activeNodes))
	return activeNodes, nil
}

// InsertNewNode add a new node.
func (lock *binder) InsertNewNode(ctx context.Context, nodeID, poolName string) error {
	node := Node{NodeID: nodeID, PoolName: poolName}
	log.Debug("bind insert node", zap.Any("node", node))
	result := lock.getDBWithContext(ctx).FirstOrCreate(&node, node)
	if result.Error != nil {
		return errors.Trace(result.Error)
	}
	return nil
}

// GetAllNodes return all Nodes.
func (lock *binder) GetAllNodes(ctx context.Context) ([]*Node, error) {
	var nodes []*Node
	result := lock.getDBWithContext(ctx).Order("node_id asc").Find(&nodes)
	if result.Error != nil {
		return nil, errors.Trace(result.Error)
	}
	log.Debug("bind get all nodes", zap.Any("nodes", nodes))
	return nodes, nil
}

func (lock *binder) getDBWithContext(ctx context.Context) *gorm.DB {
	return lock.db.WithContext(ctx)
}

func (lock *binder) getUnboundNodes(
	tx *gorm.DB,
	expectCount int,
	poolName string) ([]*Node, error) {
	var unboundNodes []*Node
	if err := tx.Where("is_bound = ? AND pool_name = ?", false, poolName).
		Limit(expectCount).Find(&unboundNodes).Error; err != nil {
		return nil, errors.Trace(err)
	}
	log.Debug("bind get unbound nodes", zap.Any("unbound-nodes", unboundNodes))
	return unboundNodes, nil
}