laniakea/bot.go

package laniakea

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
)

type ParseMode string

const (
	ParseMDV2 ParseMode = "MarkdownV2"
	ParseHTML ParseMode = "HTML"
	ParseMD   ParseMode = "Markdown"
)

type Bot struct {
	token  string
	logger *Logger
	debug  bool

	plugins  []*Plugin
	prefixes []string

	updateOffset int
	updateTypes  []string
	updateQueue  *Queue[*Update]
}

type MsgContext struct {
	Bot    *Bot
	Msg    *Message
	Update *Update
	FromID int
	Prefix string
	Text   string
}

func NewBot(token string) *Bot {
	logger := CreateLogger()
	logger.Level(DEBUG)
	updateQueue := CreateQueue[*Update](256)
	bot := &Bot{
		updateOffset: 0, plugins: make([]*Plugin, 0), debug: false,
		prefixes: make([]string, 0), updateTypes: make([]string, 0),
		updateQueue: updateQueue,
		token:       token, logger: logger,
	}

	return bot
}

func (b *Bot) UpdateTypes(t ...string) *Bot {
	b.updateTypes = make([]string, 0)
	b.updateTypes = append(b.updateTypes, t...)
	return b
}
func (b *Bot) AddUpdateType(t ...string) *Bot {
	b.updateTypes = append(b.updateTypes, t...)
	return b
}

func (b *Bot) AddPrefixes(prefixes ...string) *Bot {
	b.prefixes = append(b.prefixes, prefixes...)
	return b
}

func (b *Bot) Run() {
	if len(b.prefixes) == 0 {
		b.logger.Fatal("no prefixes defined")
		return
	}

	if len(b.plugins) == 0 {
		b.logger.Fatal("no plugins defined")
		return
	}

	b.logger.Info("Bot running. Press CTRL+C to exit.")

	go func() {
		for {
			_, err := b.Updates()
			if err != nil {
				b.logger.Error(err)
			}
		}
	}()

	for {
		queue := b.updateQueue
		if queue.IsEmpty() {
			continue
		}

		u := queue.Dequeue()
		if u.CallbackQuery != nil {
			b.handleCallback(u)
		} else {
			b.handleMessage(u)
		}
	}
}

// {"callback_query":{"chat_instance":"6202057960757700762","data":"aboba","from":{"first_name":"scuroneko","id":314834933,"is_bot":false,"language_code":"ru","username":"scuroneko"},"id":"1352205741990111553","message":{"chat":{"first_name":"scuroneko","id":314834933,"type":"private","username":"scuroneko"},"date":1734338107,"from":{"first_name":"Kurumi","id":7718900880,"is_bot":true,"username":"kurumi_game_bot"},"message_id":19,"reply_markup":{"inline_keyboard":[[{"callback_data":"aboba","text":"Test"},{"callback_data":"another","text":"Another"}]]},"text":"Aboba"}},"update_id":350979488}

func (b *Bot) handleMessage(update *Update) {
	ctx := &MsgContext{
		Bot:    b,
		Update: update,
	}

	for _, plugin := range b.plugins {
		if plugin.UpdateListener != nil {
			(*plugin.UpdateListener)(ctx)
		}
	}

	if update.Message == nil {
		return
	}

	ctx.FromID = update.Message.From.ID
	ctx.Msg = update.Message
	text := strings.TrimSpace(update.Message.Text)
	prefix, hasPrefix := b.checkPrefixes(text)
	if !hasPrefix {
		return
	}
	ctx.Prefix = prefix

	text = strings.TrimSpace(text[len(prefix):])

	for _, plugin := range b.plugins {

		// Check every command
		for cmd := range plugin.Commands {
			if !strings.HasPrefix(text, cmd) {
				continue
			}

			ctx.Text = text[len(cmd):]

			plugin.Execute(cmd, ctx)
		}
	}
}

func (b *Bot) handleCallback(update *Update) {
	ctx := &MsgContext{
		Bot:    b,
		Update: update,
	}

	for _, plugin := range b.plugins {
		if plugin.UpdateListener != nil {
			(*plugin.UpdateListener)(ctx)
		}
	}

	for _, plugin := range b.plugins {
		for payload := range plugin.Payloads {
			if !strings.HasPrefix(update.CallbackQuery.Data, payload) {
				continue
			}
			plugin.ExecutePayload(payload, ctx)
		}
	}
}

func (b *Bot) checkPrefixes(text string) (string, bool) {
	for _, prefix := range b.prefixes {
		if strings.HasPrefix(text, prefix) {
			return prefix, true
		}
	}
	return "", false
}

func (b *Bot) AddPlugins(plugin ...*Plugin) *Bot {
	b.plugins = append(b.plugins, plugin...)
	return b
}

func (b *Bot) Debug(debug bool) *Bot {
	b.debug = debug
	return b
}

func (ctx *MsgContext) Answer(text string) {
	_, err := ctx.Bot.SendMessage(&SendMessageP{
		ChatID:    ctx.Msg.Chat.ID,
		Text:      text,
		ParseMode: "markdown",
	})
	if err != nil {
		ctx.Bot.logger.Error(err)
	}
}

func (ctx *MsgContext) Error(err error) {
	_, sendErr := ctx.Bot.SendMessage(&SendMessageP{
		ChatID: ctx.Msg.Chat.ID,
		Text:   err.Error(),
	})

	if sendErr != nil {
		ctx.Bot.logger.Error(sendErr)
	}
}

func (b *Bot) Logger() *Logger {
	return b.logger
}

type ApiResponse struct {
	Ok          bool                   `json:"ok"`
	Result      map[string]interface{} `json:"result,omitempty"`
	Description string                 `json:"description,omitempty"`
	ErrorCode   int                    `json:"error_code,omitempty"`
}

type ApiResponseA struct {
	Ok          bool          `json:"ok"`
	Result      []interface{} `json:"result,omitempty"`
	Description string        `json:"description,omitempty"`
	ErrorCode   int           `json:"error_code,omitempty"`
}

// request is a low-level call to api.
func (b *Bot) request(methodName string, params map[string]interface{}) (map[string]interface{}, error) {
	var buf bytes.Buffer
	err := json.NewEncoder(&buf).Encode(params)
	if err != nil {
		return nil, err
	}

	if b.debug {
		b.logger.Debug(fmt.Sprintf(
			"POST https://api.telegram.org/bot%s/%s %s",
			"<TOKEN>",
			methodName,
			string(buf.Bytes()),
		))
	}
	r, err := http.Post(fmt.Sprintf("https://api.telegram.org/bot%s/%s", b.token, methodName), "application/json", &buf)
	if err != nil {
		return nil, err
	}
	data, err := io.ReadAll(r.Body)
	if err != nil {
		return nil, err
	}
	response := new(ApiResponse)

	var result map[string]interface{}

	err = json.Unmarshal(data, &response)
	if err != nil {
		responseArray := new(ApiResponseA)
		err = json.Unmarshal(data, responseArray)
		if err != nil {
			return nil, err
		}
		result = map[string]interface{}{
			"data": responseArray.Result,
		}
	} else {
		result = response.Result
	}
	if !response.Ok {
		return nil, fmt.Errorf("[%d] %s", response.ErrorCode, response.Description)
	}
	return result, err
}