From 2ed9c2f6dc77854308806ed154f0de4649bf19f7 Mon Sep 17 00:00:00 2001
From: Alexey Skobkin <skobkin-ru@ya.ru>
Date: Mon, 11 Mar 2024 23:15:27 +0300
Subject: [PATCH] Implementing simple stats counter.

---
 bot/bot.go        | 32 ++++++++++++++++--
 bot/middleware.go | 24 ++++++++++++++
 stats/stats.go    | 84 +++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 138 insertions(+), 2 deletions(-)
 create mode 100644 bot/middleware.go
 create mode 100644 stats/stats.go

diff --git a/bot/bot.go b/bot/bot.go
index e75dc17..9e1f17d 100644
--- a/bot/bot.go
+++ b/bot/bot.go
@@ -10,6 +10,7 @@ import (
 	"strings"
 	"telegram-ollama-reply-bot/extractor"
 	"telegram-ollama-reply-bot/llm"
+	"telegram-ollama-reply-bot/stats"
 )
 
 var (
@@ -18,12 +19,11 @@ var (
 	ErrHandlerInit    = errors.New("cannot initialize handler")
 )
 
-const contextUserKey = "user"
-
 type Bot struct {
 	api       *telego.Bot
 	llm       *llm.LlmConnector
 	extractor *extractor.Extractor
+	stats     *stats.Stats
 }
 
 func NewBot(api *telego.Bot, llm *llm.LlmConnector, extractor *extractor.Extractor) *Bot {
@@ -31,6 +31,7 @@ func NewBot(api *telego.Bot, llm *llm.LlmConnector, extractor *extractor.Extract
 		api:       api,
 		llm:       llm,
 		extractor: extractor,
+		stats:     stats.NewStats(),
 	}
 }
 
@@ -66,6 +67,10 @@ func (b *Bot) Run() error {
 	defer bh.Stop()
 	defer b.api.StopLongPolling()
 
+	// Middlewares
+	bh.Use(b.chatTypeStatsCounter)
+
+	// Handlers
 	bh.Handle(b.startHandler, th.CommandEqual("start"))
 	bh.Handle(b.heyHandler, th.CommandEqual("hey"))
 	bh.Handle(b.summarizeHandler, th.CommandEqual("summarize"))
@@ -79,6 +84,8 @@ func (b *Bot) Run() error {
 func (b *Bot) heyHandler(bot *telego.Bot, update telego.Update) {
 	slog.Info("/hey")
 
+	b.stats.HeyRequest()
+
 	chatID := tu.ID(update.Message.Chat.ID)
 
 	b.sendTyping(chatID)
@@ -112,6 +119,8 @@ func (b *Bot) heyHandler(bot *telego.Bot, update telego.Update) {
 func (b *Bot) summarizeHandler(bot *telego.Bot, update telego.Update) {
 	slog.Info("/summarize", update.Message.Text)
 
+	b.stats.SummarizeRequest()
+
 	chatID := tu.ID(update.Message.Chat.ID)
 
 	b.sendTyping(chatID)
@@ -208,6 +217,25 @@ func (b *Bot) startHandler(bot *telego.Bot, update telego.Update) {
 	}
 }
 
+func (b *Bot) statsHandler(bot *telego.Bot, update telego.Update) {
+	slog.Info("/stats")
+
+	chatID := tu.ID(update.Message.Chat.ID)
+
+	b.sendTyping(chatID)
+
+	_, err := bot.SendMessage(b.reply(update.Message, tu.Message(
+		chatID,
+		"Current bot stats:\r\n"+
+			"```json\r\n"+
+			b.stats.String()+"\r\n"+
+			"```",
+	)).WithParseMode("Markdown"))
+	if err != nil {
+		slog.Error("Cannot send a message", err)
+	}
+}
+
 func (b *Bot) reply(originalMessage *telego.Message, newMessage *telego.SendMessageParams) *telego.SendMessageParams {
 	return newMessage.WithReplyParameters(&telego.ReplyParameters{
 		MessageID: originalMessage.MessageID,
diff --git a/bot/middleware.go b/bot/middleware.go
new file mode 100644
index 0000000..e3a7467
--- /dev/null
+++ b/bot/middleware.go
@@ -0,0 +1,24 @@
+package bot
+
+import (
+	"github.com/mymmrac/telego"
+	"github.com/mymmrac/telego/telegohandler"
+)
+
+func (b *Bot) chatTypeStatsCounter(bot *telego.Bot, update telego.Update, next telegohandler.Handler) {
+	message := update.Message
+
+	if message == nil {
+		next(bot, update)
+	}
+
+	switch message.Chat.Type {
+	case telego.ChatTypeGroup, telego.ChatTypeSupergroup:
+		b.stats.GroupRequest()
+	case telego.ChatTypePrivate:
+		b.stats.PrivateRequest()
+		b.stats.PrivateRequest()
+	}
+
+	next(bot, update)
+}
diff --git a/stats/stats.go b/stats/stats.go
new file mode 100644
index 0000000..44f8eda
--- /dev/null
+++ b/stats/stats.go
@@ -0,0 +1,84 @@
+package stats
+
+import (
+	"encoding/json"
+	"sync"
+	"time"
+)
+
+type Stats struct {
+	mu sync.Mutex
+
+	RunningSince time.Time
+
+	GroupRequests   uint64
+	PrivateRequests uint64
+
+	HeyRequests       uint64
+	SummarizeRequests uint64
+}
+
+func NewStats() *Stats {
+	return &Stats{
+		RunningSince: time.Now(),
+
+		GroupRequests:   0,
+		PrivateRequests: 0,
+
+		HeyRequests:       0,
+		SummarizeRequests: 0,
+	}
+}
+
+func (s *Stats) MarshalJSON() ([]byte, error) {
+	return json.Marshal(struct {
+		Uptime string `json:"uptime"`
+
+		GroupRequests   uint64 `json:"group_requests"`
+		PrivateRequests uint64 `json:"private_requests"`
+
+		HeyRequests       uint64 `json:"hey_requests"`
+		SummarizeRequests uint64 `json:"summarize_requests"`
+	}{
+		Uptime: time.Now().Sub(s.RunningSince).String(),
+
+		GroupRequests:   s.GroupRequests,
+		PrivateRequests: s.PrivateRequests,
+
+		HeyRequests:       s.HeyRequests,
+		SummarizeRequests: s.SummarizeRequests,
+	})
+}
+
+func (s *Stats) String() string {
+	data, err := json.MarshalIndent(s, "", "  ")
+	if err != nil {
+		return "{\"error\": \"cannot serialize stats\"}"
+	}
+
+	return string(data)
+}
+
+func (s *Stats) GroupRequest() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.GroupRequests++
+}
+
+func (s *Stats) PrivateRequest() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.PrivateRequests++
+}
+
+func (s *Stats) HeyRequest() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.HeyRequests++
+}
+
+func (s *Stats) SummarizeRequest() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.SummarizeRequests++
+}