Cyrene/backend/platform-bridge/internal/adapter/qq/adapter.go

package qq

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"

	"git.yeij.top/AskaEth/Cyrene/platform-bridge/internal/bridge"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

// Adapter implements PlatformAdapter for QQ via OneBot v11 WebSocket.
// Supports two modes:
//   - "server" (正向 WS): adapter starts a WS server, NapCat connects as client.
//   - "client" (反向 WS): adapter connects to NapCat's WS server as a client.
type Adapter struct {
	configName     string // instance name, e.g. "qq-home"
	mode           string // "client" or "server"
	port           string
	accessToken    string
	remoteURL      string // NapCat OneBot WS server URL, used in client mode
	sendIntervalMs int    // minimum interval between consecutive messages
	selfID         string // bot's own QQ number, populated from incoming messages
	conn           *websocket.Conn
	connMu         sync.Mutex
	connected      bool
	connectedAt    time.Time // when the connection was established (for historical message detection)
	srv            *http.Server // HTTP server for WS upgrades (server mode only)

	groupNames   map[int64]string    // group ID → group name cache
	groupNamesMu sync.RWMutex

	pendingResponses map[string]chan *OBv11APIResponse
	respMu           sync.Mutex
}

func NewAdapter(configName, mode, port, accessToken, remoteURL string, sendIntervalMs int) *Adapter {
	if mode == "" {
		mode = "server"
	}
	return &Adapter{
		configName:       configName,
		mode:             mode,
		port:             port,
		accessToken:      accessToken,
		remoteURL:        remoteURL,
		sendIntervalMs:   sendIntervalMs,
		pendingResponses: make(map[string]chan *OBv11APIResponse),
		groupNames:       make(map[int64]string),
	}
}

func (a *Adapter) PlatformName() string   { return "qq" }
func (a *Adapter) ConfigName() string    { return a.configName }
func (a *Adapter) SendIntervalMs() int   { return a.sendIntervalMs }
func (a *Adapter) SelfID() string        { return a.selfID }

// ConnectedAt returns the time the connection was established, for historical message detection.
func (a *Adapter) ConnectedAt() time.Time {
	a.connMu.Lock()
	defer a.connMu.Unlock()
	return a.connectedAt
}

// GroupName resolves a group ID to its name. Returns empty string if unknown.
func (a *Adapter) GroupName(groupID int64) string {
	a.groupNamesMu.RLock()
	n, ok := a.groupNames[groupID]
	a.groupNamesMu.RUnlock()
	if ok {
		return n
	}
	// Try to resolve via API.
	a.fetchGroupName(groupID)
	return ""
}

// SetGroupName caches a group name for a group ID.
func (a *Adapter) SetGroupName(groupID int64, name string) {
	a.groupNamesMu.Lock()
	a.groupNames[groupID] = name
	a.groupNamesMu.Unlock()
}

// fetchGroupName tries to resolve a group name via NapCat HTTP API (client mode).
func (a *Adapter) fetchGroupName(groupID int64) {
	if a.mode != "client" || a.remoteURL == "" {
		return
	}
	// Derive HTTP base from WS URL: ws://host:port → http://host:port
	httpBase := strings.Replace(a.remoteURL, "ws://", "http://", 1)
	httpBase = strings.Replace(httpBase, "wss://", "https://", 1)
	// Strip path suffix if present
	if idx := strings.LastIndex(httpBase, "/"); idx > 8 {
		httpBase = httpBase[:idx]
	}

	go func() {
		url := fmt.Sprintf("%s/get_group_info?group_id=%d", httpBase, groupID)
		if a.accessToken != "" {
			url += "&access_token=" + a.accessToken
		}
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
		if err != nil {
			return
		}
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return
		}
		defer resp.Body.Close()
		var result struct {
			Data struct {
				GroupName string `json:"group_name"`
			} `json:"data"`
		}
		if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
			return
		}
		if result.Data.GroupName != "" {
			a.SetGroupName(groupID, result.Data.GroupName)
		}
	}()
}

func (a *Adapter) Capabilities() bridge.PlatformCapabilities {
	return bridge.PlatformCapabilities{
		MaxMessageLength:   200,
		SupportsMarkdown:   true,
		SupportsImage:      true,
		SupportsVoice:      false,
		SupportsEmoji:      true,
		SupportsReaction:   false,
		SupportsTypingHint: false,
		RecommendBurstMax:  4,
	}
}

// checkAuth 验证 WebSocket 升级请求的 access_token。
// NapCat 通过两种方式传递: query 参数 ?access_token=xxx 或 Authorization: Bearer xxx 头.
func (a *Adapter) checkAuth(r *http.Request) bool {
	if a.accessToken == "" {
		return true // 未配置 token，允许所有连接
	}
	// 1) query 参数
	if r.URL.Query().Get("access_token") == a.accessToken {
		return true
	}
	// 2) Authorization: Bearer <token>
	auth := r.Header.Get("Authorization")
	if strings.HasPrefix(auth, "Bearer ") && strings.TrimPrefix(auth, "Bearer ") == a.accessToken {
		return true
	}
	return false
}

