nignoggobot/commands/admin.go

package commands

import (
	"bytes"
	"fmt"
	"image/color"
	"image/png"
	"nignoggobot/metrics"
	"sort"
	"strings"

	"gonum.org/v1/plot"
	"gonum.org/v1/plot/plotter"
	"gonum.org/v1/plot/vg"
	"gonum.org/v1/plot/vg/draw"
	"gonum.org/v1/plot/vg/vgimg"

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

var AdminIDs = []int64{126131628}

func IsAdmin(userID int64) bool {
	for _, id := range AdminIDs {
		if id == userID {
			return true
		}
	}
	return false
}

type StatsCommand struct{}

func (StatsCommand) Name() string { return "stats" }
func (StatsCommand) Help() string { return "Show bot usage stats (admin only)." }
func (StatsCommand) Execute(update tgbotapi.Update, bot *tgbotapi.BotAPI) {
	if !IsAdmin(update.Message.From.ID) {
		return
	}
	stats := metrics.GetStats()
	var sb strings.Builder

	type cmdStat struct {
		Name      string
		AllTime   int
		Last30Sum int
		Last30    [30]int
	}
	var statsList []cmdStat
	for name, stat := range stats.Commands {
		sum := 0
		for _, v := range stat.Last30 {
			sum += v
		}
		if sum > 0 {
			statsList = append(statsList, cmdStat{Name: name, AllTime: stat.AllTime, Last30Sum: sum, Last30: stat.Last30})
		}
	}

	// Sort by last 30 days usage for the graph
	sort.Slice(statsList, func(i, j int) bool {
		return statsList[i].Last30Sum > statsList[j].Last30Sum
	})

	// Prepare data for stacked bar chart
	nCmds := len(statsList)
	nDays := 30
	if nCmds > 0 {
		colors := []color.Color{
			color.RGBA{0x1f, 0x77, 0xb4, 0xff}, // blue
			color.RGBA{0xff, 0x7f, 0x0e, 0xff}, // orange
			color.RGBA{0x2c, 0xa0, 0x2c, 0xff}, // green
			color.RGBA{0xd6, 0x27, 0x28, 0xff}, // red
			color.RGBA{0x94, 0x67, 0xbd, 0xff}, // purple
			color.RGBA{0x8c, 0x56, 0x4b, 0xff}, // brown
			color.RGBA{0xe3, 0x77, 0xc2, 0xff}, // pink
			color.RGBA{0x7f, 0x7f, 0x7f, 0xff}, // gray
			color.RGBA{0xbc, 0xbd, 0x22, 0xff}, // yellow-green
			color.RGBA{0x17, 0xbe, 0xcf, 0xff}, // cyan
		}
		for len(colors) < nCmds {
			colors = append(colors, color.RGBA{uint8(50 * len(colors)), uint8(100 * len(colors)), uint8(150 * len(colors)), 0xff})
		}

		p := plot.New()
		p.Title.Text = "Command Usage (Last 30 Days)"
		p.Y.Label.Text = "Count"
		p.X.Label.Text = "Days Ago"
		p.Legend.Top = true
		p.Legend.Left = false
		p.Legend.XOffs = 0
		p.Legend.YOffs = 0

		// Prepare the stacked values
		stacks := make([][]float64, nCmds)
		for i := range stacks {
			stacks[i] = make([]float64, nDays)
		}
		for cmdIdx, stat := range statsList {
			for day := 0; day < nDays; day++ {
				stacks[cmdIdx][day] = float64(stat.Last30[day]) // 0=today, 29=oldest
			}
		}

		barCharts := make([]*plotter.BarChart, nCmds)
		barWidth := vg.Points(10)
		for cmdIdx := 0; cmdIdx < nCmds; cmdIdx++ {
			values := make(plotter.Values, nDays)
			for day := 0; day < nDays; day++ {
				values[day] = stacks[cmdIdx][day]
			}
			bar, err := plotter.NewBarChart(values, barWidth)
			if err != nil {
				continue
			}
			bar.Color = colors[cmdIdx]
			bar.Offset = 0
			if cmdIdx > 0 {
				bar.StackOn(barCharts[cmdIdx-1])
			}
			p.Add(bar)
			p.Legend.Add(statsList[cmdIdx].Name, bar)
			barCharts[cmdIdx] = bar
		}
		labels := make([]string, nDays)
		for i := 0; i < nDays; i++ {
			labels[i] = fmt.Sprintf("%d", i) // 0=Today, 29=Oldest
		}
		p.NominalX(labels...)

		img := vgimg.New(800, 400)
		dc := draw.New(img)
		p.Draw(dc)
		buf := new(bytes.Buffer)
		png.Encode(buf, img.Image())
		photo := tgbotapi.FileBytes{Name: "stats.png", Bytes: buf.Bytes()}
		photoMsg := tgbotapi.NewPhoto(update.Message.Chat.ID, photo)
		photoMsg.Caption = "Command usage for the last 30 days (0 = today)"
		bot.Send(photoMsg)
	}

	// Text stats as before
	sort.Slice(statsList, func(i, j int) bool {
		return statsList[i].AllTime > statsList[j].AllTime
	})
	sb.WriteString("Command usage (all time):\n")
	for _, s := range statsList {
		sb.WriteString(fmt.Sprintf("/%s: %d\n", s.Name, s.AllTime))
	}

	sort.Slice(statsList, func(i, j int) bool {
		return statsList[i].Last30Sum > statsList[j].Last30Sum
	})
	sb.WriteString("\nCommand usage (last 30 days):\n")
	for _, s := range statsList {
		sb.WriteString(fmt.Sprintf("/%s: %d\n", s.Name, s.Last30Sum))
	}

	groups, privs, blocked := 0, 0, 0
	for _, chat := range stats.Chats {
		if chat.IsGroup {
			groups++
		} else {
			privs++
		}
		if chat.Blocked {
			blocked++
		}
	}
	sb.WriteString(fmt.Sprintf("\nGroups: %d\nPrivate chats: %d\nBlocked/muted: %d\n", groups, privs, blocked))
	msg := tgbotapi.NewMessage(update.Message.Chat.ID, sb.String())
	bot.Send(msg)
}

type BroadcastCommand struct{}

func (BroadcastCommand) Name() string { return "broadcast" }
func (BroadcastCommand) Help() string {
	return "Broadcast a <message> to all chats (admin only)."
}
func (BroadcastCommand) Execute(update tgbotapi.Update, bot *tgbotapi.BotAPI) {
	if !IsAdmin(update.Message.From.ID) {
		return
	}
	args := strings.TrimSpace(update.Message.CommandArguments())
	if args == "" {
		msg := tgbotapi.NewMessage(update.Message.Chat.ID, "Usage: /broadcast <message>")
		bot.Send(msg)
		return
	}
	count := 0
	metrics.BroadcastMessage(func(chatID int64) error {
		msg := tgbotapi.NewMessage(chatID, args)
		_, err := bot.Send(msg)
		if err == nil {
			count++
		}
		return err
	})
	msg := tgbotapi.NewMessage(update.Message.Chat.ID, fmt.Sprintf("Broadcast sent to %d chats.", count))
	bot.Send(msg)
}

func init() {
	Register(StatsCommand{})
	Register(BroadcastCommand{})
}