snowluma/client.go

package snowluma

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

// HandlerFunc 是事件处理函数。
type HandlerFunc func(c *Context)

// CommandFunc 是命令处理函数，args 为按空白切分的参数。
type CommandFunc func(c *Context, args []string)

// Connector 是一种通信方式（正向 WS / 反向 WS / HTTP 上报 / HTTP API）。
// 常规连接直接用 Client 的 Set* 方法配置，本接口留给自定义传输扩展。
type Connector interface {
	Name() string
	Serve(ctx context.Context, bot *Client) error
}

type handlerSet struct {
	onEvent     []HandlerFunc
	onMessage   []HandlerFunc
	onPrivate   []HandlerFunc
	onGroup     []HandlerFunc
	onNotice    []HandlerFunc
	onRequest   []HandlerFunc
	onMeta      []HandlerFunc
	noticeTypes map[string][]HandlerFunc
	commands    map[string]CommandFunc
}

// Client 是与 SnowLuma 交互的客户端对象：Set* 配置连接，On* 注册事件处理，
// Run 启动后直接调用其内嵌的 API 方法即可操作 SnowLuma。
type Client struct {
	// API 内嵌了全部类型化 API 方法，走 Client 的默认调用通道。
	API

	log        logger
	token      string // OneBot access_token
	secret     string // HTTP 上报 HMAC-SHA1 密钥
	connectors []Connector
	httpAPI    *HTTPAPI
	current    atomic.Pointer[Session]
	lastBeat   atomic.Int64

	sessMu   sync.Mutex
	sessions map[*Session]struct{}

	mu sync.Mutex
	hs handlerSet
}

// New 创建一个与 SnowLuma 交互的客户端对象。
// 用 Set* 方法配置连接后调用 Run 启动：
//
//	bot := snowluma.New().SetToken("xxx").SetForwardWS("ws://127.0.0.1:3001")
//	bot.OnMessage(func(c *snowluma.Context) { c.ReplyText("hi") })
//	bot.Run(ctx)
func New() *Client {
	c := &Client{log: slog.Default(), sessions: map[*Session]struct{}{}}
	c.API = API{call: c}
	c.hs.noticeTypes = map[string][]HandlerFunc{}
	c.hs.commands = map[string]CommandFunc{}
	return c
}

// SetToken 设置 OneBot access_token（WS 握手与 HTTP 调用自动携带，可留空）。
func (c *Client) SetToken(token string) *Client {
	c.mu.Lock()
	c.token = token
	if c.httpAPI != nil {
		c.httpAPI.token.Store(token)
	}
	c.mu.Unlock()
	return c
}

// SetSecret 设置 HTTP 上报的 HMAC-SHA1 校验密钥，留空表示不校验。
func (c *Client) SetSecret(secret string) *Client {
	c.mu.Lock()
	c.secret = secret
	c.mu.Unlock()
	return c
}

// SetLogger 设置自定义日志器，默认 slog.Default()。
func (c *Client) SetLogger(l *slog.Logger) *Client {
	if l == nil {
		l = slog.Default()
	}
	c.log = l
	return c
}

// SetForwardWS 配置正向 WebSocket：本框架作为客户端连接 SnowLuma 的 WS 服务。
// 可传多个地址同时连接多个 SnowLuma 实例；重复调用会替换之前的正向 WS 配置。
func (c *Client) SetForwardWS(urls ...string) *Client {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.connectors = dropConnectors[*ForwardWS](c.connectors)
	for _, u := range urls {
		c.connectors = append(c.connectors, &ForwardWS{url: u})
	}
	return c
}

// SetReverseWS 配置反向 WebSocket：本框架监听 addr，由 SnowLuma 主动连入。
// 每个监听地址可承接任意多个 SnowLuma 实例；重复调用会替换之前的反向 WS 配置。
func (c *Client) SetReverseWS(addrs ...string) *Client {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.connectors = dropConnectors[*ReverseWS](c.connectors)
	for _, a := range addrs {
		c.connectors = append(c.connectors, &ReverseWS{addr: a})
	}
	return c
}

