bot.go

package bot

import (
	"context"
	"fmt"
	"log/slog"
	"os"
	"strings"
	"time"

	"github.com/ai-shift/tgmd"
	"github.com/firebase/genkit/go/genkit"
	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5"
	"github.com/jackc/pgx/v5/pgxpool"

	"mimi/internal/bot/llm"
	"mimi/internal/bot/llm/agent"
	"mimi/internal/persist"
	"mimi/internal/provider/logseq"
)

func Start(ctx context.Context, token string, logseqPath string, g *genkit.Genkit) error {
	slog.Info("starting Telegram Bot")
	bot, err := tgbotapi.NewBotAPI(token)
	if err != nil {
		return fmt.Errorf("failed to initialize bot api with %w", err)
	}
	slog.Info("Authorized account", "username", bot.Self.UserName)

	pool, err := pgxpool.New(ctx, os.Getenv("DATABASE_URL"))
	if err != nil {
		return fmt.Errorf("failed to connect to postgres with %w", err)
	}

	graph := logseq.NewRegexGraph(logseqPath)
	handler := UpdateHandler{
		bot:     bot,
		g:       graph,
		llm:     llm.New(ctx, pool, graph, g),
		queries: persist.New(pool),
	}

	// Register bot commands
	commands := []tgbotapi.BotCommand{
		{
			Command:     "clear",
			Description: "Clear all message history",
		},
	}
	cmdConfig := tgbotapi.NewSetMyCommands(commands...)
	if _, err := bot.Request(cmdConfig); err != nil {
		slog.Warn("failed to set bot commands", "error", err)
	}

	u := tgbotapi.NewUpdate(0)
	u.Timeout = 60

	updates := bot.GetUpdatesChan(u)
	slog.Info("Telegram bot started")
	for {
		select {
		case <-ctx.Done():
			return nil
		case update := <-updates:
			if update.Message != nil && len(update.Message.Text) > 0 {
				go func() {
					slog.Info("got new message in Telegram bot")
					if err := handler.handleMessage(ctx, update.Message); err != nil {
						// Log comprehensive error details for debugging
						slog.Error("failed to handle message",
							"error", err,
							"chat_id", update.Message.Chat.ID,
							"chat_type", update.Message.Chat.Type,
							"chat_title", update.Message.Chat.Title,
							"chat_username", update.Message.Chat.UserName,
							"user_id", update.Message.From.ID,
							"user_username", update.Message.From.UserName,
							"user_first_name", update.Message.From.FirstName,
							"user_last_name", update.Message.From.LastName,
							"message_id", update.Message.MessageID,
							"message_text", update.Message.Text,
							"message_date", update.Message.Date,
						)

						// Send error message without parse mode to avoid formatting issues
						errorMsg := fmt.Sprintf("Error: %v", err)
						msg := tgbotapi.NewMessage(update.Message.Chat.ID, errorMsg)
						msg.ParseMode = "" // Send as plain text
						_, sendErr := bot.Send(msg)
						if sendErr != nil {
							slog.Error("failed to send error message to user",
								"original_error", err,
								"send_error", sendErr,
								"chat_id", update.Message.Chat.ID,
							)
						}
					}
				}()
			}
		}
	}
}

type UpdateHandler struct {
	bot     *tgbotapi.BotAPI
	g       logseq.RegexGraph
	llm     llm.LLM
	queries *persist.Queries
}

