web_huakui_server/pub/common/tgbot/tgbot.go

package tgbot

import (
	"fmt"
	"hkdj-web-server/pub/common/log"
	"hkdj-web-server/pub/common/srvdef"
	"time"

	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api"
)

type RSP_TYPE int

const (
	TextType RSP_TYPE = 0
	FileType RSP_TYPE = 1
)

type TgCmdRsp struct {
	Buf     []byte
	BufType RSP_TYPE
}

// TgBot Telegram机器人.
type TgBot struct {
	// 机器人标示;
	token string
	// // 管理员分组;
	// admins []int64
	// // 管理群分组;
	// groups []int64
	// 机器人引擎;
	engine *tgbotapi.BotAPI
	// 命令存储集;
	command map[string]func(string, int64) TgCmdRsp
	help    string
	exitRcv chan interface{}
}

// BotCmd 命令.
type BotCmd struct {
	// 指令;
	Key string
	// 摘要;
	Summary string
	// 处理;
	Handle func(string, int64) TgCmdRsp
}

type BotCmdList []*BotCmd

// NewBotCmd 新的命令.
func NewBotCmd(key, summary string, handle func(string, int64) TgCmdRsp) *BotCmd {
	return &BotCmd{
		Key:     key,
		Summary: summary,
		Handle:  handle,
	}
}

// NewTelegramBot 新的Telegram机器人.
func NewTelegramBot(token string) *TgBot {
	return &TgBot{
		token:   token,
		exitRcv: make(chan interface{}),
	}
}

// Init 机器人初始化.
func (tb *TgBot) Init() *srvdef.SrvErr {
	var err error
	tb.engine, err = tgbotapi.NewBotAPI(tb.token)
	if err != nil {
		log.Warn("(tgbot:Init) tgbotapi.NewBotAPI failed", log.E(err))
		return srvdef.New(srvdef.BusinessLevel, "", "")
	}
	tb.command = make(map[string]func(string, int64) TgCmdRsp)
	tb.command["help"] = tb.menu
	return nil
}

// Release 机器人退出释放
func (tb *TgBot) Release() *srvdef.SrvErr {
	tb.engine.StopReceivingUpdates()
	close(tb.exitRcv)
	return nil
}

func (tb *TgBot) menu(string, int64) TgCmdRsp {
	return TgCmdRsp{Buf: []byte(tb.help), BufType: TextType}
}

// RegisterCommands 注册命令并且自动生成help.
func (tb *TgBot) RegisterCommands(commands []*BotCmd) {
	for _, cmd := range commands {
		tb.command[cmd.Key] = cmd.Handle
		tb.help += fmt.Sprintln("/" + cmd.Key + ": " + cmd.Summary)
	}
}

// Receive 接收消息(阻塞).
func (tb *TgBot) Receive() error {
	log.Info("(tgbot:Receive) start")
	// 设置更新频率;
	cfg := tgbotapi.NewUpdate(0)
	//cfg.Timeout = 1

	// 得到更新管道;
	updates, err := tb.engine.GetUpdatesChan(cfg)
	if err != nil {
		log.Warn("(tgbot:Receive) tb.engine.GetUpdatesChan failed", log.E(err))
		return srvdef.New(srvdef.BusinessLevel, "", "")
	}
	for {
		select {
		case update := <-updates:
			tb.ProcessMessage(&update)
		case <-tb.exitRcv:
			log.Info("(tgbot:Receive) exist")
			break
		default:
		}
		time.Sleep(time.Second * 1)
	}
	return nil
}

func (tb *TgBot) ProcessMessage(update *tgbotapi.Update) {
	// 忽略非消息更新;
	if update.Message == nil {
		log.Info("(tgbot:ProcessMessage) message is nil")
		return
	}
	info := fmt.Sprintf("[Name: %s%s][ID: %d][Text: %s]\n",
		update.Message.From.FirstName, update.Message.From.LastName,
		update.Message.Chat.ID, update.Message.Text)
	log.Info(info)
	if update.Message.IsCommand() {
		message := tgbotapi.NewMessage(update.Message.Chat.ID, "")
		handle, ok := tb.command[update.Message.Command()]
		if !ok {
			message.Text = "unknown command"
		} else {
			param := update.Message.CommandArguments()
			rsp := handle(param, update.Message.Chat.ID)
			switch rsp.BufType {
			case FileType:

			default:
				message.Text = string(rsp.Buf)
			}
			//message.Text = handle(param)
		}
		tb.engine.Send(message)
	}
}

func (tb *TgBot) StopReceive() {
	tb.engine.StopReceivingUpdates()
	close(tb.exitRcv)
}

// Send 发送消息.
func (tb *TgBot) Send(ids []int64, text string) {
	for _, id := range ids {
		message := tgbotapi.NewMessage(id, text)
		tb.engine.Send(message)
	}
}