backend/internal/grpc/runner/bus.go

package runner

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

	"with_you/proto/runner"

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

// taskBusMsg Redis 任务频道消息格式
type taskBusMsg struct {
	OriginInstanceID string          `json:"origin_instance_id"`
	TargetRunnerID   string          `json:"target_runner_id"`
	TaskID           string          `json:"task_id"`
	TaskType         string          `json:"task_type"`
	TaskPayload      json.RawMessage `json:"task_payload"`
	CreatedAt        int64           `json:"created_at"`
	TimeoutSeconds   int32           `json:"timeout_seconds"`
}

// resultBusMsg Redis 结果频道消息格式
type resultBusMsg struct {
	OriginInstanceID string          `json:"origin_instance_id"`
	TaskID           string          `json:"task_id"`
	Status           string          `json:"status"`
	CompletedAt      int64           `json:"completed_at"`
	Data             json.RawMessage `json:"data"`
	ErrorCode        string          `json:"error_code"`
	ErrorMessage     string          `json:"error_message"`
}

// TaskBus 组合 RunnerHub + Redis Pub/Sub，实现跨实例任务分发和结果路由
type TaskBus struct {
	hub          *RunnerHub
	registry     *RunnerRegistry
	rdb          *redis.Client
	instanceID   string
	taskChan     string
	resultChan   string
	logger       *zap.Logger

	remoteTasks  sync.Map // taskID → originInstanceID（从其他实例接收的任务）

	sub    *redis.PubSub
	cancel context.CancelFunc
	mu     sync.Mutex
	closed bool
}

// NewTaskBus 创建 TaskBus 实例
func NewTaskBus(hub *RunnerHub, rdb *redis.Client, registry *RunnerRegistry, instanceID, taskChan, resultChan string, logger *zap.Logger) *TaskBus {
	return &TaskBus{
		hub:        hub,
		registry:   registry,
		rdb:        rdb,
		instanceID: instanceID,
		taskChan:   taskChan,
		resultChan: resultChan,
		logger:     logger,
	}
}

// Start 启动 Redis 订阅
func (b *TaskBus) Start() error {
	if b.rdb == nil {
		b.logger.Info("TaskBus: no Redis client, running in standalone mode")
		return nil
	}

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

	sub := b.rdb.Subscribe(ctx, b.taskChan, b.resultChan)
	b.sub = sub

	if _, err := sub.Receive(ctx); err != nil {
		cancel()
		return err
	}

	go b.subscribeLoop(ctx, sub)

	b.logger.Info("TaskBus started",
		zap.String("instance_id", b.instanceID),
		zap.String("task_channel", b.taskChan),
		zap.String("result_channel", b.resultChan),
	)
	return nil
}

// Stop 关闭 TaskBus
func (b *TaskBus) Stop() error {
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.closed {
		return nil
	}
	b.closed = true

	if b.cancel != nil {
		b.cancel()
	}
	if b.sub != nil {
		b.sub.Close()
	}
	if b.registry != nil {
		b.registry.Stop()
	}

	b.logger.Info("TaskBus stopped", zap.String("instance_id", b.instanceID))
	return nil
}

// InstanceID 返回实例 ID
func (b *TaskBus) InstanceID() string {
	return b.instanceID
}

