velox-bot/internal/tldr/tldrbench/load.go

// Package tldrbench is the toolkit behind the local /tldr benchmark: it loads
// a chat dump from disk, scores a model's overview and renders a dashboard.
// None of it is used by the bot itself, and nothing here talks to the network,
// so a private chat dump never leaves the machine running the benchmark.
package tldrbench

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/bwmarrin/discordgo"
)

// LoadChat reads a chat dump and returns its messages oldest first. Three
// formats are recognised by content:
//
//   - a DiscordChatExporter JSON export (a file whose root has "messages"),
//   - text copied out of the Discord client ("Name — Today at 14:32" header
//     lines followed by the message lines),
//   - plain "Author: message" lines, where a line that does not start with an
//     author is the continuation of the previous message.
func LoadChat(path string) ([]*discordgo.Message, error) {
	raw, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var msgs []*discordgo.Message
	if strings.EqualFold(filepath.Ext(path), ".json") || strings.HasPrefix(strings.TrimSpace(string(raw)), "{") {
		msgs, err = parseExport(raw)
	} else {
		msgs = parseText(string(raw))
	}
	if err != nil {
		return nil, fmt.Errorf("%s: %w", path, err)
	}
	if len(msgs) == 0 {
		return nil, fmt.Errorf("%s: no messages found", path)
	}
	return msgs, nil
}

type exportFile struct {
	Messages []struct {
		ID        string `json:"id"`
		Type      string `json:"type"`
		Timestamp string `json:"timestamp"`
		Content   string `json:"content"`
		Author    struct {
			ID       string `json:"id"`
			Name     string `json:"name"`
			Nickname string `json:"nickname"`
			IsBot    bool   `json:"isBot"`
		} `json:"author"`
		Attachments []struct {
			FileName string `json:"fileName"`
		} `json:"attachments"`
		Embeds []struct {
			Title string `json:"title"`
		} `json:"embeds"`
		Reference *struct {
			MessageID string `json:"messageId"`
		} `json:"reference"`
	} `json:"messages"`
}

func parseExport(raw []byte) ([]*discordgo.Message, error) {
	var f exportFile
	if err := json.Unmarshal(raw, &f); err != nil {
		return nil, fmt.Errorf("not a chat export: %w", err)
	}
	byID := make(map[string]*discordgo.Message, len(f.Messages))
	msgs := make([]*discordgo.Message, 0, len(f.Messages))
	for i, e := range f.Messages {
		name := e.Author.Nickname
		if name == "" {
			name = e.Author.Name
		}
		id := e.ID
		if id == "" {
			id = strconv.Itoa(i + 1)
		}
		m := &discordgo.Message{
			ID:      id,
			Content: e.Content,
			Author:  &discordgo.User{ID: firstNonEmpty(e.Author.ID, name), Username: name, Bot: e.Author.IsBot},
		}
		if ts, err := time.Parse(time.RFC3339Nano, e.Timestamp); err == nil {
			m.Timestamp = ts
		}
		switch e.Type {
		case "", "Default":
		case "Reply":
			m.Type = discordgo.MessageTypeReply
			if e.Reference != nil {
				m.ReferencedMessage = byID[e.Reference.MessageID]
			}
		default:
			m.Type = discordgo.MessageTypeGuildMemberJoin // anything that is not a human message
		}
		for _, a := range e.Attachments {
			m.Attachments = append(m.Attachments, &discordgo.MessageAttachment{Filename: a.FileName})
		}
		for _, em := range e.Embeds {
			m.Embeds = append(m.Embeds, &discordgo.MessageEmbed{Type: discordgo.EmbedTypeArticle, Title: em.Title})
		}
		byID[id] = m
		msgs = append(msgs, m)
	}
	return msgs, nil
}

func firstNonEmpty(a, b string) string {
	if a != "" {
		return a
	}
	return b
}

var (
	// "Name — Today at 14:32", "Name — 12/10/2026 14:32", "Name — 14:32".
	clientHeader = regexp.MustCompile(`^(.{1,40}?) [—–-] (?:Today at |Yesterday at |hoje às |ontem às )?(?:\d{1,4}[/.\-]\d{1,2}[/.\-]\d{1,4},? ?)?(?:at |às )?\d{1,2}:\d{2}(?:\s?[AP]M)?\s*$`)
	// Lone timestamps the client shows next to grouped messages: "[14:33]" or "14:33".
	loneTime   = regexp.MustCompile(`^\[?\d{1,2}:\d{2}(?:\s?[AP]M)?\]?$`)
	authorLine = regexp.MustCompile(`^([^:\n]{1,32}): (.+)$`)
)

func parseText(s string) []*discordgo.Message {
	lines := strings.Split(strings.ReplaceAll(s, "\r\n", "\n"), "\n")

	headers := 0
	for _, ln := range lines {
		if clientHeader.MatchString(strings.TrimSpace(ln)) {
			headers++
		}
	}
	if headers >= 2 {
		return parseClientCopy(lines)
	}
	return parseAuthorLines(lines)
}

type builder struct{ msgs []*discordgo.Message }

func (b *builder) add(author, text string) {
	text = strings.TrimSpace(text)
	if text == "" {
		return
	}
	id := strconv.Itoa(len(b.msgs) + 1)
	b.msgs = append(b.msgs, &discordgo.Message{
		ID:      id,
		Content: text,
		Author:  &discordgo.User{ID: author, Username: author},
	})
}

func parseClientCopy(lines []string) []*discordgo.Message {
	var b builder
	author := ""
	for _, ln := range lines {
		t := strings.TrimSpace(ln)
		switch {
		case t == "" || loneTime.MatchString(t):
		case clientHeader.MatchString(t):
			author = clientHeader.FindStringSubmatch(t)[1]
		case author != "":
			b.add(author, t)
		}
	}
	return b.msgs
}

func parseAuthorLines(lines []string) []*discordgo.Message {
	var b builder
	var author string
	var cur []string
	flush := func() {
		if author != "" {
			b.add(author, strings.Join(cur, "\n"))
		}
		author, cur = "", nil
	}
	for _, ln := range lines {
		t := strings.TrimSpace(ln)
		if t == "" {
			continue
		}
		if m := authorLine.FindStringSubmatch(t); m != nil {
			flush()
			author, cur = strings.TrimSpace(m[1]), []string{m[2]}
			continue
		}
		if author != "" {
			cur = append(cur, t)
		}
	}
	flush()
	return b.msgs
}