kami_gateway/internal/utils/proxy_pool.go

package utils

import (
	"context"
	"errors"
	"fmt"
	"github.com/dubonzi/otelresty"
	"github.com/go-resty/resty/v2"
	"maps"
	"net/http"
	"net/url"
	"slices"
	"strings"
	"sync"
	"time"

	"gateway/internal/cache"
	"gateway/internal/config"
	"gateway/internal/otelTrace"

	"github.com/duke-git/lancet/v2/slice"
	"github.com/duke-git/lancet/v2/strutil"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"
	"go.uber.org/zap"
)

// ProxyStrategy 代理策略接口
type ProxyStrategy interface {
	GetProxy(ctx context.Context, proxyRequest ProxyRequest) (string, error)
	GetHTTPClient(ctx context.Context, proxyRequest ProxyRequest) (*http.Client, error)
}

// ProxyInfo 代理信息
type ProxyInfo struct {
	proxy      string
	expireAt   time.Time
	authKey    string
	authPwd    string
	channel    string // 添加通道字段
	orderID    string
	lastUsedAt time.Time // 添加最后使用时间字段
	usageCount int       // 添加使用次数字段
}

// ProxyRequest 代理请求信息
type ProxyRequest struct {
	OrderNo    string
	Channel    string // 添加通道字段
	OrderPerIP int
}

// DefaultProxyStrategy 默认代理策略
type DefaultProxyStrategy struct {
	proxyURL string
	authKey  string
	authPwd  string
	mu       sync.RWMutex
	current  *ProxyInfo
}

// OrderBasedProxyStrategy 基于订单号和通道的代理策略
type OrderBasedProxyStrategy struct {
	Id           string
	authKey      string
	proxyURL     string
	authPwd      string
	OrderPerIP   int
	mu           sync.RWMutex
	totalProxies []*ProxyInfo
	proxies      map[string]*ProxyInfo // key: channel_orderID
	stopCh       chan struct{}         // 用于停止清理协程
}

// NewDefaultProxyStrategy 创建默认代理策略
func NewDefaultProxyStrategy(proxyURL, authKey, authPwd string) *DefaultProxyStrategy {
	return &DefaultProxyStrategy{
		proxyURL: proxyURL,
		authKey:  authKey,
		authPwd:  authPwd,
	}
}

// NewOrderBasedProxyStrategy 创建基于订单号和通道的代理策略
func NewOrderBasedProxyStrategy(proxyURL, authKey, authPwd string) *OrderBasedProxyStrategy {
	strategy := &OrderBasedProxyStrategy{
		proxyURL:     proxyURL,
		authKey:      authKey,
		authPwd:      authPwd,
		totalProxies: make([]*ProxyInfo, 0),
		proxies:      make(map[string]*ProxyInfo),
		stopCh:       make(chan struct{}),
	}

	// 启动清理协程
	go strategy.startCleanupRoutine()

	return strategy
}

// GetProxy 获取代理（默认策略）
func (p *DefaultProxyStrategy) GetProxy(ctx context.Context, _ ProxyRequest) (string, error) {
	proxy, err := p.innerGetProxy(ctx, ProxyRequest{})
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("http://%s:%s@%s", p.authKey, p.authPwd, proxy), nil
}

// GetProxy 获取代理（默认策略）
func (p *DefaultProxyStrategy) innerGetProxy(ctx context.Context, _ ProxyRequest) (string, error) {
	p.mu.RLock()
	if p.current != nil && time.Now().Before(p.current.expireAt) {
		if err := p.checkProxyAvailable(ctx, p.current.proxy); err == nil {
			proxy := p.current.proxy
			p.mu.RUnlock()
			otelTrace.Logger.WithContext(ctx).Info("使用缓存代理", zap.String("proxy", proxy))
			return proxy, nil
		}
	}
	p.mu.RUnlock()
	return p.getNewProxy(ctx)
}

