ragflow/internal/channels/wecom.go

//
//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//

package channels

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gorilla/websocket"

	"ragflow/internal/channels/core"
)

const (
	defaultWeComAPIBaseURL  = "https://qyapi.weixin.qq.com/cgi-bin"
	defaultWeComWSURL       = "wss://openws.work.weixin.qq.com"
	defaultWeComWebhookHost = "0.0.0.0"
	defaultWeComWebhookPort = 3002
	weComHTTPTimeout        = 30 * time.Second
	weComWSHandshakeTimeout = 10 * time.Second
	weComSubscribeTimeout   = 10 * time.Second
	weComWSWriteTimeout     = 10 * time.Second
	weComReconnectDelay     = 3 * time.Second
	weComHeartbeatInterval  = 25 * time.Second
	weComTokenSafetyMargin  = 60 * time.Second
)

var (
	errWeComAuthentication = errors.New("WeCom websocket authentication failed")
	errWeComDisconnected   = errors.New("WeCom websocket disconnected by server event")
	weComRequestSequence   atomic.Uint64
)

type wecomAccount struct {
	AccountID      string
	ConnectionType string
	CorpID         string
	AgentID        int64
	Secret         string
	Token          string
	AESKey         string
	BotID          string
	WebhookHost    string
	WebhookPort    int
	APIBaseURL     string
	WSURL          string
}

type wecomChannel struct {
	account wecomAccount
	crypto  *wecomCrypto
	client  *http.Client

	lifecycleMu sync.Mutex
	mu          sync.Mutex
	handler     core.MessageHandler
	cancel      context.CancelFunc
	done        chan struct{}
	ws          *websocket.Conn
	server      *wecomWebhookServer
	wsWriteMu   sync.Mutex
	tokenMu     sync.Mutex
	accessToken string
	tokenExpiry time.Time
}

type wecomAPIResponse struct {
	ErrCode int    `json:"errcode"`
	ErrMsg  string `json:"errmsg"`
}

type wecomAccessTokenResponse struct {
	wecomAPIResponse
	AccessToken string `json:"access_token"`
	ExpiresIn   int    `json:"expires_in"`
}

// newWeComChannel creates a WeCom channel with default endpoints filled in.
func newWeComChannel(account wecomAccount) (*wecomChannel, error) {
	account.ConnectionType = strings.ToLower(strings.TrimSpace(account.ConnectionType))
	if account.ConnectionType == "" {
		account.ConnectionType = "webhook"
	}
	account.APIBaseURL = strings.TrimRight(strings.TrimSpace(account.APIBaseURL), "/")
	if account.APIBaseURL == "" {
		account.APIBaseURL = defaultWeComAPIBaseURL
	}
	account.WSURL = strings.TrimSpace(account.WSURL)
	if account.WSURL == "" {
		account.WSURL = defaultWeComWSURL
	}
	account.WebhookHost = strings.TrimSpace(account.WebhookHost)
	if account.WebhookHost == "" {
		account.WebhookHost = defaultWeComWebhookHost
	}
	if account.WebhookPort == 0 {
		account.WebhookPort = defaultWeComWebhookPort
	}

	channel := &wecomChannel{
		account: account,
		client:  &http.Client{Timeout: weComHTTPTimeout},
	}
	if account.ConnectionType == "webhook" {
		crypto, err := newWeComCrypto(account.Token, account.AESKey, account.CorpID)
		if err != nil {
			return nil, err
		}
		channel.crypto = crypto
	}
	return channel, nil
}

