-
Notifications
You must be signed in to change notification settings - Fork 0
/
bot.go
113 lines (96 loc) · 2.31 KB
/
bot.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
package bot
import (
"fmt"
"log"
"strings"
"github.com/rahagi/pepeg-bot2/internal/helper/common"
"github.com/rahagi/pepeg-bot2/internal/helper/logger"
"github.com/rahagi/pepeg-bot2/internal/irc"
"github.com/rahagi/pepeg-bot2/internal/irc/message"
"github.com/rahagi/pepeg-bot2/markov/generator"
"github.com/rahagi/pepeg-bot2/markov/trainer"
)
const (
MARKOV_MAX_WORDS = 30
MARKOV_DEFAULT_COUNTER = 50
)
type HandlerFunc func(irc.IRCClient, *message.Payload) error
type HandlerMap map[string]HandlerFunc
type Bot interface {
// Init starts the bot by receiving messages from the IRC server
Init()
// Register register a command handler
RegisterHandler(cmd string, h HandlerFunc)
}
type bot struct {
i irc.IRCClient
h HandlerMap
g generator.Generator
t trainer.Trainer
countUntilGenerate int
enableLogging bool
learningOnly bool
}
// NewBot create new bot
func NewBot(ircClient irc.IRCClient, enableLogging bool, g generator.Generator, t trainer.Trainer, learningOnly bool) Bot {
handlerMap := make(HandlerMap)
return &bot{
i: ircClient,
h: handlerMap,
g: g,
t: t,
countUntilGenerate: MARKOV_DEFAULT_COUNTER,
learningOnly: learningOnly,
enableLogging: enableLogging,
}
}
func (b *bot) Init() {
messages := b.i.Receive()
for p := range messages {
b.countUntilGenerate -= 1
b.defaultHandler(p)
cmd, _ := common.PickCommand(p.Message)
if hf, ok := b.h[cmd]; ok {
if err := hf(b.i, p); err != nil {
log.Printf("bot: failed to handle command %s: %v\n", p.Message, err)
}
}
}
}
func (b *bot) RegisterHandler(cmd string, h HandlerFunc) {
b.h[cmd] = h
}
func (b *bot) defaultHandler(p *message.Payload) {
if strings.HasPrefix(p.Message, "PING") {
b.i.Pong()
return
}
if b.countUntilGenerate <= 0 {
if !b.learningOnly {
m := b.generateMarkov(p)
b.i.Chat(m)
}
b.resetCounter()
}
if p.User != "" {
fm := p.String()
go b.t.AddChain(fm)
if b.enableLogging {
b.log(p)
}
}
}
func (b *bot) resetCounter() {
b.countUntilGenerate = MARKOV_DEFAULT_COUNTER
}
func (b *bot) generateMarkov(p *message.Payload) string {
res, err := b.g.Generate(p.String(), MARKOV_MAX_WORDS)
if err != nil {
return ""
}
return res
}
func (b *bot) log(p *message.Payload) {
logPath := fmt.Sprintf("./log/%s.log", b.i.Channel())
logger.Tee(p.String(), logPath)
}