// wsHandler 统一的 WebSocket 连接处理 — 单连接承载 API 调用 + 事件推送 (OneBot 正向 WS).
func (a *Adapter) wsHandler(w http.ResponseWriter, r *http.Request) {
	if !a.checkAuth(r) {
		http.Error(w, "Forbidden: invalid access_token", http.StatusForbidden)
		return
	}
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		fmt.Printf("[qq] upgrade error: %v\n", err)
		return
	}
	a.connMu.Lock()
	// 关闭旧连接 (NapCat 重连)
	if a.conn != nil {
		a.conn.Close()
	}
	a.conn = conn
	a.connected = true
	a.connectedAt = time.Now()
	a.connMu.Unlock()
	fmt.Println("[qq] NapCat/OneBot connected (正向WS)")
}

// legacyHandler 兼容旧的路径 /ws/qq 和 /ws/qq/event.
func (a *Adapter) legacyHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Printf("[qq] legacy WS path %s connected (consider changing NapCat URL to root /)\n", r.URL.Path)
	a.wsHandler(w, r)
}

func (a *Adapter) Connect(ctx context.Context) error {
	if a.mode == "client" {
		return a.connectClient(ctx)
	}
	return a.connectServer()
}

func (a *Adapter) connectClient(ctx context.Context) error {
	url := a.remoteURL
	if a.accessToken != "" {
		sep := "?"
		if strings.Contains(url, "?") {
			sep = "&"
		}
		url += sep + "access_token=" + a.accessToken
	}

	dialer := websocket.DefaultDialer
	conn, _, err := dialer.DialContext(ctx, url, nil)
	if err != nil {
		return fmt.Errorf("dial NapCat WS %s: %w", url, err)
	}

	a.connMu.Lock()
	if a.conn != nil {
		a.conn.Close()
	}
	a.conn = conn
	a.connected = true
	a.connectedAt = time.Now()
	a.connMu.Unlock()

	fmt.Printf("[qq] connected to NapCat OneBot WS (client mode): %s\n", a.remoteURL)
	return nil
}

func (a *Adapter) connectServer() error {
	mux := http.NewServeMux()
	mux.HandleFunc("/", a.wsHandler)                       // NapCat 正向 WS 标准路径
	mux.HandleFunc("/ws/qq", a.legacyHandler)               // 向下兼容旧配置
	mux.HandleFunc("/ws/qq/event", a.legacyHandler)         // 向下兼容旧配置

	addr := ":" + a.port
	a.srv = &http.Server{Addr: addr, Handler: mux}
	go func() {
		fmt.Printf("[qq] WebSocket server on %s (waiting for NapCat forward WS connection)\n", addr)
		if a.accessToken != "" {
			fmt.Println("[qq] access_token 已配置，将验证连接请求")
		}
		if err := a.srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			fmt.Printf("[qq] server error: %v\n", err)
		}
	}()
	return nil
}

func (a *Adapter) Disconnect(ctx context.Context) error {
	a.connMu.Lock()
	if a.conn != nil {
		a.conn.Close()
		a.conn = nil
	}
	a.connected = false
	srv := a.srv
	a.srv = nil
	a.connMu.Unlock()

	if srv != nil {
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
		defer cancel()
		if err := srv.Shutdown(shutdownCtx); err != nil {
			return fmt.Errorf("qq server shutdown: %w", err)
		}
	}
	return nil
}

func (a *Adapter) IsConnected() bool {
	a.connMu.Lock()
	defer a.connMu.Unlock()
	return a.connected
}

func (a *Adapter) HealthCheck() error {
	if !a.IsConnected() {
		return fmt.Errorf("QQ bot not connected")
	}
	return nil
}

