backend/internal/grpc/runner/hub.go

package runner

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

	"with_you/proto/runner"

	"go.uber.org/zap"
)

// RunnerConnection 表示一个已连接的 Runner
type RunnerConnection struct {
	ID           string
	Version      string
	Capabilities map[string]string
	Stream       runner.RunnerHub_ConnectServer
	LastSeen     time.Time
	sendChan     chan *runner.StreamMessage
	logger       *zap.Logger
}

// RunnerHub 管理 Runner 连接和任务分发
type RunnerHub struct {
	runner.UnimplementedRunnerHubServer // 嵌入以实现 gRPC 接口

	mu             sync.RWMutex
	runners        map[string]*RunnerConnection
	taskManager    *TaskManager
	registerChan   chan *RunnerConnection
	unregisterChan chan string
	logger         *zap.Logger
}

// NewRunnerHub 创建新的 RunnerHub
func NewRunnerHub(taskManager *TaskManager, logger *zap.Logger) *RunnerHub {
	return &RunnerHub{
		runners:        make(map[string]*RunnerConnection),
		taskManager:    taskManager,
		registerChan:   make(chan *RunnerConnection, 100),
		unregisterChan: make(chan string, 100),
		logger:         logger,
	}
}

// Connect 实现 RunnerHub 服务的 Connect 方法
func (h *RunnerHub) Connect(stream runner.RunnerHub_ConnectServer) error {
	var conn *RunnerConnection
	ctx := stream.Context()

	defer func() {
		if conn != nil {
			h.unregisterRunner(conn.ID)
		}
	}()

	// 启动发送协程
	sendDone := make(chan struct{})

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-sendDone:
			return nil
		default:
			msg, err := stream.Recv()
			if err != nil {
				h.logger.Debug("Stream receive error", zap.Error(err))
				return err
			}

			switch m := msg.Message.(type) {
			case *runner.StreamMessage_Register:
				conn = h.handleRegister(m.Register, stream)
				if conn != nil {
					// 启动发送协程
					go h.sendLoop(conn, sendDone)
				}
			case *runner.StreamMessage_Heartbeat:
				h.handleHeartbeat(conn, m.Heartbeat)
			case *runner.StreamMessage_Result:
				h.handleTaskResult(m.Result)
			}
		}
	}
}

// sendLoop 处理向 Runner 发送消息
func (h *RunnerHub) sendLoop(conn *RunnerConnection, done chan struct{}) {
	for {
		select {
		case msg, ok := <-conn.sendChan:
			if !ok {
				// Channel 已关闭
				return
			}
			if msg == nil {
				// 忽略 nil 消息
				continue
			}
			if err := conn.Stream.Send(msg); err != nil {
				h.logger.Error("Failed to send message to runner",
					zap.String("runner_id", conn.ID),
					zap.Error(err))
				return
			}
		case <-done:
			return
		}
	}
}

func (h *RunnerHub) handleRegister(req *runner.RegisterRequest, stream runner.RunnerHub_ConnectServer) *RunnerConnection {
	h.logger.Info("Runner registering",
		zap.String("runner_id", req.RunnerId),
		zap.String("version", req.Version),
		zap.Any("capabilities", req.Capabilities))

	conn := &RunnerConnection{
		ID:           req.RunnerId,
		Version:      req.Version,
		Capabilities: req.Capabilities,
		Stream:       stream,
		LastSeen:     time.Now(),
		sendChan:     make(chan *runner.StreamMessage, 100),
		logger:       h.logger,
	}

	h.mu.Lock()
	// 如果已存在旧连接，先关闭
	if oldConn, ok := h.runners[req.RunnerId]; ok {
		close(oldConn.sendChan)
	}
	h.runners[req.RunnerId] = conn
	h.mu.Unlock()

	// 发送注册响应
	if err := conn.send(&runner.StreamMessage{
		Message: &runner.StreamMessage_RegisterResponse{
			RegisterResponse: &runner.RegisterResponse{
				Success:           true,
				SessionId:         conn.ID,
				Message:           "Registered successfully",
				HeartbeatInterval: 30,
			},
		},
	}); err != nil {
		h.logger.Error("Failed to send register response", zap.Error(err))
	}

	h.logger.Info("Runner registered successfully", zap.String("runner_id", req.RunnerId))
	return conn
}

