EveryPublish/server/internal/ws/dispatcher.go

package ws

import (
	"encoding/json"
	"log"
	"strconv"
	"time"

	"github.com/google/uuid"
	"gorm.io/gorm"

	"everypublish/server/internal/config"
	"everypublish/server/internal/models"
	"everypublish/server/internal/task"
	"everypublish/shared/proto"
)

// Dispatcher 任务下发器（MVP：进程内轮询 + 审批/重试即时触发；接口预留二期换 asynq）
type Dispatcher struct {
	DB       *gorm.DB
	Hub      *Hub
	Cfg      *config.Config
	Interval time.Duration
}

// NewDispatcher 构造下发器
func NewDispatcher(db *gorm.DB, hub *Hub, cfg *config.Config) *Dispatcher {
	return &Dispatcher{DB: db, Hub: hub, Cfg: cfg, Interval: time.Second}
}

// Start 后台轮询：把到期且排队的任务下发给在线设备
func (d *Dispatcher) Start(stop <-chan struct{}) {
	if d.Interval <= 0 {
		d.Interval = time.Second
	}
	ticker := time.NewTicker(d.Interval)
	defer ticker.Stop()
	for {
		select {
		case <-stop:
			return
		case <-ticker.C:
			d.poll()
		}
	}
}

func (d *Dispatcher) poll() {
	var rows []models.Task
	d.DB.Where("status = ?", task.Queued).
		Where("(schedule_at is null or schedule_at <= ?)", time.Now()).
		Order("priority desc, id asc").Limit(20).Find(&rows)
	for i := range rows {
		if d.TryDispatch(rows[i].ID) {
			log.Printf("[dispatch] task %d pushed", rows[i].ID)
		}
	}
}

// TryDispatch 立即下发单个任务（审批/重试后调用）；无在线设备返回 false
func (d *Dispatcher) TryDispatch(taskID uint64) bool {
	var row models.Task
	if err := d.DB.First(&row, taskID).Error; err != nil || row.Status != task.Queued {
		return false
	}
	if row.ScheduleAt != nil && row.ScheduleAt.After(time.Now()) {
		return false
	}
	devID := d.pickDevice(&row)
	if devID == 0 {
		return false
	}
	push := d.buildPush(&row, devID)
	if push == nil {
		return false
	}
	// 先落库 dispatched，再推送，避免 ack 早于状态落库的竞态
	if next, ok := task.Next(row.Status, task.ActionDispatch); ok {
		_ = d.DB.Model(&row).Update("status", next).Error
		row.Status = next
	}
	env := proto.NewEnvelope(uuid.NewString(), proto.TypeTaskPush, push)
	if err := d.Hub.SendToAgent(devID, env); err != nil {
		// 下发失败（设备离线/队列满）回退 queued，避免卡在 dispatched
		_ = d.DB.Model(&row).Update("status", task.Queued).Error
		return false
	}
	d.Hub.BroadcastWS(row.WorkspaceID, "task.status", taskEvent(row))
	return true
}

// pickDevice 选设备：账号已绑定设备则必须在线上才下发（保持一账号一IP），
// 未绑定设备的账号才兜底任意在线设备。
func (d *Dispatcher) pickDevice(row *models.Task) uint64 {
	var accountIDs []uint64
	_ = json.Unmarshal([]byte(row.AccountIDs), &accountIDs)
	if len(accountIDs) > 0 {
		var accounts []models.Account
		d.DB.Where("id in ? and workspace_id = ?", accountIDs, row.WorkspaceID).Find(&accounts)
		hasAssigned := false
		for _, a := range accounts {
			if a.AgentDeviceID == 0 {
				continue
			}
			hasAssigned = true
			if d.Hub.AgentOnline(a.AgentDeviceID) {
				return a.AgentDeviceID
			}
		}
		if hasAssigned {
			return 0 // 已绑定设备但离线：不下发，保持 IP 隔离
		}
	}
	online := d.Hub.OnlineDevices(row.WorkspaceID)
	if len(online) > 0 {
		return online[0]
	}
	return 0
}

// buildPush 组装 TaskPush（素材生成一次性签名直链）
func (d *Dispatcher) buildPush(row *models.Task, devID uint64) *proto.TaskPush {
	var materialIDs []uint64
	_ = json.Unmarshal([]byte(row.MaterialIDs), &materialIDs)
	urls := make([]string, 0, len(materialIDs))
	for _, mid := range materialIDs {
		var mat models.Material
		if err := d.DB.First(&mat, mid).Error; err != nil {
			continue
		}
		token := models.TransferToken{MaterialID: mat.ID, Token: uuid.NewString(), ExpiresAt: time.Now().Add(10 * time.Minute)}
		if err := d.DB.Create(&token).Error; err != nil {
			continue
		}
		urls = append(urls, d.Cfg.BaseURL+"/api/v1/files/"+token.Token)
	}
	var accountIDs []uint64
	_ = json.Unmarshal([]byte(row.AccountIDs), &accountIDs)
	var tags []string
	_ = json.Unmarshal([]byte(row.Tags), &tags)
	push := &proto.TaskPush{
		TaskID:       strconv.FormatUint(row.ID, 10),
		Title:        row.Title,
		Content:      row.Content,
		Tags:         tags,
		MaterialURLs: urls,
		Priority:     row.Priority,
	}
	if len(accountIDs) > 0 {
		var acc models.Account
		if err := d.DB.First(&acc, accountIDs[0]).Error; err == nil {
			push.Platform = acc.Platform
			push.AccountID = strconv.FormatUint(acc.ID, 10)
			push.AccountName = acc.AccountName
		}
	}
	if row.ScheduleAt != nil {
		push.ScheduleAt = row.ScheduleAt.UnixMilli()
	}
	return push
}