WebHelper/api/socket.go

package api

import (
	"bufio"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"webhelper/global"

	"github.com/gorilla/websocket"
	"golang.org/x/net/proxy"
)

// WebSocketClient WebSocket 客户端
type WebSocketClient struct {
	conn        *websocket.Conn
	isConnected bool
	mu          sync.RWMutex
	onMessage   func(message string)
	onError     func(err error)
	onClose     func()
}

// WebSocketConfig WebSocket 配置
type WebSocketConfig struct {
	URL      string            `json:"url"`
	Headers  map[string]string `json:"headers"`
	Proxy    string            `json:"proxy"`    // 代理地址，支持 http:// 和 socks5://
	Insecure bool              `json:"insecure"`
}

// WebSocketMessage WebSocket 消息
type WebSocketMessage struct {
	Type    string `json:"type"` // text, binary, ping, pong, error, system
	Content string `json:"content"`
	Time    string `json:"time"`
}

var (
	wsClient     *WebSocketClient
	wsClientOnce sync.Once
)

// NewWebSocketClient 创建 WebSocket 客户端
func NewWebSocketClient() *WebSocketClient {
	return &WebSocketClient{
		isConnected: false,
	}
}

// Connect 连接 WebSocket
func (c *WebSocketClient) Connect(config *WebSocketConfig, onMessage func(string), onError func(error), onClose func()) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.isConnected {
		return fmt.Errorf("已连接")
	}

	// 解析 URL
	u, err := url.Parse(config.URL)
	if err != nil {
		return fmt.Errorf("URL 格式错误: %v", err)
	}

	// 创建 HTTP Header（只添加用户自定义头）
	header := http.Header{}
	for key, value := range config.Headers {
		header.Set(key, value)
	}

	// 创建 Dialer
	dialer := &websocket.Dialer{
		HandshakeTimeout: 10 * time.Second,
	}

	// 配置代理
	proxyAddr := global.NormalizeProxy(config.Proxy)
	if proxyAddr != "" {
		proxyURL, err := url.Parse(proxyAddr)
		if err != nil {
			return fmt.Errorf("代理地址格式错误: %v", err)
		}
		dialer.Proxy = http.ProxyURL(proxyURL)
	}

	// 连接
	global.Log.Infof("正在连接 WebSocket: %s (proxy=%s)", u.String(), proxyAddr)
	conn, _, err := dialer.Dial(u.String(), header)
	if err != nil {
		return fmt.Errorf("连接失败: %v", err)
	}

	c.conn = conn
	c.isConnected = true
	c.onMessage = onMessage
	c.onError = onError
	c.onClose = onClose

	global.Log.Infof("WebSocket 连接成功: %s", u.String())

	// 启动消息监听
	go c.readMessages()

	return nil
}

// readMessages 读取消息
func (c *WebSocketClient) readMessages() {
	defer func() {
		c.mu.Lock()
		c.isConnected = false
		c.mu.Unlock()
		if c.onClose != nil {
			c.onClose()
		}
	}()

	for {
		_, message, err := c.conn.ReadMessage()
		if err != nil {
			if c.onError != nil {
				c.onError(err)
			}
			return
		}
		if c.onMessage != nil {
			c.onMessage(string(message))
		}
	}
}

// Send 发送消息
func (c *WebSocketClient) Send(message string) error {
	c.mu.RLock()
	defer c.mu.RUnlock()

	if !c.isConnected || c.conn == nil {
		return fmt.Errorf("未连接")
	}

	return c.conn.WriteMessage(websocket.TextMessage, []byte(message))
}

// Close 关闭连接
func (c *WebSocketClient) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.conn != nil {
		err := c.conn.Close()
		c.isConnected = false
		return err
	}
	return nil
}

// IsConnected 是否已连接
func (c *WebSocketClient) IsConnected() bool {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.isConnected
}

// GetWebSocketClient 获取 WebSocket 客户端（并发安全的单例）
func GetWebSocketClient() *WebSocketClient {
	wsClientOnce.Do(func() {
		wsClient = NewWebSocketClient()
	})
	return wsClient
}

// ConnectWebSocket 连接 WebSocket
func ConnectWebSocket(config *WebSocketConfig, onMessage func(string), onError func(error), onClose func()) error {
	client := GetWebSocketClient()
	return client.Connect(config, onMessage, onError, onClose)
}

// SendWebSocketMessage 发送 WebSocket 消息
func SendWebSocketMessage(message string) error {
	client := GetWebSocketClient()
	return client.Send(message)
}

// CloseWebSocket 关闭 WebSocket 连接
func CloseWebSocket() error {
	client := GetWebSocketClient()
	return client.Close()
}

// IsWebSocketConnected WebSocket 是否已连接
func IsWebSocketConnected() bool {
	client := GetWebSocketClient()
	return client.IsConnected()
}

// TCPClient TCP 客户端
type TCPClient struct {
	conn        net.Conn
	isConnected bool
	mu          sync.RWMutex
	onMessage   func(message string)
	onError     func(err error)
	onClose     func()
}

// TCPConfig TCP 配置
type TCPConfig struct {
	Host  string `json:"host"`
	Port  int    `json:"port"`
	Proxy string `json:"proxy"` // 代理地址，支持 socks5://
}

// TCPMessage TCP 消息
type TCPMessage struct {
	Type    string `json:"type"`
	Content string `json:"content"`
	Time    string `json:"time"`
}

var (
	tcpClient     *TCPClient
	tcpClientOnce sync.Once
)