// GetProxy 获取代理（基于订单号和通道策略）
func (p *OrderBasedProxyStrategy) GetProxy(ctx context.Context, proxyRequest ProxyRequest) (string, error) {
	proxy, err := p.getInnerProxy(ctx, proxyRequest)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("socks5://%s:%s@%s", p.authKey, p.authPwd, proxy), nil
}

// GetProxy 获取代理（基于订单号和通道策略）
func (p *OrderBasedProxyStrategy) getInnerProxy(ctx context.Context, proxyRequest ProxyRequest) (string, error) {
	if proxyRequest.OrderPerIP == 0 {
		proxyRequest.OrderPerIP = 1
	}

	p.mu.RLock()
	// 生成代理缓存键：channel_orderID
	cacheKey := fmt.Sprintf("%s:%s", proxyRequest.Channel, proxyRequest.OrderNo)

	if proxy, exists := p.proxies[cacheKey]; exists {
		p.mu.RUnlock()
		if err := p.checkProxyAvailable(ctx, proxy.proxy); err == nil {
			// 更新最后使用时间和使用次数
			p.mu.Lock()
			proxy.lastUsedAt = time.Now()
			proxy.usageCount++
			p.mu.Unlock()

			otelTrace.Logger.WithContext(ctx).Info("使用订单缓存代理",
				zap.String("proxy", proxy.proxy),
				zap.String("orderID", proxyRequest.OrderNo),
				zap.String("channel", proxyRequest.Channel),
				zap.Int("usageCount", proxy.usageCount))
			return proxy.proxy, nil
		}
	} else {
		p.mu.RUnlock()
	}
	return p.getNewProxy(ctx, proxyRequest)
}

// getNewProxy 获取新代理（默认策略）
func (p *DefaultProxyStrategy) getNewProxy(ctx context.Context) (string, error) {
	var lastErr error
	for i := range 3 {
		proxy, err := p.tryGetProxy(ctx)
		if err == nil {
			p.mu.Lock()
			p.current = &ProxyInfo{
				proxy:    proxy,
				expireAt: time.Now().Add(50 * time.Second),
				authKey:  p.authKey,
				authPwd:  p.authPwd,
			}
			p.mu.Unlock()
			otelTrace.Logger.WithContext(ctx).Info("获取新代理", zap.String("proxy", proxy))
			return proxy, nil
		}
		lastErr = err
		otelTrace.Logger.WithContext(ctx).Warn("获取代理失败，准备重试",
			zap.Int("重试次数", i+1),
			zap.Error(err))

		select {
		case <-ctx.Done():
			return "", ctx.Err()
		case <-time.After(time.Second * time.Duration(i+1)):
			continue
		}
	}
	return "", fmt.Errorf("获取代理失败，已重试3次: %v", lastErr)
}

// 查找使用次数少于指定次数的 ip（跨通道查找）
func (p *OrderBasedProxyStrategy) getUnderLimitedProxy(ctx context.Context, perIP int, excludeChannels string) (string, error) {
	p.mu.RLock()
	//需要找到指定通道所有的代理
	counts := map[*ProxyInfo]int{}
	for channel, info := range p.proxies {
		if strings.Contains(channel, excludeChannels) && info.usageCount > perIP {
			counts[info]++
		}
	}
	proxies := slice.Filter(p.totalProxies, func(index int, item *ProxyInfo) bool {
		for info := range counts {
			if info.proxy == item.proxy {
				return false
			}
		}
		return true
	})
	p.mu.RUnlock()
	for _, proxy := range proxies {
		if err := p.checkProxyAvailable(ctx, proxy.proxy); err == nil {
			return proxy.proxy, nil
		} else {
			// 清理不可用的代理
			func() {
				p.mu.Lock()
				defer p.mu.Unlock()
				maps.DeleteFunc(p.proxies, func(key string, value *ProxyInfo) bool {
					return value.proxy == proxy.proxy
				})
				p.totalProxies = slices.DeleteFunc(p.totalProxies, func(info *ProxyInfo) bool {
					return info.proxy == proxy.proxy
				})
				otelTrace.Logger.WithContext(ctx).Warn("代理IP不可用", zap.String("proxy", proxy.proxy), zap.Error(err))
			}()
		}
	}
	return "", fmt.Errorf("没有找到使用次数少于指定次数的 ip")
}