// SetHTTPAPI 配置 HTTP API：执行 action 时 POST 到 SnowLuma 的 HTTP 服务。
// 没有任何 WS 会话时 Client.Call 走该通道；重复调用替换之前的配置。
func (c *Client) SetHTTPAPI(baseURL string) *Client {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.connectors = dropConnectors[*HTTPAPI](c.connectors)
	api := NewHTTPAPI(baseURL)
	api.token.Store(c.token)
	c.httpAPI = api
	c.connectors = append(c.connectors, api)
	return c
}

// SetHTTPReport 配置 HTTP POST 上报：本框架监听 addr，SnowLuma 把事件 POST 过来，
// handler 里的 Reply/Approve 以 OneBot 快速操作形式随响应返回。
// 可配置多个监听地址；重复调用会替换之前的上报配置。
func (c *Client) SetHTTPReport(addrs ...string) *Client {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.connectors = dropConnectors[*HTTPReport](c.connectors)
	for _, a := range addrs {
		c.connectors = append(c.connectors, &HTTPReport{addr: a})
	}
	return c
}

// Use 注册自定义 Connector（高级用法；常规连接请用 Set* 方法）。
func (c *Client) Use(conns ...Connector) {
	c.mu.Lock()
	defer c.mu.Unlock()
	for _, conn := range conns {
		if api, ok := conn.(*HTTPAPI); ok {
			c.httpAPI = api
			api.token.Store(c.token)
		}
	}
	c.connectors = append(c.connectors, conns...)
}

// Token 返回当前 access_token。
func (c *Client) Token() string { return c.token }

// dropConnectors 过滤掉类型为 T 的连接器。
func dropConnectors[T Connector](conns []Connector) []Connector {
	out := make([]Connector, 0, len(conns))
	for _, conn := range conns {
		if _, ok := any(conn).(T); ok {
			continue
		}
		out = append(out, conn)
	}
	return out
}

// SelfID 返回当前会话的机器人 QQ 号（连接后有效，否则 0）。
func (c *Client) SelfID() int64 {
	if s := c.current.Load(); s != nil {
		return s.SelfID()
	}
	return 0
}

// LastHeartbeat 返回最近一次心跳时间；从未收到过心跳则返回零值。
func (c *Client) LastHeartbeat() time.Time {
	unix := c.lastBeat.Load()
	if unix == 0 {
		return time.Time{}
	}
	return time.Unix(0, unix)
}

// Run 启动所有已配置的连接并阻塞，直到 ctx 取消（正常返回 nil）或某个连接致命出错。
// 取消后会主动关闭所有会话与监听，Run 返回即代表端口与连接已释放。
// 请在 Run 之前完成 Set* 配置与 On* 注册。
func (c *Client) Run(ctx context.Context) error {
	c.mu.Lock()
	conns := append([]Connector(nil), c.connectors...)
	c.mu.Unlock()
	if len(conns) == 0 {
		return errors.New("snowluma: no connection configured, call SetForwardWS/SetReverseWS/SetHTTPReport first")
	}
	runCtx, cancel := context.WithCancel(ctx)
	defer cancel()

	var wg sync.WaitGroup
	errCh := make(chan error, len(conns))
	for _, conn := range conns {
		conn := conn
		wg.Add(1)
		go func() {
			defer wg.Done()
			err := conn.Serve(runCtx, c)
			if err == nil || errors.Is(err, context.Canceled) || errors.Is(err, http.ErrServerClosed) {
				return
			}
			select {
			case errCh <- fmt.Errorf("connector %s: %w", conn.Name(), err):
			default:
			}
		}()
	}

	var runErr error
	select {
	case runErr = <-errCh:
		c.log.Error("snowluma: connector fatal, shutting down", "err", runErr)
	case <-ctx.Done():
	}
	cancel()
	c.closeSessions()
	shCtx, shCancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer shCancel()
	for _, conn := range conns {
		if sh, ok := conn.(interface{ Shutdown(context.Context) error }); ok {
			_ = sh.Shutdown(shCtx)
		}
	}
	wg.Wait()
	if runErr != nil {
		return runErr
	}
	return nil
}

