bbgo_origin/pkg/notifier/telegramnotifier/telegram.go

package telegramnotifier

import (
	"fmt"
	"strconv"
	"time"

	"github.com/sirupsen/logrus"
	"gopkg.in/tucnak/telebot.v2"

	"github.com/c9s/bbgo/pkg/types"
)

var log = logrus.WithField("service", "telegram")

type Notifier struct {
	bot *telebot.Bot

	// Subscribers stores the Chat objects for broadcasting
	Subscribers map[int64]time.Time `json:"chats"`

	// Owner
	// when owner and owner chat is not nil, notification will send to the owner chat
	Owner     *telebot.User `json:"owner"`
	OwnerChat *telebot.Chat `json:"chat"`

	broadcast bool
}

type Option func(notifier *Notifier)

func UseBroadcast() Option {
	return func(notifier *Notifier) {
		notifier.broadcast = true
	}
}

// New
func New(bot *telebot.Bot, options ...Option) *Notifier {
	notifier := &Notifier{ bot: bot }

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

	return notifier
}

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 {
		switch a := arg.(type) {

		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 texts, pureArgs = filterPlaintextMessages(args)
	var message string

	switch a := obj.(type) {

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

	case types.PlainText:
		message = a.PlainText()

	case types.Stringer:
		message = a.String()

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

	}

	if n.broadcast {
		n.Broadcast(message)
		for _, text := range texts {
			n.Broadcast(text)
		}
	} else if n.OwnerChat != nil {
		n.SendToOwner(message)
		for _, text := range texts {
			n.SendToOwner(text)
		}
	}
}

func (n *Notifier) AddSubscriber(m *telebot.Message) {
	if n.Subscribers == nil {
		n.Subscribers = make(map[int64]time.Time)
	}

	n.Subscribers[m.Chat.ID] = m.Time()
}

func (n *Notifier) SetOwner(owner *telebot.User, chat *telebot.Chat) {
	n.Owner = owner
	n.OwnerChat = chat
}

func (n *Notifier) SendToOwner(message string) {
	if n.OwnerChat == nil {
		return
	}

	if _, err := n.bot.Send(n.OwnerChat, message); err != nil {
		log.WithError(err).Error("telegram send error")
	}
}

func (n *Notifier) Broadcast(message string) {
	if n.Subscribers == nil {
		return
	}

	for chatID := range n.Subscribers {
		chat, err := n.bot.ChatByID(strconv.FormatInt(chatID, 10))
		if err != nil {
			log.WithError(err).Error("can not get chat by ID")
			continue
		}

		if _, err := n.bot.Send(chat, message); err != nil {
			log.WithError(err).Error("failed to send message")
		}
	}
}