qbtrade/pkg/notifier/larknotifier/lark.go

package larknotifier

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"git.qtrade.icu/lychiyu/qbtrade/pkg/types"
	"github.com/sirupsen/logrus"
	"golang.org/x/time/rate"
	"net/http"
	"reflect"
	"time"
)

var apiLimiter = rate.NewLimiter(rate.Every(time.Second), 5)
var log = logrus.WithField("service", "lark")

type notifyTask struct {
	message string
}

type Notifier struct {
	token string
	taskC chan notifyTask
}

type Option func(notifier *Notifier)

type TextMessage struct {
	MsgType string `json:"msg_type"`
	Content struct {
		Text string `json:"text"`
	} `json:"content"`
}

func getLocDateTimeString() string {
	loc, err := time.LoadLocation("Asia/Shanghai")
	if err != nil {
		fmt.Println("Error loading location:", err)
		return ""
	}

	// 获取当前东八区时间
	currentTime := time.Now().In(loc)

	// 格式化时间
	formattedTime := currentTime.Format("2006-01-02 15:04:05")
	return formattedTime
}

func sendLarkMessage(token, text string) error {
	// 创建消息结构体
	message := TextMessage{
		MsgType: "text",
		Content: struct {
			Text string `json:"text"`
		}{
			Text: fmt.Sprintf("通知时间：%s\n内容：\n%s", getLocDateTimeString(), text),
		},
	}

	// 将消息结构体序列化为 JSON
	jsonData, err := json.Marshal(message)
	if err != nil {
		return fmt.Errorf("failed to marshal JSON: %v", err)
	}
	webhookURL := fmt.Sprintf("https://open.larksuite.com/open-apis/bot/v2/hook/%s", token)
	// 创建 HTTP POST 请求
	req, err := http.NewRequest("POST", webhookURL, bytes.NewBuffer(jsonData))
	if err != nil {
		return fmt.Errorf("failed to create request: %v", err)
	}
	req.Header.Set("Content-Type", "application/json")

	// 发送 HTTP 请求
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return fmt.Errorf("failed to send request: %v", err)
	}
	defer resp.Body.Close()

	// 检查响应状态码
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("request failed with status code: %d", resp.StatusCode)
	}

	return nil
}

func New(token string, options ...Option) *Notifier {
	notifier := &Notifier{
		token: token,
		taskC: make(chan notifyTask, 100),
	}

	for _, o := range options {
		o(notifier)
	}

	go notifier.worker()

	return notifier
}

func (n *Notifier) worker() {
	ctx := context.Background()
	for {
		select {
		case <-ctx.Done():
			return
		case task := <-n.taskC:
			apiLimiter.Wait(ctx)
			n.consume(task)
		}
	}
}

func (n *Notifier) consume(task notifyTask) {
	if task.message != "" {
		if err := sendLarkMessage(n.token, task.message); err != nil {
			log.WithError(err).Error("lark send error")
		}
	}
}

func (n *Notifier) Notify(obj interface{}, args ...interface{}) {
	n.NotifyTo("", obj, args...)
}

func filterPlaintextMessages(args []interface{}) (texts []string, pureArgs []interface{}) {
	var firstObjectOffset = -1
	for idx, arg := range args {
		rt := reflect.TypeOf(arg)
		if rt.Kind() == reflect.Ptr {
			switch a := arg.(type) {

			case nil:
				texts = append(texts, "nil")
				if firstObjectOffset == -1 {
					firstObjectOffset = idx
				}

			case types.PlainText:
				texts = append(texts, a.PlainText())
				if firstObjectOffset == -1 {
					firstObjectOffset = idx
				}

			case types.Stringer:
				texts = append(texts, a.String())
				if firstObjectOffset == -1 {
					firstObjectOffset = idx
				}
			}
		}
	}

	pureArgs = args
	if firstObjectOffset > -1 {
		pureArgs = args[:firstObjectOffset]
	}

	return texts, pureArgs
}

func (n *Notifier) NotifyTo(channel string, obj interface{}, args ...interface{}) {
	var message string

	switch a := obj.(type) {

	case string:
		message = fmt.Sprintf(a, args...)

	default:
		log.Errorf("unsupported notification format: %T %+v", a, a)
	}

	select {
	case n.taskC <- notifyTask{
		message: message,
	}:
	default:
		log.Error("[lark] cannot send task to notify")
	}
}

func (n *Notifier) SendPhoto(buffer *bytes.Buffer) {
	n.SendPhotoTo("", buffer)
}

func (n *Notifier) SendPhotoTo(channel string, buffer *bytes.Buffer) {
	// TODO
}