func (c *Client) closeSessions() {
	c.sessMu.Lock()
	defer c.sessMu.Unlock()
	for s := range c.sessions {
		_ = s.Close()
	}
}

// Call 实现 Caller：优先走当前 WS 会话，其次走 HTTP API，都没有则报错。
func (c *Client) Call(ctx context.Context, action string, params any, out any) error {
	if s := c.current.Load(); s != nil {
		return s.Call(ctx, action, params, out)
	}
	if c.httpAPI != nil {
		return c.httpAPI.Call(ctx, action, params, out)
	}
	return ErrNotConnected
}

func (c *Client) bindSession(s *Session) {
	c.sessMu.Lock()
	c.sessions[s] = struct{}{}
	c.sessMu.Unlock()
	c.current.Store(s)
}

func (c *Client) unbindSession(s *Session) {
	c.sessMu.Lock()
	delete(c.sessions, s)
	c.sessMu.Unlock()
	c.current.CompareAndSwap(s, nil)
}

// handleRaw 处理 WS 连接上的一条原始消息：事件则入队分发，带 echo 的则为 API 响应。
func (c *Client) handleRaw(s *Session, data []byte) {
	var probe struct {
		PostType string `json:"post_type"`
	}
	if err := json.Unmarshal(data, &probe); err != nil {
		c.log.Warn("snowluma: invalid payload", "err", err)
		return
	}
	if probe.PostType != "" {
		ev, err := decodeEvent(data)
		if err != nil {
			c.log.Warn("snowluma: decode event failed", "err", err)
			return
		}
		if ev.SelfID != 0 {
			s.selfID.Store(ev.SelfID)
		}
		if ev.Meta != nil {
			if ev.Meta.MetaType == "heartbeat" {
				c.lastBeat.Store(time.Now().UnixNano())
			}
			if ev.Meta.MetaType == "lifecycle" && ev.Meta.SubType == "connect" {
				c.log.Info("snowluma: bot connected", "self_id", ev.SelfID)
			}
		}
		s.enqueue(ev, nil)
		return
	}
	s.caller.routeResponse(data, c.log)
}

// ---- 事件处理注册 ----

// OnEvent 注册处理所有事件的函数。
func (c *Client) OnEvent(fn HandlerFunc) { c.add(&c.hs.onEvent, fn) }

// OnMessage 注册处理私聊+群聊消息事件（含 message_sent）的函数。
func (c *Client) OnMessage(fn HandlerFunc) { c.add(&c.hs.onMessage, fn) }

// OnPrivateMessage 注册处理私聊消息事件的函数。
func (c *Client) OnPrivateMessage(fn HandlerFunc) { c.add(&c.hs.onPrivate, fn) }

// OnGroupMessage 注册处理群聊消息事件的函数。
func (c *Client) OnGroupMessage(fn HandlerFunc) { c.add(&c.hs.onGroup, fn) }

// OnNotice 注册处理所有通知事件的函数。
func (c *Client) OnNotice(fn HandlerFunc) { c.add(&c.hs.onNotice, fn) }

// OnNoticeType 注册处理指定 notice_type 通知事件的函数（如 "notify"、"group_recall"）。
func (c *Client) OnNoticeType(noticeType string, fn HandlerFunc) {
	c.mu.Lock()
	c.hs.noticeTypes[noticeType] = append(c.hs.noticeTypes[noticeType], fn)
	c.mu.Unlock()
}

// OnRequest 注册处理所有请求事件（好友/群）的函数。
func (c *Client) OnRequest(fn HandlerFunc) { c.add(&c.hs.onRequest, fn) }

// OnMeta 注册处理元事件（心跳/生命周期）的函数。
func (c *Client) OnMeta(fn HandlerFunc) { c.add(&c.hs.onMeta, fn) }