func (h *RunnerHub) handleHeartbeat(conn *RunnerConnection, hb *runner.Heartbeat) {
	if conn == nil {
		return
	}
	conn.LastSeen = time.Now()

	if err := conn.send(&runner.StreamMessage{
		Message: &runner.StreamMessage_HeartbeatAck{
			HeartbeatAck: &runner.HeartbeatAck{
				Timestamp: time.Now().Unix(),
			},
		},
	}); err != nil {
		h.logger.Debug("Failed to send heartbeat ack",
			zap.String("runner_id", conn.ID),
			zap.Error(err))
	}
}

func (h *RunnerHub) handleTaskResult(result *runner.TaskResult) {
	h.logger.Debug("Received task result",
		zap.String("task_id", result.TaskId),
		zap.String("status", result.Status.String()))

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

// unregisterRunner 注销 Runner
func (h *RunnerHub) unregisterRunner(runnerID string) {
	h.mu.Lock()
	if conn, ok := h.runners[runnerID]; ok {
		close(conn.sendChan)
		delete(h.runners, runnerID)
	}
	h.mu.Unlock()

	h.logger.Info("Runner unregistered", zap.String("runner_id", runnerID))
}

// SendTask 向指定 Runner 发送任务
func (h *RunnerHub) SendTask(runnerID string, task *runner.Task) error {
	h.mu.RLock()
	conn, ok := h.runners[runnerID]
	h.mu.RUnlock()

	if !ok {
		return fmt.Errorf("runner not found: %s", runnerID)
	}

	return conn.send(&runner.StreamMessage{
		Message: &runner.StreamMessage_Task{
			Task: task,
		},
	})
}

// GetAvailableRunner 获取支持指定任务类型的可用 Runner
func (h *RunnerHub) GetAvailableRunner(taskType runner.TaskType) (*RunnerConnection, error) {
	h.mu.RLock()
	defer h.mu.RUnlock()

	taskTypeStr := taskType.String()
	for _, conn := range h.runners {
		for capKey, capValue := range conn.Capabilities {
			// 支持两种格式：
			// 1. key = "TASK_TYPE_GET_SCHEDULE", value = "true"
			// 2. key = "schedule", value = "TASK_TYPE_GET_SCHEDULE"
			if capKey == taskTypeStr && capValue == "true" {
				return conn, nil
			}
			if capValue == taskTypeStr {
				return conn, nil
			}
		}
	}
	return nil, fmt.Errorf("no available runner for task type: %s", taskType)
}

// GetAllRunners 获取所有已连接的 Runner 信息
func (h *RunnerHub) GetAllRunners() []map[string]any {
	h.mu.RLock()
	defer h.mu.RUnlock()

	result := make([]map[string]any, 0, len(h.runners))
	for id, conn := range h.runners {
		result = append(result, map[string]any{
			"id":           id,
			"version":      conn.Version,
			"capabilities": conn.Capabilities,
			"last_seen":    conn.LastSeen,
		})
	}
	return result
}

// SetTaskManager 设置任务管理器（用于解决循环依赖）
func (h *RunnerHub) SetTaskManager(tm *TaskManager) {
	h.taskManager = tm
}

// send 向 RunnerConnection 发送消息
func (c *RunnerConnection) send(msg *runner.StreamMessage) error {
	select {
	case c.sendChan <- msg:
		return nil
	default:
		return fmt.Errorf("send buffer full for runner: %s", c.ID)
	}
}

// CheckStaleRunners 检查并清理超时的 Runner 连接
func (h *RunnerHub) CheckStaleRunners(timeout time.Duration) {
	h.mu.Lock()
	defer h.mu.Unlock()

	now := time.Now()
	for id, conn := range h.runners {
		if now.Sub(conn.LastSeen) > timeout {
			h.logger.Warn("Runner timeout, removing",
				zap.String("runner_id", id),
				zap.Time("last_seen", conn.LastSeen))
			close(conn.sendChan)
			delete(h.runners, id)
		}
	}
}

// Shutdown 关闭所有 Runner 连接
func (h *RunnerHub) Shutdown(ctx context.Context) error {
	h.mu.Lock()
	defer h.mu.Unlock()

	for id, conn := range h.runners {
		h.logger.Info("Closing runner connection", zap.String("runner_id", id))
		close(conn.sendChan)
		delete(h.runners, id)
	}

	return nil
}