✨ thetharus 增加 傲娇 可爱 词库

plugin/thesaurus/chat.go

@@ -3,31 +3,65 @@ package thesaurus
 
 import (
 	"encoding/json"
+	"io"
+	"io/fs"
 	"math/rand"
 
+	"github.com/FloatTech/floatbox/ctxext"
 	ctrl "github.com/FloatTech/zbpctrl"
 	"github.com/FloatTech/zbputils/control"
+	"github.com/fumiama/jieba"
 	"github.com/sirupsen/logrus"
 	zero "github.com/wdvxdr1123/ZeroBot"
 	"github.com/wdvxdr1123/ZeroBot/message"
+	"gopkg.in/yaml.v3"
 )
 
-type kimo = map[string][]string
-
-type simai struct {
-	D map[string][]string `yaml:"傲娇"`
-	K map[string][]string `yaml:"可爱"`
-}
-
 func init() {
 	engine := control.Register("thesaurus", &ctrl.Options[*zero.Ctx]{
 		DisableOnDefault: false,
 		Brief:            "词典匹配回复",
-		Help:             "- 词典匹配回复",
+		Help:             "- 切换[kimo|傲娇|可爱]词库",
 		PublicDataFolder: "Chat",
 	})
+	engine.OnRegex(`^切换(kimo|傲娇|可爱)词库$`, zero.AdminPermission).SetBlock(true).Handle(func(ctx *zero.Ctx) {
+		c, ok := ctx.State["manager"].(*ctrl.Control[*zero.Ctx])
+		if !ok {
+			ctx.SendChain(message.Text("ERROR: 找不到 manager"))
+			return
+		}
+		gid := ctx.Event.GroupID
+		if gid == 0 {
+			gid = -ctx.Event.UserID
+		}
+		switch ctx.State["regex_matched"].([]string)[1] {
+		case "kimo":
+			c.SetData(gid, TYPKIMO)
+		case "傲娇":
+			c.SetData(gid, TYPDERE)
+		case "可爱":
+			c.SetData(gid, TYPKAWA)
+		}
+	})
 	go func() {
-		data, err := engine.GetLazyData("kimoi.json", false)
+		data, err := engine.GetLazyData("dict.txt", false)
+		if err != nil {
+			panic(err)
+		}
+		seg, err := jieba.LoadDictionary(&mockfile{data: data})
+		if err != nil {
+			panic(err)
+		}
+		smd, err := engine.GetLazyData("simai.yml", false)
+		if err != nil {
+			panic(err)
+		}
+		sm := simai{D: make(map[string][]string, 8192), K: make(map[string][]string, 16384)}
+		err = yaml.Unmarshal(smd, &sm)
+		if err != nil {
+			panic(err)
+		}
+		data, err = engine.GetLazyData("kimoi.json", false)
 		if err != nil {
 			panic(err)
 		}
@@ -36,17 +70,94 @@ func init() {
 		if err != nil {
 			panic(err)
 		}
-		chatList := make([]string, 0, 256)
+		chatList := make([]string, 0, len(kimomap))
 		for k := range kimomap {
 			chatList = append(chatList, k)
 		}
 		logrus.Infoln("[thesaurus]加载", len(chatList), "条kimoi")
-		engine.OnFullMatchGroup(chatList, zero.OnlyToMe).SetBlock(true).Handle(
-			func(ctx *zero.Ctx) {
-				key := ctx.MessageString()
-				val := kimomap[key]
-				text := val[rand.Intn(len(val))]
-				ctx.SendChain(message.Reply(ctx.Event.MessageID), message.Text(text)) // 来自于 https://github.com/Kyomotoi/AnimeThesaurus 的回复 经过二次修改
-			})
+
+		chatListD := make([]string, 0, len(sm.D))
+		for k := range sm.D {
+			chatListD = append(chatListD, k)
+		}
+		chatListK := make([]string, 0, len(sm.K))
+		for k := range sm.K {
+			chatListK = append(chatListK, k)
+		}
+		logrus.Infoln("[thesaurus]加载", len(chatListD), "条傲娇词库", len(chatListK), "条可爱词库")
+
+		engine.OnMessage(canmatch(TYPKIMO),
+			ctxext.JiebaFullMatch(seg, getmsg, chatList...),
+		).SetBlock(false).Handle(randreply(kimomap))
+		engine.OnMessage(canmatch(TYPDERE),
+			ctxext.JiebaFullMatch(seg, getmsg, chatListD...),
+		).SetBlock(false).Handle(randreply(sm.D))
+		engine.OnMessage(canmatch(TYPKAWA),
+			ctxext.JiebaFullMatch(seg, getmsg, chatListK...),
+		).SetBlock(false).Handle(randreply(sm.K))
 	}()
 }
+
+type kimo = map[string][]string
+
+type simai struct {
+	D map[string][]string `yaml:"傲娇"`
+	K map[string][]string `yaml:"可爱"`
+}
+
+type mockfile struct {
+	p    uintptr
+	data []byte
+}
+
+func (*mockfile) Stat() (fs.FileInfo, error) {
+	return nil, nil
+}
+func (f *mockfile) Read(buf []byte) (int, error) {
+	if int(f.p) >= len(f.data) {
+		return 0, io.EOF
+	}
+	n := copy(buf, f.data[f.p:])
+	f.p += uintptr(n)
+	return n, nil
+}
+func (f *mockfile) Close() error {
+	if f.data == nil {
+		return fs.ErrClosed
+	}
+	f.data = nil
+	return nil
+}
+
+const (
+	TYPKIMO = iota
+	TYPDERE
+	TYPKAWA
+)
+
+func canmatch(typ int64) zero.Rule {
+	return func(ctx *zero.Ctx) bool {
+		c, ok := ctx.State["manager"].(*ctrl.Control[*zero.Ctx])
+		if !ok {
+			return false
+		}
+		gid := ctx.Event.GroupID
+		if gid == 0 {
+			gid = -ctx.Event.UserID
+		}
+		return c.GetData(gid) == typ
+	}
+}
+
+func getmsg(ctx *zero.Ctx) string {
+	return ctx.MessageString()
+}
+
+func randreply(m map[string][]string) zero.Handler {
+	return func(ctx *zero.Ctx) {
+		key := ctx.State["matched"].(string)
+		val := m[key]
+		text := val[rand.Intn(len(val))]
+		ctx.SendChain(message.Reply(ctx.Event.MessageID), message.Text(text))
+	}
+}