forked from Ressetkk/dgwidgets
-
Notifications
You must be signed in to change notification settings - Fork 0
/
widget.go
254 lines (225 loc) · 6.31 KB
/
widget.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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
package dgwidgets
import (
"errors"
"sync"
"time"
"github.com/bwmarrin/discordgo"
)
// error vars
var (
ErrAlreadyRunning = errors.New("err: Widget already running")
ErrIndexOutOfBounds = errors.New("err: Index is out of bounds")
ErrNilMessage = errors.New("err: Message is nil")
ErrNilEmbed = errors.New("err: embed is nil")
ErrNotRunning = errors.New("err: not running")
ErrTickerNotSet = errors.New("err: Timeout ticker is not set")
)
// WidgetHandler ...
type WidgetHandler func(*Widget, *discordgo.MessageReaction)
// Widget is a message embed with reactions for buttons.
// Accepts custom handlers for reactions.
type Widget struct {
sync.Mutex
Embed *discordgo.MessageEmbed
Message *discordgo.Message
Ses *discordgo.Session
ChannelID string
Timeout time.Duration
Close chan bool
// Handlers binds emoji names to functions
Handlers map[string]WidgetHandler
// keys stores the handlers keys in the order they were added
Keys []string
// Delete reactions after they are added
DeleteReactions bool
// Refresh timer after action on a widget
RefreshAfterAction bool
// Only allow listed users to use reactions.
UserWhitelist []string
running bool
ticker *time.Ticker
}
// NewWidget returns a pointer to a Widget object
// ses : discordgo session
// channelID: channelID to spawn the widget on
func NewWidget(ses *discordgo.Session, channelID string, embed *discordgo.MessageEmbed) *Widget {
return &Widget{
ChannelID: channelID,
Ses: ses,
Keys: []string{},
Handlers: map[string]WidgetHandler{},
Close: make(chan bool),
DeleteReactions: true,
Embed: embed,
}
}
// isUserAllowed returns true if the user is allowed
// to use this widget.
func (w *Widget) isUserAllowed(userID string) bool {
if w.UserWhitelist == nil || len(w.UserWhitelist) == 0 {
return true
}
for _, user := range w.UserWhitelist {
if user == userID {
return true
}
}
return false
}
// Spawn spawns the widget in channel w.ChannelID
func (w *Widget) Spawn() error {
if w.Running() {
return ErrAlreadyRunning
}
w.running = true
defer func() {
w.running = false
}()
if w.Embed == nil {
return ErrNilEmbed
}
if w.Timeout != 0 {
w.ticker = time.NewTicker(w.Timeout)
defer w.ticker.Stop()
}
// Create initial message.
msg, err := w.Ses.ChannelMessageSendEmbed(w.ChannelID, w.Embed)
if err != nil {
return err
}
w.Message = msg
// Add reaction buttons
for _, v := range w.Keys {
w.Ses.MessageReactionAdd(w.Message.ChannelID, w.Message.ID, v)
}
var reaction *discordgo.MessageReaction
for {
// Navigation timeout enabled
if w.Timeout != 0 {
select {
case k := <-nextMessageReactionAddC(w.Ses):
reaction = k.MessageReaction
case <-w.ticker.C:
return nil
case <-w.Close:
return nil
}
} else /*Navigation timeout not enabled*/ {
select {
case k := <-nextMessageReactionAddC(w.Ses):
reaction = k.MessageReaction
case <-w.Close:
return nil
}
}
// Ignore reactions sent by bot
if reaction.MessageID != w.Message.ID || w.Ses.State.User.ID == reaction.UserID {
continue
}
if v, ok := w.Handlers[reaction.Emoji.Name]; ok {
if w.isUserAllowed(reaction.UserID) {
go v(w, reaction)
}
}
if w.DeleteReactions {
go func() {
if w.isUserAllowed(reaction.UserID) {
time.Sleep(time.Millisecond * 250)
w.Ses.MessageReactionRemove(reaction.ChannelID, reaction.MessageID, reaction.Emoji.Name, reaction.UserID)
}
}()
}
}
}
// Handle adds a handler for the given emoji name
// emojiName: The unicode value of the emoji
// handler : handler function to call when the emoji is clicked
// func(*Widget, *discordgo.MessageReaction)
func (w *Widget) Handle(emojiName string, handler WidgetHandler) error {
if _, ok := w.Handlers[emojiName]; !ok {
w.Keys = append(w.Keys, emojiName)
w.Handlers[emojiName] = handler
}
// if the widget is running, append the added emoji to the message.
if w.Running() && w.Message != nil {
return w.Ses.MessageReactionAdd(w.Message.ChannelID, w.Message.ID, emojiName)
}
return nil
}
// QueryInput queries the user with ID `id` for input
// prompt : Question prompt
// userID : UserID to get message from
// timeout: How long to wait for the user's response
func (w *Widget) QueryInput(prompt string, userID string, timeout time.Duration) (*discordgo.Message, error) {
msg, err := w.Ses.ChannelMessageSend(w.ChannelID, "<@"+userID+">, "+prompt)
if err != nil {
return nil, err
}
defer func() {
w.Ses.ChannelMessageDelete(msg.ChannelID, msg.ID)
}()
timeoutChan := make(chan int)
go func() {
time.Sleep(timeout)
timeoutChan <- 0
}()
for {
select {
case userMsg := <-nextMessageCreateC(w.Ses):
if userMsg.Author.ID != userID {
continue
}
w.Ses.ChannelMessageDelete(userMsg.ChannelID, userMsg.ID)
return userMsg.Message, nil
case <-timeoutChan:
return nil, errors.New("timed out")
}
}
}
// Running returns w.running
func (w *Widget) Running() bool {
w.Lock()
running := w.running
w.Unlock()
return running
}
// UpdateEmbed updates the embed object and edits the original message
// embed: New embed object to replace w.Embed
func (w *Widget) UpdateEmbed(embed *discordgo.MessageEmbed) (*discordgo.Message, error) {
if w.Message == nil {
return nil, ErrNilMessage
}
return w.Ses.ChannelMessageEditEmbed(w.ChannelID, w.Message.ID, embed)
}
// Reset resets timeout ticker by duration. Returns ErrTickerNotSet when ticker is nil
// d: New ticker duration
func (w *Widget) Reset(d time.Duration) error {
if w.ticker != nil {
w.Lock()
w.ticker.Reset(d)
w.Unlock()
return nil
}
return ErrTickerNotSet
}
// Stop stops timeout ticker. Returns ErrTickerNotSet when ticker is nil
func (w *Widget) Stop() error {
if w.ticker != nil {
w.Lock()
w.ticker.Stop()
w.Unlock()
}
return ErrTickerNotSet
}
// RefreshTimeout refreshes timeout ticker by Widget's timeout. Returns ErrTickerNotSet when ticker is nil
func (w *Widget) RefreshTimeout() error {
return w.Reset(w.Timeout)
}
// LockToUsers adds defined userIDs to the UserWhitelist locking the widget to them
func (w *Widget) LockToUsers(userIDs ...string) error {
if len(userIDs) == 0 {
return errors.New("userID can't be empty")
}
w.UserWhitelist = append(w.UserWhitelist, userIDs...)
return nil
}