// ToUnified converts an OBv11 message to UnifiedMessage.
func (a *Adapter) ToUnified(rawMessage interface{}) (*bridge.UnifiedMessage, error) {
	msg, ok := rawMessage.(*OBv11Message)
	if !ok {
		return nil, fmt.Errorf("expected *OBv11Message, got %T", rawMessage)
	}

	content := extractText(msg)

	senderID := ""
	senderName := "unknown"
	channelType := "direct"
	channelID := ""

	switch msg.MessageType {
	case "private":
		senderID = fmt.Sprintf("%d", msg.UserID)
		channelType = "direct"
		channelID = fmt.Sprintf("private_%d", msg.UserID)
		if msg.Sender != nil {
			senderName = msg.Sender.Nickname
		}
	case "group":
		senderID = fmt.Sprintf("%d", msg.UserID)
		channelType = "group"
		channelID = fmt.Sprintf("%d", msg.GroupID)
		if msg.Sender != nil {
			if msg.Sender.Card != "" {
				senderName = msg.Sender.Card
			} else {
				senderName = msg.Sender.Nickname
			}
		}
	}

	var mentions []string
	if segments, ok := msg.Message.([]interface{}); ok {
		for _, s := range segments {
			if seg, ok := s.(map[string]interface{}); ok {
				if seg["type"] == "at" {
					if data, ok := seg["data"].(map[string]interface{}); ok {
						if qq, ok := data["qq"].(string); ok {
							mentions = append(mentions, qq)
						}
					}
				}
			}
		}
	}
	// Fallback: parse CQ at codes from string format (e.g. [CQ:at,qq=2254389756]).
	if len(mentions) == 0 {
		raw := msg.RawMessage
		if raw == "" {
			if s, ok := msg.Message.(string); ok {
				raw = s
			}
		}
		for _, m := range cqAtRegex.FindAllStringSubmatch(raw, -1) {
			if len(m) >= 2 {
				mentions = append(mentions, m[1])
			}
		}
	}

	// Resolve group name for group messages.
	groupName := ""
	if msg.MessageType == "group" {
		groupName = a.GroupName(msg.GroupID)
	}

	attachments := extractAttachments(msg)

	return &bridge.UnifiedMessage{
		SenderID:    senderID,
		SenderName:  senderName,
		Platform:    "qq",
		ChannelID:   channelID,
		ChannelType: channelType,
		Content:     content,
		ContentType: "text",
		MessageID:   fmt.Sprintf("%d", msg.MessageID),
		Mentions:    mentions,
		Attachments: attachments,
		RawData:     rawMessage,
		Timestamp:   time.Unix(msg.Time, 0),
			GroupName:   groupName,
	}, nil
}

// FromUnified converts a unified response to QQ platform messages.
func (a *Adapter) FromUnified(response *bridge.UnifiedResponse) ([]bridge.PlatformMessage, error) {
	var msgs []bridge.PlatformMessage
	for _, rm := range response.Messages {
		content := ConvertMarkdownToQQ(rm.Content)
		runes := []rune(content)
		if len(runes) > 200 {
			content = string(runes[:200])
		}
		msgs = append(msgs, bridge.PlatformMessage{
			Content:    content,
			ReplyTo:    response.ReplyTo,
			FormatMode: "text",
		})
	}
	return msgs, nil
}

// SendMessage sends a message through the QQ WebSocket connection.
func (a *Adapter) SendMessage(msgType string, userID, groupID int64, content string) error {
	a.connMu.Lock()
	conn := a.conn
	a.connMu.Unlock()
	if conn == nil {
		return fmt.Errorf("QQ bot not connected")
	}

	req := OBv11SendMsg{
		Action: "send_msg",
		Params: OBv11Params{
			MessageType: msgType,
			UserID:      userID,
			GroupID:     groupID,
			Message:     content,
		},
		Echo: fmt.Sprintf("echo_%d", time.Now().UnixNano()),
	}

	data, _ := json.Marshal(req)
	return conn.WriteMessage(websocket.TextMessage, data)
}

// ReadMessages reads OBv11 messages from the WebSocket connection.
func (a *Adapter) ReadMessages(ctx context.Context, msgCh chan<- *OBv11Message) {
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}

		a.connMu.Lock()
		conn := a.conn
		a.connMu.Unlock()
		if conn == nil {
			// Client mode: auto-reconnect when connection is lost.
			if a.mode == "client" {
				if err := a.connectClient(ctx); err != nil {
					fmt.Printf("[qq] reconnect failed: %v, retrying in 3s...\n", err)
					time.Sleep(3 * time.Second)
				}
				continue
			}
			time.Sleep(time.Second)
			continue
		}

		_, raw, err := conn.ReadMessage()
		if err != nil {
			fmt.Printf("[qq] read error: %v\n", err)
			a.connMu.Lock()
			if a.conn != nil {
				a.conn.Close()
				a.conn = nil
			}
			a.connected = false
			a.connMu.Unlock()
			time.Sleep(time.Second)
			continue
		}

		// Try to parse as OBv11 message (event from QQ).
		var msg OBv11Message
		if err := json.Unmarshal(raw, &msg); err != nil {
			var resp OBv11APIResponse
			if err := json.Unmarshal(raw, &resp); err != nil {
				fmt.Printf("[qq] unknown message: %s\n", string(raw))
				continue
			}
			if resp.Echo != "" {
				a.respMu.Lock()
				if ch, ok := a.pendingResponses[resp.Echo]; ok {
					ch <- &resp
					delete(a.pendingResponses, resp.Echo)
				}
				a.respMu.Unlock()
			}
			continue
		}

		// Capture bot's own QQ number from incoming messages.
		if msg.SelfID != 0 && a.selfID == "" {
			a.selfID = fmt.Sprintf("%d", msg.SelfID)
			fmt.Printf("[qq:%s] self ID captured: %s\n", a.configName, a.selfID)
		}

		if msg.PostType == "message" {
			select {
			case msgCh <- &msg:
			case <-ctx.Done():
				return
			}
		}
	}
}

