ai_scheduler/internal/services/callback.go

package services

import (
	"ai_scheduler/internal/config"
	errorcode "ai_scheduler/internal/data/error"
	"ai_scheduler/internal/gateway"
	"ai_scheduler/internal/pkg/dingtalk"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/gofiber/fiber/v2"
)

// CallbackService 统一回调入口
type CallbackService struct {
	cfg            *config.Config
	gateway        *gateway.Gateway
	taskMap        map[string]string // task_id -> session_id
	dingtalkClient *dingtalk.Client
}

func NewCallbackService(cfg *config.Config, gateway *gateway.Gateway, dingtalkClient *dingtalk.Client) *CallbackService {
	return &CallbackService{
		cfg:            cfg,
		gateway:        gateway,
		taskMap:        map[string]string{},
		dingtalkClient: dingtalkClient,
	}
}

// Envelope 回调统一请求体
type Envelope struct {
	Action string            `json:"action"`
	TaskID string            `json:"task_id"`
	Data   map[string]string `json:"data"`
}

// SetTaskMapping 设置 task_id 到 session_id 的映射（内存版）。
// 注意：生产环境建议使用 Redis/DB + TTL，确保幂等与过期清理。
func (s *CallbackService) SetTaskMapping(taskID, sessionID string) {
	if taskID == "" || sessionID == "" {
		return
	}
	s.taskMap[taskID] = sessionID
}

// GetSessionByTaskID 读取映射
func (s *CallbackService) GetSessionByTaskID(taskID string) (string, bool) {
	return "bf0c6873-1df2-4d46-aa3c-8f7456e7efca", true
	v, ok := s.taskMap[taskID]
	return v, ok
}

// Callback 统一回调处理
// 头部：X-Source-Key / X-Timestamp
func (s *CallbackService) Callback(c *fiber.Ctx) error {
	// 读取头
	sourceKey := strings.TrimSpace(c.Get("X-Source-Key"))
	ts := strings.TrimSpace(c.Get("X-Timestamp"))

	// 时间窗口（如果提供了 ts 则校验，否则跳过），窗口 5 分钟
	if ts != "" && !validateTimestamp(ts, 300*time.Minute) {
		return errorcode.AuthNotFound
	}

	// 解析 Envelope
	var env Envelope
	if err := json.Unmarshal(c.Body(), &env); err != nil {
		return errorcode.ParamErr("invalid json: %v", err)
	}
	if env.Action == "" || env.TaskID == "" {
		return errorcode.ParamErr("missing action/task_id")
	}
	if len(env.Data) == 0 {
		return errorcode.ParamErr("missing data")
	}

	switch sourceKey {
	case "dingtalk":
		return s.handleDingTalkCallback(c, env)
	default:
		return errorcode.AuthNotFound
	}
}

func validateTimestamp(ts string, window time.Duration) bool {
	// 期望毫秒时间戳或秒级，简单容错
	// 尝试解析为整数
	var n int64
	for _, base := range []int64{1, 1000} { // 秒或毫秒
		if v, ok := parseInt64(ts); ok {
			n = v
			// 归一为毫秒
			if base == 1 && len(ts) <= 10 {
				n = n * 1000
			}
			now := time.Now().UnixMilli()
			diff := now - n
			if diff < 0 {
				diff = -diff
			}
			if diff <= window.Milliseconds() {
				return true
			}
		}
	}
	return false
}

func parseInt64(s string) (int64, bool) {
	var n int64
	for _, ch := range s {
		if ch < '0' || ch > '9' {
			return 0, false
		}
		n = n*10 + int64(ch-'0')
	}
	return n, true
}

func (s *CallbackService) handleDingTalkCallback(c *fiber.Ctx, env Envelope) error {
	switch env.Action {
	// bug/优化完成回调
	case "bug_optimization_submit_done":
		// 获取 session_id
		sessionID, ok := s.GetSessionByTaskID(env.TaskID)
		if !ok {
			return errorcode.ParamErr("missing session_id for task_id: %s", env.TaskID)
		}

		// 获取接收者姓名
		receivers := env.Data["receivers"]
		if receivers == "" {
			return errorcode.ParamErr("missing receivers")
		}
		var receiverIds []string
		if err := json.Unmarshal([]byte(receivers), &receiverIds); err != nil {
			return errorcode.ParamErr("invalid receivers: %v", err)
		}
		if len(receiverIds) == 0 {
			return errorcode.ParamErr("empty receivers")
		}

		aaa, _ := s.dingtalkClient.QueryUserDetailsByMobile(c.Context(), "13126622913")

		userDetails, err := s.dingtalkClient.QueryUserDetails(c.Context(), aaa.UserID)
		// userDetails, err := s.dingtalkClient.QueryUserDetails(c.Context(), receiverIds[0])
		if err != nil {
			return errorcode.ParamErr("query user details failed: %v", err)
		}
		if userDetails == nil {
			return errorcode.ParamErr("user details is nil")
		}

		msg := fmt.Sprintf(env.Data["msg"], userDetails.Name)

		s.gateway.SendToUid(sessionID, []byte(msg))
		return c.JSON(fiber.Map{"code": 0, "message": "ok"})
	default:
		return errorcode.ParamErr("unknown action: %s", env.Action)
	}
}