// getNewProxy 获取新代理（基于订单号和通道策略）
func (p *OrderBasedProxyStrategy) getNewProxy(ctx context.Context, proxyRequest ProxyRequest) (string, error) {
	var lastErr error

	// 如果订单 ip设置使用次数大于1，则获取使用次数少于指定次数的 ip
	if proxyRequest.OrderPerIP >= 1 {
		// 排除当前通道，允许跨通道复用代理
		excludeChannels := proxyRequest.Channel
		proxy, err := p.getUnderLimitedProxy(ctx, proxyRequest.OrderPerIP, excludeChannels)
		if err == nil {
			cacheKey := fmt.Sprintf("%s:%s", proxyRequest.Channel, proxyRequest.OrderNo)
			p.mu.Lock()
			p.proxies[cacheKey] = &ProxyInfo{
				proxy:      proxy,
				expireAt:   time.Now().Add(50 * time.Second),
				authKey:    p.authKey,
				authPwd:    p.authPwd,
				channel:    proxyRequest.Channel,
				orderID:    proxyRequest.OrderNo,
				lastUsedAt: time.Now(),
				usageCount: 1,
			}
			p.mu.Unlock()
			return proxy, nil
		}
	}

	tmpProxies := make([]*ProxyInfo, 0)
	var err error

	for i := range 3 {
		proxies, err2 := p.tryGetProxy(ctx)
		if err2 != nil {
			otelTrace.Logger.WithContext(ctx).Warn("获取代理失败，准备重试",
				zap.Int("retryCount", i+1),
				zap.Error(err2))
			time.Sleep(time.Second * 2)
			lastErr = err2
			continue
		}

		for _, proxy := range proxies {
			p.mu.Lock()
			proxyInfo := &ProxyInfo{
				proxy:      proxy,
				expireAt:   time.Now().Add(50 * time.Second),
				authKey:    p.authKey,
				authPwd:    p.authPwd,
				channel:    proxyRequest.Channel,
				orderID:    proxyRequest.OrderNo,
				lastUsedAt: time.Now(),
				usageCount: 1,
			}
			tmpProxies = append(tmpProxies, proxyInfo)
			p.mu.Unlock()

			otelTrace.Logger.WithContext(ctx).Info("获取新代理",
				zap.String("proxy", proxy),
				zap.String("orderID", proxyRequest.OrderNo),
				zap.String("channel", proxyRequest.Channel))
		}
		lastErr = err
		break
	}

	p.totalProxies = append(p.totalProxies, tmpProxies...)

	for _, proxy := range tmpProxies {
		if err2 := p.checkProxyAvailable(ctx, proxy.proxy); err2 == nil {
			cacheKey := fmt.Sprintf("%s:%s", proxyRequest.Channel, proxyRequest.OrderNo)
			p.proxies[cacheKey] = proxy
			return proxy.proxy, nil
		}
	}

	return "", fmt.Errorf("获取代理失败，已重试2次: %v", lastErr)
}

// GetHTTPClient 获取HTTP客户端（默认策略）
func (p *DefaultProxyStrategy) GetHTTPClient(ctx context.Context, _ ProxyRequest) (*http.Client, error) {
	proxy, err := p.GetProxy(ctx, ProxyRequest{})
	if err != nil {
		return nil, err
	}
	return p.createHTTPClient(proxy)
}