// DispatchTask 实现 TaskDispatcher 接口
// 本地优先 → Redis 注册中心查找 → 远程发布
func (b *TaskBus) DispatchTask(task *runner.Task) error {
	taskTypeStr := task.Type.String()

	// 1. 尝试本地查找
	conn, err := b.hub.GetAvailableRunner(task.Type)
	if err == nil && conn != nil {
		return b.hub.SendTask(conn.ID, task)
	}

	// 2. 查询 Redis 注册中心
	if b.registry != nil {
		runnerID, instanceID, err := b.registry.GetAvailableRunner(taskTypeStr)
		if err != nil {
			return fmt.Errorf("no available runner: %w", err)
		}

		// 如果 runner 在本实例（可能本地 map 还没更新），直接发送
		if instanceID == b.instanceID {
			return b.hub.SendTask(runnerID, task)
		}

		// 发布到远程实例
		return b.publishTask(task, runnerID)
	}

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

// HandleResult 处理任务结果
// 如果任务来自远程实例，转发结果；否则交给 TaskManager
func (b *TaskBus) HandleResult(result *runner.TaskResult) {
	originInstance, ok := b.remoteTasks.Load(result.TaskId)
	if ok {
		b.remoteTasks.Delete(result.TaskId)

		// 任务来自远程实例，转发结果
		if originStr, ok := originInstance.(string); ok && originStr != b.instanceID {
			b.publishResult(result, originStr)
			return
		}
	}

	// 本地任务，交给 TaskManager
	if b.hub.taskManager != nil {
		b.hub.taskManager.HandleResult(result)
	}
}

// publishTask 发布任务到 Redis 频道
func (b *TaskBus) publishTask(task *runner.Task, targetRunnerID string) error {
	msg := taskBusMsg{
		OriginInstanceID: b.instanceID,
		TargetRunnerID:   targetRunnerID,
		TaskID:           task.TaskId,
		TaskType:         task.Type.String(),
		TaskPayload:      task.Payload,
		CreatedAt:        task.CreatedAt,
		TimeoutSeconds:   task.TimeoutSeconds,
	}

	data, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("failed to marshal task message: %w", err)
	}

	if err := b.rdb.Publish(context.Background(), b.taskChan, data).Err(); err != nil {
		return fmt.Errorf("failed to publish task to Redis: %w", err)
	}

	return nil
}

// publishResult 发布结果到 Redis 频道
func (b *TaskBus) publishResult(result *runner.TaskResult, originInstanceID string) error {
	var data json.RawMessage
	if len(result.Data) > 0 {
		data = result.Data
	} else {
		data = json.RawMessage("{}")
	}

	msg := resultBusMsg{
		OriginInstanceID: originInstanceID,
		TaskID:           result.TaskId,
		Status:           result.Status.String(),
		CompletedAt:      result.CompletedAt,
		Data:             data,
		ErrorCode:        result.ErrorCode,
		ErrorMessage:     result.ErrorMessage,
	}

	raw, err := json.Marshal(msg)
	if err != nil {
		b.logger.Error("Failed to marshal result message", zap.Error(err))
		return err
	}

	if err := b.rdb.Publish(context.Background(), b.resultChan, raw).Err(); err != nil {
		b.logger.Error("Failed to publish result to Redis", zap.Error(err))
		return err
	}

	return nil
}

// subscribeLoop 处理 Redis 订阅消息
func (b *TaskBus) subscribeLoop(ctx context.Context, sub *redis.PubSub) {
	ch := sub.Channel()
	for {
		select {
		case <-ctx.Done():
			return
		case msg, ok := <-ch:
			if !ok {
				return
			}
			switch msg.Channel {
			case b.taskChan:
				b.handleRemoteTask(msg)
			case b.resultChan:
				b.handleRemoteResult(msg)
			}
		}
	}
}

