backend/internal/grpc/runner/registry.go

package runner

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

	"github.com/redis/go-redis/v9"
	"go.uber.org/zap"
)

// RunnerRegistry 分布式 Runner 注册中心
// 使用 Redis Hash 存储 Runner 信息，Redis Set 按能力索引
type RunnerRegistry struct {
	rdb            *redis.Client
	instanceID     string
	registryPrefix string
	capPrefix      string
	onlineTTL      time.Duration
	heartbeatInterval time.Duration
	localRunners   sync.Map // runnerID → map[string]string (capabilities)
	logger         *zap.Logger
	cancel         context.CancelFunc
}

// runnerInfo 存储在 Redis Hash 中的 Runner 信息
type runnerInfo struct {
	InstanceID   string `json:"instance_id"`
	Version      string `json:"version"`
	Capabilities string `json:"capabilities"` // JSON-encoded map
}

// NewRunnerRegistry 创建 Runner 注册中心
func NewRunnerRegistry(rdb *redis.Client, instanceID, registryPrefix, capPrefix string, onlineTTL, heartbeatInterval time.Duration, logger *zap.Logger) *RunnerRegistry {
	return &RunnerRegistry{
		rdb:               rdb,
		instanceID:        instanceID,
		registryPrefix:    registryPrefix,
		capPrefix:         capPrefix,
		onlineTTL:         onlineTTL,
		heartbeatInterval: heartbeatInterval,
		logger:            logger,
	}
}

// Register 注册 Runner 到分布式注册中心
func (r *RunnerRegistry) Register(runnerID, version string, capabilities map[string]string) error {
	if r.rdb == nil {
		return nil
	}

	ctx := context.Background()
	capJSON, _ := json.Marshal(capabilities)
	key := r.registryPrefix + runnerID

	info := runnerInfo{
		InstanceID:   r.instanceID,
		Version:      version,
		Capabilities: string(capJSON),
	}
	infoJSON, _ := json.Marshal(info)

	pipe := r.rdb.Pipeline()
	pipe.Set(ctx, key, infoJSON, r.onlineTTL)

	// 为每个能力创建索引
	taskTypeStr := parseTaskTypes(capabilities)
	for _, tt := range taskTypeStr {
		capKey := r.capPrefix + tt
		pipe.SAdd(ctx, capKey, runnerID)
		pipe.Expire(ctx, capKey, r.onlineTTL*2)
	}

	if _, err := pipe.Exec(ctx); err != nil {
		return fmt.Errorf("failed to register runner in redis: %w", err)
	}

	r.localRunners.Store(runnerID, capabilities)
	return nil
}

// Unregister 从分布式注册中心移除 Runner
func (r *RunnerRegistry) Unregister(runnerID string, capabilities map[string]string) {
	if r.rdb == nil {
		return
	}

	ctx := context.Background()
	key := r.registryPrefix + runnerID
	pipe := r.rdb.Pipeline()
	pipe.Del(ctx, key)

	taskTypeStr := parseTaskTypes(capabilities)
	for _, tt := range taskTypeStr {
		capKey := r.capPrefix + tt
		pipe.SRem(ctx, capKey, runnerID)
	}

	if _, err := pipe.Exec(ctx); err != nil {
		r.logger.Error("Failed to unregister runner from redis",
			zap.String("runner_id", runnerID),
			zap.Error(err))
	}

	r.localRunners.Delete(runnerID)
}

// RefreshHeartbeat 刷新 Runner 的 TTL
func (r *RunnerRegistry) RefreshHeartbeat(runnerID string) error {
	if r.rdb == nil {
		return nil
	}

	ctx := context.Background()
	key := r.registryPrefix + runnerID
	return r.rdb.Expire(ctx, key, r.onlineTTL).Err()
}

// GetAvailableRunner 查找支持指定任务类型的 Runner
// 返回 runnerID 和所在 instanceID
func (r *RunnerRegistry) GetAvailableRunner(taskType string) (runnerID, instanceID string, err error) {
	if r.rdb == nil {
		return "", "", fmt.Errorf("registry not available")
	}

	ctx := context.Background()
	capKey := r.capPrefix + taskType

	members, err := r.rdb.SMembers(ctx, capKey).Result()
	if err != nil {
		return "", "", fmt.Errorf("failed to query runner capabilities: %w", err)
	}

	if len(members) == 0 {
		return "", "", fmt.Errorf("no runner available for task type: %s", taskType)
	}

	// 随机选一个（伪随机，map遍历本身随机性）
	for _, rid := range members {
		key := r.registryPrefix + rid
		val, err := r.rdb.Get(ctx, key).Bytes()
		if err != nil {
			continue // 已过期，跳过
		}

		var info runnerInfo
		if err := json.Unmarshal(val, &info); err != nil {
			continue
		}

		return rid, info.InstanceID, nil
	}

	return "", "", fmt.Errorf("no runner available for task type: %s", taskType)
}