// extractText retrieves plain text from an OBv11 message.
func extractText(msg *OBv11Message) string {
	if msg.RawMessage != "" {
		return msg.RawMessage
	}
	switch m := msg.Message.(type) {
	case string:
		return m
	case []interface{}:
		var text string
		for _, seg := range m {
			if s, ok := seg.(map[string]interface{}); ok {
				if s["type"] == "text" {
					if data, ok := s["data"].(map[string]interface{}); ok {
						if t, ok := data["text"].(string); ok {
							text += t
						}
					}
				}
			}
		}
		return text
	}
	return ""
}

var cqAtRegex = regexp.MustCompile(`\[CQ:at,qq=(\d+)\]`)
var cqImageRegex = regexp.MustCompile(`\[CQ:image,[^\]]*\]`)
var cqURLRegex = regexp.MustCompile(`\burl=([^,\]]+)`)
	var boldRegex = regexp.MustCompile(`\*\*(.+?)\*\*`)
	var italicRegex = regexp.MustCompile(`\*(.+?)\*`)
	var strikethroughRegex = regexp.MustCompile(`~~(.+?)~~`)

// extractAttachments extracts image URLs from OBv11Message.
// Handles both string format (CQ codes in raw_message) and array format (parsed segments).
func extractAttachments(msg *OBv11Message) []bridge.Attachment {
	var attachments []bridge.Attachment

	// Array format: iterate segments looking for type="image".
	if segments, ok := msg.Message.([]interface{}); ok {
		for _, s := range segments {
			if seg, ok := s.(map[string]interface{}); ok {
				if seg["type"] != "image" {
					continue
				}
				data, _ := seg["data"].(map[string]interface{})
				if data == nil {
					continue
				}
				url, _ := data["url"].(string)
				file, _ := data["file"].(string)
				if url == "" {
					continue
				}
				attachments = append(attachments, bridge.Attachment{
					Type:     "image",
					URL:      url,
					FileName: file,
				})
			}
		}
		return attachments
	}

	// String format: parse CQ codes from RawMessage or string Message.
	raw := msg.RawMessage
	if raw == "" {
		if s, ok := msg.Message.(string); ok {
			raw = s
		}
	}
	matches := cqImageRegex.FindAllString(raw, -1)
	for _, m := range matches {
		urlMatch := cqURLRegex.FindStringSubmatch(m)
		if len(urlMatch) >= 2 {
			attachments = append(attachments, bridge.Attachment{
				Type: "image",
				URL:  urlMatch[1],
			})
		}
	}
	return attachments
}

// ConvertMarkdownToQQ converts markdown to QQ plain-text format.
func ConvertMarkdownToQQ(md string) string {
	md = boldRegex.ReplaceAllString(md, "$1")
	md = italicRegex.ReplaceAllString(md, "$1")
	md = strikethroughRegex.ReplaceAllString(md, "$1")
	md = removeHeadings(md)
	md = removeCodeBlocks(md)
	return md
}

func removeHeadings(s string) string {
	for _, line := range splitLines(s) {
		if len(line) > 0 && line[0] == '#' {
			s = replaceLine(s, line, stripPrefix(line, "# "))
		}
	}
	return s
}

func removeCodeBlocks(s string) string {
	result := ""
	inCode := false
	for _, line := range splitLines(s) {
		if hasPrefix(line, "```") {
			inCode = !inCode
			continue
		}
		if inCode {
			result += line + "\n"
		} else {
			result += line + "\n"
		}
	}
	return result
}

func splitLines(s string) []string {
	var lines []string
	start := 0
	for i, c := range s {
		if c == '\n' {
			lines = append(lines, s[start:i])
			start = i + 1
		}
	}
	if start < len(s) {
		lines = append(lines, s[start:])
	}
	return lines
}

func hasPrefix(s, prefix string) bool {
	return len(s) >= len(prefix) && s[:len(prefix)] == prefix
}

func stripPrefix(s, prefix string) string {
	if len(s) >= len(prefix) && s[:len(prefix)] == prefix {
		return s[len(prefix):]
	}
	return s
}

func replaceLine(s, old, new string) string {
	idx := indexOf(s, old)
	if idx < 0 {
		return s
	}
	return s[:idx] + new + s[idx+len(old):]
}

func indexOf(s, sub string) int {
	for i := 0; i <= len(s)-len(sub); i++ {
		if s[i:i+len(sub)] == sub {
			return i
		}
	}
	return -1
}