func (h UpdateHandler) handleMessage(ctx context.Context, m *tgbotapi.Message) error {
	slog.Info("new message", "chatId", m.Chat.ID, "text", m.Text)

	// Handle /clear command
	if m.Text == "/clear" {
		if err := h.queries.DeleteChatMessages(ctx, m.Chat.ID); err != nil {
			return fmt.Errorf("failed to delete chat messages with %w", err)
		}
		msg := tgbotapi.NewMessage(m.Chat.ID, "Chat history has been cleared.")
		if _, err := h.bot.Send(msg); err != nil {
			return fmt.Errorf("failed to send confirmation message with %w", err)
		}
		return nil
	}

	// Set bot typing status
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()
	go func() {
		for {
			_, _ = h.bot.Send(tgbotapi.NewChatAction(m.Chat.ID, tgbotapi.ChatTyping))
			select {
			case <-ctx.Done():
				ticker.Stop()
				return
			case <-ticker.C:
				continue
			}
		}
	}()

	// Generate LLM answer
	result, err := h.llm.Answer(ctx, m.Chat.ID, m.Text)
	if err != nil {
		return fmt.Errorf("failed to get answer from LLM with %w", err)
	}

	switch data := result.Data.(type) {
	case agent.DataText:
		// Response to the user's query
		slog.Info("got LLM text answer", "length", len(data.Text))
		if err := sendLongMessage(h.bot, m.Chat.ID, data.Text); err != nil {
			return fmt.Errorf("failed to send LLM response with %w", err)
		}
	case agent.DataFile:
		slog.Info("got LLM file answer", "size", len(data.Blob))
		// TODO: Send as file
		tmpfile, err := os.CreateTemp("", fmt.Sprintf("*-%s", data.Name))
		if err != nil {
			return fmt.Errorf("failed to create temporary file with %w", err)
		}
		defer os.Remove(tmpfile.Name())

		if _, err := tmpfile.Write(data.Blob); err != nil {
			return fmt.Errorf("failed to write to temporary file with %w", err)
		}
		if err := tmpfile.Close(); err != nil {
			return fmt.Errorf("failed to close temporary file with %w", err)
		}

		tmpfile, err = os.Open(tmpfile.Name())
		if err != nil {
			return fmt.Errorf("failed to open temporary file after write with %w", err)
		}

		f := tgbotapi.FileReader{
			Name:   tmpfile.Name(),
			Reader: tmpfile,
		}
		req := tgbotapi.NewDocument(m.Chat.ID, f)
		_, err = h.bot.Send(req)
		if err != nil {
			return fmt.Errorf("failed to send document with %w", err)
		}
	default:
		return fmt.Errorf("unexpected answer type '%#v'", data)
	}

	return nil
}

// sendLongMessage splits text into chunks and may send several messages
// to prevent error of exceeding Telegram's limit
func sendLongMessage(bot *tgbotapi.BotAPI, chatID int64, text string) error {
	var buf []string
	var curLen int
	for _, line := range strings.Split(text, "\n") {
		if curLen+len(line) <= 3000 {
			// Under the limit, continue accumulating
			buf = append(buf, line)
			curLen += len(line)
			continue
		}
		// The time to send is come
		if err := sendShortMessage(bot, chatID, strings.Join(buf, "\n")); err != nil {
			slog.Warn("failed to send message with MarkdownV2, falling back to plain text",
				"error", err,
				"error_type", fmt.Sprintf("%T", err),
				"chat_id", chatID,
				"original_text", text,
				"text_length", len(text),
			)
			if err := sendPlainMessage(bot, chatID, text); err != nil {
				return err
			}
		}
		// Clean up state
		buf = buf[:0]
		curLen = 0
	}
	if len(buf) == 0 {
		return nil
	}
	return sendShortMessage(bot, chatID, strings.Join(buf, "\n"))
}

func sendShortMessage(bot *tgbotapi.BotAPI, chatID int64, text string) error {
	formattedText := tgmd.Telegramify(text)
	msg := tgbotapi.NewMessage(chatID, formattedText)
	msg.ParseMode = "MarkdownV2"
	_, err := bot.Send(msg)
	if err != nil {
		return err
	}
	return nil
}

func sendPlainMessage(bot *tgbotapi.BotAPI, chatID int64, text string) error {
	plainMsg := tgbotapi.NewMessage(chatID, text)
	plainMsg.ParseMode = ""
	_, err := bot.Send(plainMsg)
	return err
}