// GetRunnerInstance 获取 Runner 所在的实例 ID
func (r *RunnerRegistry) GetRunnerInstance(runnerID string) (string, error) {
	if r.rdb == nil {
		return "", fmt.Errorf("registry not available")
	}

	ctx := context.Background()
	key := r.registryPrefix + runnerID
	val, err := r.rdb.Get(ctx, key).Bytes()
	if err != nil {
		return "", fmt.Errorf("runner not found: %s", runnerID)
	}

	var info runnerInfo
	if err := json.Unmarshal(val, &info); err != nil {
		return "", fmt.Errorf("failed to parse runner info: %w", err)
	}

	return info.InstanceID, nil
}

// GetAllRunners 获取所有注册的 Runner 信息
func (r *RunnerRegistry) GetAllRunners() ([]map[string]any, error) {
	if r.rdb == nil {
		return nil, nil
	}

	ctx := context.Background()
	pattern := r.registryPrefix + "*"
	keys, err := r.rdb.Keys(ctx, pattern).Result()
	if err != nil {
		return nil, err
	}

	result := make([]map[string]any, 0, len(keys))
	for _, key := range keys {
		val, err := r.rdb.Get(ctx, key).Bytes()
		if err != nil {
			continue
		}
		var info runnerInfo
		if err := json.Unmarshal(val, &info); err != nil {
			continue
		}

		var caps map[string]string
		json.Unmarshal([]byte(info.Capabilities), &caps)

		runnerID := key[len(r.registryPrefix):]
		result = append(result, map[string]any{
			"id":           runnerID,
			"instance_id":  info.InstanceID,
			"version":      info.Version,
			"capabilities": caps,
		})
	}
	return result, nil
}

// StartHeartbeat 启动心跳续期
func (r *RunnerRegistry) StartHeartbeat() {
	if r.rdb == nil {
		return
	}

	ctx, cancel := context.WithCancel(context.Background())
	r.cancel = cancel

	go func() {
		ticker := time.NewTicker(r.heartbeatInterval)
		defer ticker.Stop()

		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				r.heartbeat()
			}
		}
	}()
}

func (r *RunnerRegistry) heartbeat() {
	ctx := context.Background()
	pipe := r.rdb.Pipeline()
	count := 0

	r.localRunners.Range(func(key, _ any) bool {
		runnerID := key.(string)
		regKey := r.registryPrefix + runnerID
		pipe.Expire(ctx, regKey, r.onlineTTL)
		count++
		return true
	})

	if count > 0 {
		if _, err := pipe.Exec(ctx); err != nil {
			r.logger.Error("Failed to heartbeat runner registry", zap.Error(err))
		}
	}
}

// Stop 停止心跳并清理本地 Runner 的 Redis 注册
func (r *RunnerRegistry) Stop() {
	if r.cancel != nil {
		r.cancel()
	}

	if r.rdb == nil {
		return
	}

	r.localRunners.Range(func(key, value any) bool {
		runnerID := key.(string)
		capabilities, _ := value.(map[string]string)
		r.Unregister(runnerID, capabilities)
		return true
	})
}

// parseTaskTypes 从 capabilities map 中提取支持的任务类型列表
func parseTaskTypes(capabilities map[string]string) []string {
	var types []string
	for capKey, capValue := range capabilities {
		// key = "TASK_TYPE_GET_SCHEDULE", value = "true"
		if capValue == "true" && isTaskType(capKey) {
			types = append(types, capKey)
			continue
		}
		// key = "schedule", value = "TASK_TYPE_GET_SCHEDULE"
		if isTaskType(capValue) {
			types = append(types, capValue)
		}
	}
	return types
}

// isTaskType 检查字符串是否为有效的 TaskType 枚举值
func isTaskType(s string) bool {
	switch s {
	case "TASK_TYPE_LOGIN", "TASK_TYPE_GET_SCHEDULE", "TASK_TYPE_GET_GRADES",
		"TASK_TYPE_GET_EXAMS", "TASK_TYPE_GET_USER_INFO":
		return true
	}
	return false
}