// newWeComChannelFromConfig builds a WeCom channel from chat_channel.config.credential.
func newWeComChannelFromConfig(accountID string, cfg map[string]any) (*wecomChannel, error) {
	connectionType := strings.ToLower(firstString(cfg, "connection_type"))
	if connectionType == "" {
		connectionType = "webhook"
	}
	if connectionType != "websocket" {
		// Python treats every non-websocket value as webhook. Keep the same
		// effective runtime path while still applying webhook validation.
		connectionType = "webhook"
	}

	account := wecomAccount{
		AccountID:      accountID,
		ConnectionType: connectionType,
		Secret:         firstString(cfg, "secret"),
		WebhookHost:    firstString(cfg, "webhook_host"),
		APIBaseURL:     firstString(cfg, "api_base_url"),
		WSURL:          firstString(cfg, "ws_url"),
	}
	if account.Secret == "" {
		return nil, fmt.Errorf("wecom account %q missing required field: secret", accountID)
	}

	if connectionType == "websocket" {
		account.BotID = firstString(cfg, "bot_id")
		if account.BotID == "" {
			return nil, fmt.Errorf("wecom account %q missing required field: bot_id", accountID)
		}
		return newWeComChannel(account)
	}

	account.CorpID = firstString(cfg, "corp_id")
	account.Token = firstString(cfg, "token")
	account.AESKey = firstString(cfg, "aes_key")
	missing := make([]string, 0, 3)
	if account.CorpID == "" {
		missing = append(missing, "corp_id")
	}
	if account.Token == "" {
		missing = append(missing, "token")
	}
	if account.AESKey == "" {
		missing = append(missing, "aes_key")
	}
	if len(missing) > 0 {
		return nil, fmt.Errorf("wecom account %q missing required fields: %s", accountID, strings.Join(missing, ", "))
	}
	if len(account.AESKey) != 43 {
		return nil, fmt.Errorf("wecom account %q aes_key (EncodingAESKey) must be 43 characters, got %d", accountID, len(account.AESKey))
	}

	agentID, err := parseWeComInt(cfg["agent_id"], "agent_id")
	if err != nil {
		return nil, fmt.Errorf("wecom account %q %w", accountID, err)
	}
	account.AgentID = agentID
	var configuredWebhookPort *int
	if raw, ok := cfg["webhook_port"]; ok {
		port, err := parseWeComInt(raw, "webhook_port")
		if err != nil {
			return nil, fmt.Errorf("wecom account %q %w", accountID, err)
		}
		account.WebhookPort = int(port)
		configuredWebhookPort = &account.WebhookPort
	}
	channel, err := newWeComChannel(account)
	if err != nil {
		return nil, err
	}
	if configuredWebhookPort != nil {
		channel.account.WebhookPort = *configuredWebhookPort
	}
	return channel, nil
}

// ChannelID returns the platform identifier used by the chat-channel runtime.
func (c *wecomChannel) ChannelID() string {
	return "wecom"
}

// AccountID returns the chat_channel.id bound to this WeCom runtime instance.
func (c *wecomChannel) AccountID() string {
	return c.account.AccountID
}

// SetMessageHandler installs the RAGFlow bridge invoked for inbound messages.
func (c *wecomChannel) SetMessageHandler(handler core.MessageHandler) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.handler = handler
}

// Start starts either the shared webhook listener or the WebSocket client.
func (c *wecomChannel) Start(ctx context.Context) error {
	c.lifecycleMu.Lock()
	defer c.lifecycleMu.Unlock()

	if c.account.ConnectionType == "webhook" {
		c.mu.Lock()
		started := c.server != nil
		c.mu.Unlock()
		if started {
			return nil
		}
		server, err := acquireWeComWebhookServer(c)
		if err != nil {
			return err
		}
		c.mu.Lock()
		c.server = server
		c.mu.Unlock()
		log.Printf("[wecom:%s] registered webhook at /wecom/%s/callback", c.account.AccountID, c.account.AccountID)
		return nil
	}

	c.mu.Lock()
	if c.cancel != nil {
		c.mu.Unlock()
		return nil
	}
	runCtx, cancel := context.WithCancel(ctx)
	done := make(chan struct{})
	c.cancel = cancel
	c.done = done
	c.mu.Unlock()

	go func() {
		defer close(done)
		c.runWebSocket(runCtx)
		c.mu.Lock()
		if c.done == done {
			c.cancel = nil
			c.ws = nil
		}
		c.mu.Unlock()
	}()
	return nil
}

