ZeroBot-Plugin/plugin/thesaurus/chat.go

// Package thesaurus 修改过的单纯回复插件
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"
)

func init() {
	engine := control.Register("thesaurus", &ctrl.Options[*zero.Ctx]{
		DisableOnDefault: false,
		Brief:            "词典匹配回复",
		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("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)
		}
		kimomap := make(kimo, 256)
		err = json.Unmarshal(data, &kimomap)
		if err != nil {
			panic(err)
		}
		chatList := make([]string, 0, len(kimomap))
		for k := range kimomap {
			chatList = append(chatList, k)
		}
		logrus.Infoln("[thesaurus]加载", len(chatList), "条kimoi")

		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))
	}
}