searcher.go

// Package saucenao P站ID/saucenao/ascii2d搜图
package saucenao

import (
	"fmt"
	"net/http"
	"os"
	"reflect"
	"strconv"

	"github.com/fumiama/terasu/http2"
	"github.com/sirupsen/logrus"
	zero "github.com/wdvxdr1123/ZeroBot"
	"github.com/wdvxdr1123/ZeroBot/message"

	"github.com/FloatTech/AnimeAPI/ascii2d"
	"github.com/FloatTech/AnimeAPI/pixiv"
	"github.com/jozsefsallai/gophersauce"

	"github.com/FloatTech/floatbox/binary"
	"github.com/FloatTech/floatbox/file"
	ctrl "github.com/FloatTech/zbpctrl"
	"github.com/FloatTech/zbputils/control"
	"github.com/FloatTech/zbputils/ctxext"
	"github.com/FloatTech/zbputils/img/pool"
)

const (
	enableHex = 0x10
	unableHex = 0x7fffffff_fffffffd
)

var (
	saucenaocli *gophersauce.Client
)

func init() { // 插件主体
	engine := control.AutoRegister(&ctrl.Options[*zero.Ctx]{
		DisableOnDefault: false,
		Brief:            "以图搜图",
		Help: "- 以图搜图 | 搜索图片 | 以图识图[图片]\n" +
			"- 搜图[P站图片ID]",
		PrivateDataFolder: "saucenao",
	})
	apikeyfile := engine.DataFolder() + "apikey.txt"
	if file.IsExist(apikeyfile) {
		key, err := os.ReadFile(apikeyfile)
		if err != nil {
			panic(err)
		}
		saucenaocli, err = gophersauce.NewClient(&gophersauce.Settings{
			MaxResults: 1,
			APIKey:     binary.BytesToString(key),
		})
		if err != nil {
			panic(err)
		}
	}
	// 根据 PID 搜图
	engine.OnRegex(`^搜图(\d+)$`).SetBlock(true).
		Handle(func(ctx *zero.Ctx) {
			id, _ := strconv.ParseInt(ctx.State["regex_matched"].([]string)[1], 10, 64)
			ctx.SendChain(message.Text("少女祈祷中......"))
			// 获取P站插图信息
			illust, err := pixiv.Works(id)
			if err != nil {
				ctx.SendChain(message.Text("ERROR: ", err))
				return
			}
			if illust.Pid > 0 {
				name := strconv.FormatInt(illust.Pid, 10)
				var imgs message.Message
				for i := range illust.ImageUrls {
					f := file.BOTPATH + "/" + illust.Path(i)
					n := name + "_p" + strconv.Itoa(i)
					var m *pool.Image
					if file.IsNotExist(f) {
						m, err = pool.GetImage(n)
						if err == nil {
							imgs = append(imgs, message.Image(m.String()))
							continue
						}
						logrus.Debugln("[saucenao]开始下载", n)
						logrus.Debugln("[saucenao]urls:", illust.ImageUrls)
						err1 := illust.DownloadToCache(i)
						if err1 == nil {
							m.SetFile(f)
							_, _ = m.Push(ctxext.SendToSelf(ctx), ctxext.GetMessage(ctx))
						}
						if err1 != nil {
							logrus.Debugln("[saucenao]下载err:", err1)
						}
					}
					imgs = append(imgs, message.Image("file:///"+f))
				}
				txt := message.Text(
					"标题: ", illust.Title, "\n",
					"插画ID: ", illust.Pid, "\n",
					"画师: ", illust.UserName, "\n",
					"画师ID: ", illust.UserID, "\n",
					"直链: ", "https://pixivel.moe/detail?id=", illust.Pid,
				)
				if imgs != nil {
					ctx.Send(message.Message{ctxext.FakeSenderForwardNode(ctx, txt),
						ctxext.FakeSenderForwardNode(ctx, imgs...)})
				} else {
					// 图片下载失败，仅发送文字结果
					ctx.SendChain(txt)
				}
			} else {
				ctx.SendChain(message.Text("图片不存在!"))
			}
		})
	// 以图搜图
	engine.OnKeywordGroup([]string{"以图搜图", "搜索图片", "以图识图"}, zero.MustProvidePicture).SetBlock(true).
		Handle(func(ctx *zero.Ctx) {
			// 开始搜索图片
			pics, ok := ctx.State["image_url"].([]string)
			showPic := false
			if !ok {
				ctx.SendChain(message.Text("ERROR: 未获取到图片链接"))
				return
			}
			c, ok := ctx.State["manager"].(*ctrl.Control[*zero.Ctx])
			if ok && c.GetData(ctx.Event.GroupID)&enableHex == enableHex {
				showPic = true
			}
			ctx.SendChain(message.Text("少女祈祷中..."))
			for _, pic := range pics {
				if saucenaocli != nil {
					resp, err := saucenaocli.FromURL(pic)
					if err == nil && resp.Count() > 0 {
						result := resp.First()
						s, err := strconv.ParseFloat(result.Header.Similarity, 64)
						if err == nil {
							rr := reflect.ValueOf(&result.Data).Elem()
							b := binary.NewWriterF(func(w *binary.Writer) {
								r := rr.Type()
								for i := 0; i < r.NumField(); i++ {
									if !rr.Field(i).IsZero() {
										w.WriteString("\n")
										w.WriteString(r.Field(i).Name)
										w.WriteString(": ")
										w.WriteString(fmt.Sprint(rr.Field(i).Interface()))
									}
								}
							})
							resp, err := http2.Head(result.Header.Thumbnail)
							msg := make(message.Message, 0, 3)
							if s > 80.0 {
								msg = append(msg, message.Text("我有把握是这个!"))
							} else {
								msg = append(msg, message.Text("也许是这个?"))
							}
							if showPic {
								if err == nil {
									_ = resp.Body.Close()
									if resp.StatusCode == http.StatusOK {
										msg = append(msg, message.Image(result.Header.Thumbnail))
									} else {
										msg = append(msg, message.Image(pic))
									}
								} else {
									msg = append(msg, message.Image(pic))
								}
							}
							msg = append(msg, message.Text("\n图源: ", result.Header.IndexName, binary.BytesToString(b)))
							ctx.Send(message.Message{ctxext.FakeSenderForwardNode(ctx, msg...)})
							if s > 80.0 {
								continue
							}
						}
					}
				} else {
					ctx.SendChain(message.Text("请私聊发送 设置 saucenao api key [apikey] 以启用 saucenao 搜图 (方括号不需要输入), key 请前往 https://saucenao.com/user.php?page=search-api 获取"))
				}
				// ascii2d 搜索
				result, err := ascii2d.ASCII2d(pic)
				if err != nil {
					ctx.SendChain(message.Text("ERROR: ", err))
					continue
				}
				msg := message.Message{ctxext.FakeSenderForwardNode(ctx, message.Text("ascii2d搜图结果"))}
				for i := 0; i < len(result) && i < 5; i++ {
					var resultMsgs message.Message
					if showPic {
						resultMsgs = append(resultMsgs, message.Image(result[i].Thumb))
					}
					resultMsgs = append(resultMsgs, message.Text(fmt.Sprintf(
						"标题: %s\n图源: %s\n画师: %s\n画师链接: %s\n图片链接: %s",
						result[i].Name,
						result[i].Type,
						result[i].AuthNm,
						result[i].Author,
						result[i].Link,
					)))
					msg = append(msg, ctxext.FakeSenderForwardNode(ctx, resultMsgs...))
				}
				if id := ctx.Send(msg).ID(); id == 0 {
					ctx.SendChain(message.Text("ERROR: 可能被风控了"))
				}
			}
		})
	engine.OnRegex(`^设置\s?saucenao\s?api\s?key\s?([0-9a-f]{40})$`, zero.SuperUserPermission, zero.OnlyPrivate).SetBlock(true).
		Handle(func(ctx *zero.Ctx) {
			var err error
			saucenaocli, err = gophersauce.NewClient(&gophersauce.Settings{
				MaxResults: 1,
				APIKey:     ctx.State["regex_matched"].([]string)[1],
			})
			if err != nil {
				ctx.SendChain(message.Text("ERROR: ", err))
				return
			}
			err = os.WriteFile(apikeyfile, binary.StringToBytes(saucenaocli.APIKey), 0644)
			if err != nil {
				ctx.SendChain(message.Text("ERROR: ", err))
				return
			}
			ctx.SendChain(message.Text("成功!"))
		})
	engine.OnRegex(`^(开启|打开|启用|关闭|关掉|禁用)搜图显示图片$`, zero.AdminPermission).SetBlock(true).
		Handle(func(ctx *zero.Ctx) {
			gid := ctx.Event.GroupID
			if gid <= 0 {
				// 个人用户设为负数
				gid = -ctx.Event.UserID
			}
			option := ctx.State["regex_matched"].([]string)[1]
			c, ok := ctx.State["manager"].(*ctrl.Control[*zero.Ctx])
			if !ok {
				ctx.SendChain(message.Text("找不到服务!"))
				return
			}
			data := c.GetData(ctx.Event.GroupID)
			switch option {
			case "开启", "打开", "启用":
				data |= enableHex
			case "关闭", "关掉", "禁用":
				data &= unableHex
			default:
				return
			}
			err := c.SetData(gid, data)
			if err != nil {
				ctx.SendChain(message.Text("出错啦: ", err))
				return
			}
			ctx.SendChain(message.Text("已", option, "搜图显示图片"))
		})
}