// GetHTTPClient 获取HTTP客户端（基于订单号和通道策略）
func (p *OrderBasedProxyStrategy) GetHTTPClient(ctx context.Context, proxyRequest ProxyRequest) (*http.Client, error) {
	proxy, err := p.GetProxy(ctx, proxyRequest)
	if err != nil {
		return nil, err
	}
	return p.createHTTPClient(proxy)
}

// createHTTPClient 创建HTTP客户端
func (p *DefaultProxyStrategy) createHTTPClient(proxy string) (*http.Client, error) {
	proxyURL, err := url.Parse(proxy)
	if err != nil {
		return nil, fmt.Errorf("解析代理URL失败: %v", err)
	}

	transport := &http.Transport{
		Proxy: http.ProxyURL(proxyURL),
	}

	return &http.Client{
		Transport: transport,
		Timeout:   30 * time.Second,
	}, nil
}

// createHTTPClient 创建HTTP客户端
func (p *OrderBasedProxyStrategy) createHTTPClient(proxy string) (*http.Client, error) {
	proxyURL, err := url.Parse(proxy)
	if err != nil {
		return nil, fmt.Errorf("解析代理URL失败: %v", err)
	}

	return &http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyURL(proxyURL),
		},
		Timeout: 2 * time.Second,
	}, nil
}

// tryGetProxy 尝试获取代理IP
func (p *DefaultProxyStrategy) tryGetProxy(ctx context.Context) (string, error) {
	client := resty.New()
	otelresty.TraceClient(client)
	proxyURL, err := client.R().SetContext(ctx).Get(p.proxyURL)
	if err != nil {
		return "", fmt.Errorf("请求代理服务失败: %v", err)
	}
	return proxyURL.String(), nil
}

// tryGetProxy 尝试获取代理IP
func (p *OrderBasedProxyStrategy) tryGetProxy(ctx context.Context) ([]string, error) {
	client := resty.New()
	otelresty.TraceClient(client)
	proxyURL, err := client.R().SetContext(ctx).Get(p.proxyURL)
	if err != nil {
		return []string{}, fmt.Errorf("请求代理服务失败: %v", err)
	}
	proxyIPs := strutil.SplitAndTrim(proxyURL.String(), "\n")

	slice.ForEach(proxyIPs, func(index int, item string) {
		proxyIPs[index] = strutil.Trim(item)
	})
	return proxyIPs, nil
}

// checkProxyAvailable 检查代理IP是否可用
func (p *DefaultProxyStrategy) checkProxyAvailable(ctx context.Context, proxyIP string) error {
	client := resty.New()
	otelresty.TraceClient(client)
	client.SetProxy(fmt.Sprintf("socks5://%s:%s@%s", p.authKey, p.authPwd, proxyIP))
	response, err := client.R().SetContext(ctx).Get("https://www.qq.com")
	if err != nil {
		return fmt.Errorf("代理连接测试失败: %v", err)
	}
	if response.IsSuccess() {
		return nil
	}
	return errors.New("代理响应状态码异常")
}

// checkProxyAvailable 检查代理IP是否可用
func (p *OrderBasedProxyStrategy) checkProxyAvailable(ctx context.Context, proxyIP string) error {
	// 为代理检测设置独立的短超时，避免阻塞其他操作
	ctxWithTimeout, cancel := context.WithTimeout(ctx, 3*time.Second)
	defer cancel()

	client := resty.New()
	otelresty.TraceClient(client)
	client.SetTimeout(3 * time.Second) // 设置客户端超时
	client.SetProxy(fmt.Sprintf("socks5://%s:%s@%s", p.authKey, p.authPwd, proxyIP))

	response, err := client.R().SetContext(ctxWithTimeout).Get("https://www.qq.com")
	if err != nil {
		return fmt.Errorf("代理连接测试失败: %v", err)
	}
	if response.IsSuccess() {
		return nil
	}
	return errors.New("代理响应状态码异常")
}