// OnCommand 注册文本命令：消息文本等于 name 或以 name+空白 开头时触发。
// 命令名自带头尾风格（如 "/ping" 或 "ping"）。
func (c *Client) OnCommand(name string, fn CommandFunc) {
	c.mu.Lock()
	c.hs.commands[name] = fn
	c.mu.Unlock()
}

func (c *Client) add(dst *[]HandlerFunc, fn HandlerFunc) {
	c.mu.Lock()
	*dst = append(*dst, fn)
	c.mu.Unlock()
}

func (c *Client) snapshot() handlerSet {
	c.mu.Lock()
	defer c.mu.Unlock()
	hs := handlerSet{
		onEvent:     slicesClone(c.hs.onEvent),
		onMessage:   slicesClone(c.hs.onMessage),
		onPrivate:   slicesClone(c.hs.onPrivate),
		onGroup:     slicesClone(c.hs.onGroup),
		onNotice:    slicesClone(c.hs.onNotice),
		onRequest:   slicesClone(c.hs.onRequest),
		onMeta:      slicesClone(c.hs.onMeta),
		noticeTypes: mapsCloneOf(c.hs.noticeTypes),
		commands:    mapsCloneCmd(c.hs.commands),
	}
	return hs
}

// dispatch 按事件类型把事件交给已注册的处理函数；HTTP 上报模式下 quick 非空。
func (c *Client) dispatch(ev *Event, caller Caller, quick *quickBox) {
	if caller == nil {
		caller = c
	}
	ctx := &Context{Bot: c, Event: ev, API: API{call: caller}, quick: quick}
	hs := c.snapshot()

	run := func(fns []HandlerFunc) {
		for _, fn := range fns {
			func() {
				defer func() {
					if r := recover(); r != nil {
						c.log.Error("snowluma: handler panic", "post_type", ev.PostType, "panic", fmt.Sprint(r))
					}
				}()
				fn(ctx)
			}()
		}
	}

	run(hs.onEvent)
	switch ev.PostType {
	case "message", "message_sent":
		run(hs.onMessage)
		if ev.Message != nil {
			if ev.Message.MessageType == "group" {
				run(hs.onGroup)
			} else {
				run(hs.onPrivate)
			}
			c.matchCommand(ctx, hs, run)
		}
	case "notice":
		run(hs.onNotice)
		if ev.Notice != nil {
			run(hs.noticeTypes[ev.Notice.NoticeType])
		}
	case "request":
		run(hs.onRequest)
	case "meta_event":
		run(hs.onMeta)
	}
}

// matchCommand 命中已注册命令则调用对应处理函数。
func (c *Client) matchCommand(ctx *Context, hs handlerSet, run func([]HandlerFunc)) {
	if len(hs.commands) == 0 || ctx.Event.Message == nil {
		return
	}
	text := strings.TrimSpace(SegmentsText(ctx.Segments()))
	if text == "" {
		return
	}
	for name, fn := range hs.commands {
		if text != name && !strings.HasPrefix(text, name+" ") {
			continue
		}
		rest := strings.TrimSpace(strings.TrimPrefix(text, name))
		args := []string{}
		if rest != "" {
			args = strings.Fields(rest)
		}
		func() {
			defer func() {
				if r := recover(); r != nil {
					c.log.Error("snowluma: command handler panic", "cmd", name, "panic", fmt.Sprint(r))
				}
			}()
			fn(ctx, args)
		}()
		return
	}
}

// 小工具：避免每个文件都引 slices/maps 包。
func slicesClone(s []HandlerFunc) []HandlerFunc {
	if s == nil {
		return nil
	}
	out := make([]HandlerFunc, len(s))
	copy(out, s)
	return out
}

func mapsCloneOf(m map[string][]HandlerFunc) map[string][]HandlerFunc {
	out := make(map[string][]HandlerFunc, len(m))
	for k, v := range m {
		out[k] = slicesClone(v)
	}
	return out
}

func mapsCloneCmd(m map[string]CommandFunc) map[string]CommandFunc {
	out := make(map[string]CommandFunc, len(m))
	for k, v := range m {
		out[k] = v
	}
	return out
}