// handleRemoteTask 处理从其他实例发来的任务
func (b *TaskBus) handleRemoteTask(msg *redis.Message) {
	var taskMsg taskBusMsg
	if err := json.Unmarshal([]byte(msg.Payload), &taskMsg); err != nil {
		b.logger.Error("TaskBus: failed to unmarshal task message", zap.Error(err))
		return
	}

	// 过滤自己发送的消息
	if taskMsg.OriginInstanceID == b.instanceID {
		return
	}

	// 检查目标 Runner 是否在本实例
	conn, err := b.hub.GetRunnerByID(taskMsg.TargetRunnerID)
	if err != nil {
		b.logger.Warn("TaskBus: target runner not found locally",
			zap.String("runner_id", taskMsg.TargetRunnerID),
			zap.String("task_id", taskMsg.TaskID))
		return
	}

	// 记录任务的来源实例，用于结果回传
	b.remoteTasks.Store(taskMsg.TaskID, taskMsg.OriginInstanceID)

	// 构造 Task 并发送
	task := &runner.Task{
		TaskId:         taskMsg.TaskID,
		Type:           parseTaskType(taskMsg.TaskType),
		CreatedAt:      taskMsg.CreatedAt,
		TimeoutSeconds: taskMsg.TimeoutSeconds,
		Payload:        taskMsg.TaskPayload,
	}

	if err := b.hub.SendTask(conn.ID, task); err != nil {
		b.logger.Error("TaskBus: failed to send remote task to runner",
			zap.String("runner_id", conn.ID),
			zap.String("task_id", taskMsg.TaskID),
			zap.Error(err))
		b.remoteTasks.Delete(taskMsg.TaskID)
	}
}

// handleRemoteResult 处理从其他实例发来的任务结果
func (b *TaskBus) handleRemoteResult(msg *redis.Message) {
	var resultMsg resultBusMsg
	if err := json.Unmarshal([]byte(msg.Payload), &resultMsg); err != nil {
		b.logger.Error("TaskBus: failed to unmarshal result message", zap.Error(err))
		return
	}

	// 过滤自己发送的消息
	if resultMsg.OriginInstanceID == b.instanceID {
		return
	}

	// 只有 origin 匹配本实例才需要处理
	// （originInstanceID 标识的是发起任务的实例）
	// 收到其他实例转发的结果，且 origin 是我们 → 交给 TaskManager
	// 但这里收到的是其他实例广播的结果，我们只是中继
	// 实际上 result 是发到 runner:result channel 的，所有实例都会收到
	// 我们只关心 origin 是本实例的结果
	if resultMsg.OriginInstanceID != b.instanceID {
		return
	}

	// 交给 TaskManager 处理
	result := &runner.TaskResult{
		TaskId:       resultMsg.TaskID,
		Status:       parseTaskStatus(resultMsg.Status),
		CompletedAt:  resultMsg.CompletedAt,
		Data:         resultMsg.Data,
		ErrorCode:    resultMsg.ErrorCode,
		ErrorMessage: resultMsg.ErrorMessage,
	}

	if b.hub.taskManager != nil {
		b.hub.taskManager.HandleResult(result)
	}
}

// parseTaskType 将字符串转换为 runner.TaskType 枚举
func parseTaskType(s string) runner.TaskType {
	switch s {
	case "TASK_TYPE_LOGIN":
		return runner.TaskType_TASK_TYPE_LOGIN
	case "TASK_TYPE_GET_SCHEDULE":
		return runner.TaskType_TASK_TYPE_GET_SCHEDULE
	case "TASK_TYPE_GET_GRADES":
		return runner.TaskType_TASK_TYPE_GET_GRADES
	case "TASK_TYPE_GET_EXAMS":
		return runner.TaskType_TASK_TYPE_GET_EXAMS
	case "TASK_TYPE_GET_USER_INFO":
		return runner.TaskType_TASK_TYPE_GET_USER_INFO
	default:
		return runner.TaskType_TASK_TYPE_UNKNOWN
	}
}

// parseTaskStatus 将字符串转换为 runner.TaskStatus 枚举
func parseTaskStatus(s string) runner.TaskStatus {
	switch s {
	case "TASK_STATUS_SUCCESS":
		return runner.TaskStatus_TASK_STATUS_SUCCESS
	case "TASK_STATUS_FAILED":
		return runner.TaskStatus_TASK_STATUS_FAILED
	case "TASK_STATUS_TIMEOUT":
		return runner.TaskStatus_TASK_STATUS_TIMEOUT
	case "TASK_STATUS_CANCELLED":
		return runner.TaskStatus_TASK_STATUS_CANCELLED
	default:
		return runner.TaskStatus_TASK_STATUS_UNKNOWN
	}
}