// startCleanupRoutine 启动清理协程
func (p *OrderBasedProxyStrategy) startCleanupRoutine() {
	ticker := time.NewTicker(5 * time.Minute)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			p.cleanupUnusedProxies()
		case <-p.stopCh:
			return
		}
	}
}

// cleanupUnusedProxies 清理未使用的代理
func (p *OrderBasedProxyStrategy) cleanupUnusedProxies() {
	now := time.Now()

	// 分两阶段清理，减少锁持有时间

	// 第一阶段：收集需要删除的键（只持有读锁）
	var keysToDelete []string
	p.mu.RLock()
	for key, info := range p.proxies {
		if info.expireAt.Before(now) {
			keysToDelete = append(keysToDelete, key)
		}
	}
	p.mu.RUnlock()

	// 第二阶段：快速删除（只短暂持有写锁）
	if len(keysToDelete) > 0 {
		p.mu.Lock()
		for _, key := range keysToDelete {
			delete(p.proxies, key)
		}

		// 同时清理 totalProxies 中的过期代理
		p.totalProxies = slices.DeleteFunc(p.totalProxies, func(info *ProxyInfo) bool {
			return info.expireAt.Before(now)
		})
		p.mu.Unlock()
	}
}

// Stop 停止清理协程
func (p *OrderBasedProxyStrategy) Stop() {
	close(p.stopCh)
}

var OrderBasedProxyStrategyInstance *OrderBasedProxyStrategy
var DMProxyStrategyInstance *DMProxyStrategy

func StartProxyPool() {
	proxyConfig := config.GetProxyInfo()

	//// 初始化订单代理策略
	OrderBasedProxyStrategyInstance = NewOrderBasedProxyStrategy(
		proxyConfig.Url,
		proxyConfig.AuthKey,
		proxyConfig.AuthPwd,
	)

	// 初始化大漠代理策略
	DMProxyStrategyInstance = NewDMProxyStrategy(
		"http://need1.dmdaili.com:7771/dmgetip.asp?apikey=42bab0ac&pwd=1b567c2f286a08e391b5805565fa0882&getnum=20&httptype=1&geshi=&fenge=1&fengefu=&operate=all",
		cache.GetRedisClient().Client,
	)
}

func GetProxy(ctx context.Context, orderId string, channel string) (string, error) {
	ctx, span := otelTrace.Span(ctx, "GetProxy", "GetProxy", trace.WithAttributes(
		attribute.String("orderId", orderId),
		attribute.String("channel", channel),
	))
	defer span.End()

	// 为代理获取添加超时机制，防止在 OnBeforeRequest 中长时间阻塞
	ctxWithTimeout, cancel := context.WithTimeout(ctx, 8*time.Second)
	defer cancel()

	done := make(chan struct {
		proxy string
		err   error
	}, 1)

	go func() {
		proxyConfig := config.GetProxy()
		var proxy string
		var err error

		if proxyConfig == "dm" {
			proxy, err = DMProxyStrategyInstance.GetProxy(ctxWithTimeout, ProxyRequest{
				OrderNo:    orderId,
				Channel:    channel,
				OrderPerIP: 1,
			})
		} else {
			proxy, err = OrderBasedProxyStrategyInstance.GetProxy(ctxWithTimeout, ProxyRequest{
				OrderNo:    orderId,
				Channel:    channel,
				OrderPerIP: 1,
			})
		}

		select {
		case done <- struct {
			proxy string
			err   error
		}{proxy, err}:
		case <-ctxWithTimeout.Done():
		}
	}()

	select {
	case result := <-done:
		return result.proxy, result.err
	case <-ctxWithTimeout.Done():
		otelTrace.Logger.WithContext(ctx).Warn("获取代理超时",
			zap.String("orderId", orderId),
			zap.String("channel", channel),
			zap.Error(ctxWithTimeout.Err()))
		return "", fmt.Errorf("获取代理超时: %v", ctxWithTimeout.Err())
	}
}