// Stop releases the webhook listener or stops the WebSocket reconnect loop.
func (c *wecomChannel) Stop(ctx context.Context) error {
	c.lifecycleMu.Lock()
	defer c.lifecycleMu.Unlock()

	if c.account.ConnectionType == "webhook" {
		c.mu.Lock()
		server := c.server
		c.server = nil
		c.mu.Unlock()
		if server != nil {
			return releaseWeComWebhookServer(ctx, server, c.account.AccountID)
		}
		return nil
	}

	c.mu.Lock()
	cancel := c.cancel
	done := c.done
	conn := c.ws
	c.mu.Unlock()
	if cancel == nil {
		return nil
	}
	cancel()
	if conn != nil {
		_ = conn.Close()
	}
	select {
	case <-done:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// Send sends an application message or a WebSocket bot message.
func (c *wecomChannel) Send(ctx context.Context, msg core.OutgoingMessage) error {
	if c.account.ConnectionType == "websocket" {
		if err := c.sendWebSocketMessage(msg); err != nil {
			log.Printf("[wecom:%s] websocket send failed: %v", c.account.AccountID, err)
		}
		return nil
	}
	if err := c.sendApplicationMessage(ctx, msg); err != nil {
		log.Printf("[wecom:%s] application send failed: %v", c.account.AccountID, err)
	}
	return nil
}

func (c *wecomChannel) runWebSocket(ctx context.Context) {
	for ctx.Err() == nil {
		err := c.runWebSocketOnce(ctx)
		if ctx.Err() != nil {
			return
		}
		if errors.Is(err, errWeComAuthentication) || errors.Is(err, errWeComDisconnected) {
			log.Printf("[wecom:%s] websocket stopped: %v", c.account.AccountID, err)
			return
		}
		if err != nil {
			log.Printf("[wecom:%s] websocket loop error: %v", c.account.AccountID, err)
		}
		if !waitForWeComRetry(ctx) {
			return
		}
	}
}

func (c *wecomChannel) runWebSocketOnce(ctx context.Context) error {
	dialer := *websocket.DefaultDialer
	dialer.HandshakeTimeout = weComWSHandshakeTimeout
	conn, _, err := dialer.DialContext(ctx, c.account.WSURL, nil)
	if err != nil {
		return err
	}
	defer conn.Close()

	c.mu.Lock()
	c.ws = conn
	c.mu.Unlock()
	defer func() {
		c.mu.Lock()
		if c.ws == conn {
			c.ws = nil
		}
		c.mu.Unlock()
	}()

	connectionDone := make(chan struct{})
	defer close(connectionDone)
	go func() {
		select {
		case <-ctx.Done():
			_ = conn.Close()
		case <-connectionDone:
		}
	}()

	if err := c.subscribeWebSocket(conn); err != nil {
		return err
	}
	heartbeatDone := make(chan struct{})
	defer close(heartbeatDone)
	go c.runWebSocketHeartbeat(ctx, conn, heartbeatDone)

	for ctx.Err() == nil {
		if err := conn.SetReadDeadline(time.Now().Add(2 * weComHeartbeatInterval)); err != nil {
			return err
		}
		_, payload, err := conn.ReadMessage()
		if err != nil {
			return err
		}
		disconnect, err := c.handleWebSocketPayload(ctx, payload)
		if err != nil {
			log.Printf("[wecom:%s] invalid websocket payload: %v", c.account.AccountID, err)
			continue
		}
		if disconnect {
			return errWeComDisconnected
		}
	}
	return ctx.Err()
}

func (c *wecomChannel) subscribeWebSocket(conn *websocket.Conn) error {
	payload := map[string]any{
		"cmd":     "aibot_subscribe",
		"headers": map[string]any{"req_id": newWeComRequestID()},
		"body": map[string]any{
			"bot_id": c.account.BotID,
			"secret": c.account.Secret,
		},
	}
	if err := c.writeWebSocketJSON(conn, payload); err != nil {
		return err
	}
	_ = conn.SetReadDeadline(time.Now().Add(weComSubscribeTimeout))
	defer conn.SetReadDeadline(time.Time{})
	var response map[string]any
	if err := conn.ReadJSON(&response); err != nil {
		return fmt.Errorf("WeCom websocket subscribe acknowledgement: %w", err)
	}
	if cmd := weComRawStringValue(response["cmd"]); cmd != "aibot_subscribe" {
		log.Printf("[wecom:%s] unexpected subscribe response command %q", c.account.AccountID, cmd)
	}
	errCode := weComIntValue(response["errcode"])
	if errCode == 853000 {
		return fmt.Errorf("%w: invalid bot_id or secret", errWeComAuthentication)
	}
	if errCode != 0 {
		return fmt.Errorf("WeCom websocket subscribe failed: errcode=%d errmsg=%s", errCode, weComStringValue(response["errmsg"]))
	}
	log.Printf("[wecom:%s] websocket subscribed", c.account.AccountID)
	return nil
}

func (c *wecomChannel) runWebSocketHeartbeat(ctx context.Context, conn *websocket.Conn, done <-chan struct{}) {
	ticker := time.NewTicker(weComHeartbeatInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-done:
			return
		case <-ticker.C:
			if err := c.writeWebSocketJSON(conn, map[string]any{
				"cmd":     "ping",
				"headers": map[string]any{"req_id": newWeComRequestID()},
			}); err != nil {
				_ = conn.Close()
				return
			}
		}
	}
}

func (c *wecomChannel) handleWebSocketPayload(ctx context.Context, payload []byte) (bool, error) {
	var object map[string]any
	if err := json.Unmarshal(payload, &object); err != nil {
		return false, err
	}
	cmd := weComRawStringValue(object["cmd"])
	switch cmd {
	case "aibot_msg_callback":
		c.handleWebSocketMessage(ctx, object)
	case "aibot_event_callback":
		body := weComMapValue(object["body"])
		event := weComMapValue(body["event"])
		if weComRawStringValue(event["eventtype"]) == "disconnected_event" {
			return true, nil
		}
	case "aibot_subscribe", "aibot_respond_msg", "aibot_respond_welcome_msg", "aibot_send_msg", "ping":
		if code := weComIntValue(object["errcode"]); code != 0 {
			log.Printf("[wecom:%s] websocket response error: %s", c.account.AccountID, string(payload))
		}
	}
	return false, nil
}

func (c *wecomChannel) handleWebSocketMessage(ctx context.Context, raw map[string]any) {
	headers := weComMapValue(raw["headers"])
	body := weComMapValue(raw["body"])
	if weComRawStringValue(body["msgtype"]) != "text" {
		return
	}
	sender := weComMapValue(body["from"])
	senderID := weComRawStringValue(sender["userid"])
	if senderID == "" {
		return
	}
	chatID := weComRawStringValue(body["chatid"])
	if chatID == "" {
		chatID = senderID
	}
	messageID := weComRawStringValue(headers["req_id"])
	if messageID == "" {
		messageID = weComRawStringValue(body["msgid"])
	}
	chatType := "p2p"
	if weComRawStringValue(body["chattype"]) == "group" {
		chatType = "group"
	}
	text := weComRawStringValue(weComMapValue(body["text"])["content"])
	c.handleTextMessage(ctx, chatID, senderID, messageID, text, chatType, raw)
}

func (c *wecomChannel) sendWebSocketMessage(msg core.OutgoingMessage) error {
	if msg.Text == "" {
		return errors.New("WeCom websocket reply text is required")
	}
	c.mu.Lock()
	conn := c.ws
	c.mu.Unlock()
	if conn == nil {
		return errors.New("WeCom websocket is not connected")
	}
	return c.writeWebSocketJSON(conn, map[string]any{
		"cmd":     "aibot_send_msg",
		"headers": map[string]any{"req_id": newWeComRequestID()},
		"body": map[string]any{
			"chatid":  msg.ChatID,
			"msgtype": "markdown",
			"markdown": map[string]any{
				"content": msg.Text,
			},
		},
	})
}

func (c *wecomChannel) writeWebSocketJSON(conn *websocket.Conn, payload any) error {
	c.wsWriteMu.Lock()
	defer c.wsWriteMu.Unlock()
	if err := conn.SetWriteDeadline(time.Now().Add(weComWSWriteTimeout)); err != nil {
		return err
	}
	return conn.WriteJSON(payload)
}

func (c *wecomChannel) sendApplicationMessage(ctx context.Context, msg core.OutgoingMessage) error {
	if msg.ChatID == "" {
		return errors.New("WeCom chat_id is required")
	}
	accessToken, err := c.getAccessToken(ctx)
	if err != nil {
		return err
	}
	endpoint, err := url.Parse(c.account.APIBaseURL + "/message/send")
	if err != nil {
		return err
	}
	query := endpoint.Query()
	query.Set("access_token", accessToken)
	endpoint.RawQuery = query.Encode()

	var response wecomAPIResponse
	if err := c.requestJSON(ctx, http.MethodPost, endpoint.String(), map[string]any{
		"touser":  msg.ChatID,
		"msgtype": "text",
		"agentid": c.account.AgentID,
		"text":    map[string]any{"content": msg.Text},
		"safe":    0,
	}, &response); err != nil {
		return err
	}
	if response.ErrCode != 0 {
		if response.ErrCode == 40014 || response.ErrCode == 42001 {
			c.clearAccessToken()
		}
		return fmt.Errorf("WeCom message/send failed: errcode=%d errmsg=%s", response.ErrCode, response.ErrMsg)
	}
	return nil
}

func (c *wecomChannel) getAccessToken(ctx context.Context) (string, error) {
	c.tokenMu.Lock()
	defer c.tokenMu.Unlock()
	if c.accessToken != "" && time.Now().Before(c.tokenExpiry) {
		return c.accessToken, nil
	}

	endpoint, err := url.Parse(c.account.APIBaseURL + "/gettoken")
	if err != nil {
		return "", err
	}
	query := endpoint.Query()
	query.Set("corpid", c.account.CorpID)
	query.Set("corpsecret", c.account.Secret)
	endpoint.RawQuery = query.Encode()
	var response wecomAccessTokenResponse
	if err := c.requestJSON(ctx, http.MethodGet, endpoint.String(), nil, &response); err != nil {
		return "", err
	}
	if response.ErrCode != 0 || response.AccessToken == "" {
		return "", fmt.Errorf("WeCom gettoken failed: errcode=%d errmsg=%s", response.ErrCode, response.ErrMsg)
	}
	expiresIn := response.ExpiresIn
	if expiresIn == 0 {
		expiresIn = 7200
	}
	c.accessToken = response.AccessToken
	c.tokenExpiry = time.Now().Add(time.Duration(expiresIn)*time.Second - weComTokenSafetyMargin)
	return c.accessToken, nil
}

func (c *wecomChannel) clearAccessToken() {
	c.tokenMu.Lock()
	defer c.tokenMu.Unlock()
	c.accessToken = ""
	c.tokenExpiry = time.Time{}
}

func (c *wecomChannel) requestJSON(ctx context.Context, method, endpoint string, body map[string]any, out any) error {
	var reader io.Reader
	if body != nil {
		payload, err := json.Marshal(body)
		if err != nil {
			return err
		}
		reader = bytes.NewReader(payload)
	}
	req, err := http.NewRequestWithContext(ctx, method, endpoint, reader)
	if err != nil {
		return err
	}
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	resp, err := c.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	payload, err := io.ReadAll(io.LimitReader(resp.Body, 1024*1024))
	if err != nil {
		return err
	}
	if out == nil {
		return nil
	}
	if err := json.Unmarshal(payload, out); err != nil {
		return fmt.Errorf("WeCom API returned invalid JSON: %w", err)
	}
	return nil
}

func (c *wecomChannel) handleTextMessage(ctx context.Context, chatID, senderID, messageID, text, chatType string, raw map[string]any) {
	if strings.TrimSpace(text) == "" {
		return
	}
	c.mu.Lock()
	handler := c.handler
	c.mu.Unlock()
	if handler == nil {
		return
	}
	if err := handler(ctx, core.IncomingMessage{
		Channel:   c.ChannelID(),
		AccountID: c.account.AccountID,
		ChatID:    chatID,
		ChatType:  chatType,
		MessageID: messageID,
		SenderID:  senderID,
		Text:      text,
		Raw:       raw,
	}); err != nil {
		log.Printf("[wecom:%s] message handler error: %v", c.account.AccountID, err)
	}
}

func parseWeComInt(value any, field string) (int64, error) {
	switch v := value.(type) {
	case int:
		return int64(v), nil
	case int64:
		return v, nil
	case float64:
		if v != float64(int64(v)) {
			return 0, fmt.Errorf("%s must be an integer", field)
		}
		return int64(v), nil
	case json.Number:
		result, err := v.Int64()
		if err != nil {
			return 0, fmt.Errorf("%s must be an integer", field)
		}
		return result, nil
	case string:
		result, err := strconv.ParseInt(strings.TrimSpace(v), 10, 64)
		if err != nil {
			return 0, fmt.Errorf("%s must be an integer", field)
		}
		return result, nil
	default:
		return 0, fmt.Errorf("%s must be an integer", field)
	}
}

func weComMapValue(value any) map[string]any {
	result, _ := value.(map[string]any)
	return result
}

func weComStringValue(value any) string {
	if value == nil {
		return ""
	}
	return strings.TrimSpace(fmt.Sprint(value))
}

// weComRawStringValue mirrors Python's str(value or "") conversion for
// message fields without changing user-visible whitespace.
func weComRawStringValue(value any) string {
	if value == nil {
		return ""
	}
	if text, ok := value.(string); ok {
		return text
	}
	return fmt.Sprint(value)
}

func weComIntValue(value any) int {
	switch v := value.(type) {
	case float64:
		return int(v)
	case int:
		return v
	case json.Number:
		result, _ := v.Int64()
		return int(result)
	case string:
		result, _ := strconv.Atoi(v)
		return result
	default:
		return 0
	}
}

func newWeComRequestID() string {
	return fmt.Sprintf("req-%d-%d", time.Now().UnixNano(), weComRequestSequence.Add(1))
}

func waitForWeComRetry(ctx context.Context) bool {
	timer := time.NewTimer(weComReconnectDelay)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return false
	case <-timer.C:
		return true
	}
}