// NewTCPClient 创建 TCP 客户端
func NewTCPClient() *TCPClient {
	return &TCPClient{
		isConnected: false,
	}
}

// Connect 连接 TCP
func (c *TCPClient) Connect(config *TCPConfig, onMessage func(string), onError func(error), onClose func()) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.isConnected {
		return fmt.Errorf("已连接")
	}

	addr := fmt.Sprintf("%s:%d", config.Host, config.Port)

	// 配置代理
	proxyAddr := global.NormalizeProxy(config.Proxy)
	var conn net.Conn
	var err error

	if proxyAddr != "" {
		// 通过代理连接
		proxyURL, parseErr := url.Parse(proxyAddr)
		if parseErr != nil {
			return fmt.Errorf("代理地址格式错误: %v", parseErr)
		}

		switch proxyURL.Scheme {
		case "socks5", "socks4":
			// SOCKS 代理
			auth := &proxy.Auth{}
			if proxyURL.User != nil {
				auth.User = proxyURL.User.Username()
				if p, ok := proxyURL.User.Password(); ok {
					auth.Password = p
				}
			}
			dialer, proxyErr := proxy.SOCKS5("tcp", proxyURL.Host, auth, proxy.Direct)
			if proxyErr != nil {
				return fmt.Errorf("创建 SOCKS 代理失败: %v", proxyErr)
			}
			conn, err = dialer.Dial("tcp", addr)
		default:
			// HTTP 代理（通过 CONNECT 隧道）
			conn, err = httpConnectProxy(addr, proxyAddr)
		}
	} else {
		// 直连
		conn, err = net.DialTimeout("tcp", addr, 10*time.Second)
	}

	if err != nil {
		return fmt.Errorf("连接失败: %v", err)
	}

	c.conn = conn
	c.isConnected = true
	c.onMessage = onMessage
	c.onError = onError
	c.onClose = onClose

	global.Log.Infof("TCP 连接成功: %s (proxy=%s)", addr, proxyAddr)

	// 启动消息监听
	go c.readMessages()

	return nil
}

// httpConnectProxy 通过 HTTP CONNECT 代理建立 TCP 连接
func httpConnectProxy(target, proxyAddr string) (net.Conn, error) {
	proxyURL, err := url.Parse(proxyAddr)
	if err != nil {
		return nil, err
	}

	// 连接代理
	conn, err := net.DialTimeout("tcp", proxyURL.Host, 10*time.Second)
	if err != nil {
		return nil, err
	}

	// 发送 CONNECT 请求
	connectReq := fmt.Sprintf("CONNECT %s HTTP/1.1\r\nHost: %s\r\n\r\n", target, target)
	_, err = conn.Write([]byte(connectReq))
	if err != nil {
		conn.Close()
		return nil, err
	}

	// 读取响应（直到遇到空行）
	reader := bufio.NewReader(conn)
	line, err := reader.ReadString('\n')
	if err != nil {
		conn.Close()
		return nil, err
	}

	// 检查 200 状态
	if !strings.Contains(line, "200") {
		conn.Close()
		return nil, fmt.Errorf("代理连接失败: %s", strings.TrimSpace(line))
	}

	// 读取剩余响应头直到空行
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			break
		}
		if line == "\r\n" || line == "\n" {
			break
		}
	}

	return conn, nil
}

// tcpBufSize TCP 读取缓冲区大小，64KB 覆盖绝大多数调试场景
const tcpBufSize = 64 << 10

// readMessages 读取消息
func (c *TCPClient) readMessages() {
	defer func() {
		c.mu.Lock()
		c.isConnected = false
		c.mu.Unlock()
		if c.onClose != nil {
			c.onClose()
		}
	}()

	buf := make([]byte, tcpBufSize)
	for {
		n, err := c.conn.Read(buf)
		if err != nil {
			if c.onError != nil {
				c.onError(err)
			}
			return
		}
		if c.onMessage != nil {
			c.onMessage(string(buf[:n]))
		}
	}
}

// Send 发送消息
func (c *TCPClient) Send(message string) error {
	c.mu.RLock()
	defer c.mu.RUnlock()

	if !c.isConnected || c.conn == nil {
		return fmt.Errorf("未连接")
	}

	_, err := c.conn.Write([]byte(message))
	return err
}

// Close 关闭连接
func (c *TCPClient) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.conn != nil {
		err := c.conn.Close()
		c.isConnected = false
		return err
	}
	return nil
}

// IsConnected 是否已连接
func (c *TCPClient) IsConnected() bool {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.isConnected
}

// GetTCPClient 获取 TCP 客户端（并发安全的单例）
func GetTCPClient() *TCPClient {
	tcpClientOnce.Do(func() {
		tcpClient = NewTCPClient()
	})
	return tcpClient
}

// ConnectTCP 连接 TCP
func ConnectTCP(config *TCPConfig, onMessage func(string), onError func(error), onClose func()) error {
	client := GetTCPClient()
	return client.Connect(config, onMessage, onError, onClose)
}

// SendTCPMessage 发送 TCP 消息
func SendTCPMessage(message string) error {
	client := GetTCPClient()
	return client.Send(message)
}

// CloseTCP 关闭 TCP 连接
func CloseTCP() error {
	client := GetTCPClient()
	return client.Close()
}

// IsTCPConnected TCP 是否已连接
func IsTCPConnected() bool {
	client := GetTCPClient()
	return client.isConnected
}