common/config/config.go

package config

import (
	"errors"
	"genspark2api/common/env"
	"genspark2api/yescaptcha"
	"math/rand"
	"os"
	"strings"
	"sync"
	"time"
)

var ApiSecret = os.Getenv("API_SECRET")
var ApiSecrets = strings.Split(os.Getenv("API_SECRET"), ",")

var GSCookie = os.Getenv("GS_COOKIE")

//var GSCookies = strings.Split(os.Getenv("GS_COOKIE"), ",")

// var IpBlackList = os.Getenv("IP_BLACK_LIST")
var IpBlackList = strings.Split(os.Getenv("IP_BLACK_LIST"), ",")

var AutoDelChat = env.Int("AUTO_DEL_CHAT", 0)
var ProxyUrl = env.String("PROXY_URL", "")
var AutoModelChatMapType = env.Int("AUTO_MODEL_CHAT_MAP_TYPE", 1)
var YesCaptchaClientKey = env.String("YES_CAPTCHA_CLIENT_KEY", "")

// var CheatUrl = env.String("CHEAT_URL", "https://gs-cheat.aytsao.cn/genspark/create/req/body")
var RecaptchaProxyUrl = env.String("RECAPTCHA_PROXY_URL", "")

// 隐藏思考过程
var ReasoningHide = env.Int("REASONING_HIDE", 0)

// 前置message
var PRE_MESSAGES_JSON = env.String("PRE_MESSAGES_JSON", "")

var RateLimitCookieLockDuration = env.Int("RATE_LIMIT_COOKIE_LOCK_DURATION", 10*60)

// 路由前缀
var RoutePrefix = env.String("ROUTE_PREFIX", "")
var ModelChatMapStr = env.String("MODEL_CHAT_MAP", "")
var ModelChatMap = make(map[string]string)
var SessionImageChatMap = make(map[string]string)
var GlobalSessionManager *SessionManager

var SessionImageChatMapStr = env.String("SESSION_IMAGE_CHAT_MAP", "")
var YescaptchaClient *yescaptcha.Client

var AllDialogRecordEnable = os.Getenv("ALL_DIALOG_RECORD_ENABLE")
var RequestOutTime = os.Getenv("REQUEST_OUT_TIME")
var StreamRequestOutTime = os.Getenv("STREAM_REQUEST_OUT_TIME")
var SwaggerEnable = os.Getenv("SWAGGER_ENABLE")
var OnlyOpenaiApi = os.Getenv("ONLY_OPENAI_API")

var DebugEnabled = os.Getenv("DEBUG") == "true"

var RateLimitKeyExpirationDuration = 20 * time.Minute

var RequestOutTimeDuration = 5 * time.Minute

var (
	RequestRateLimitNum            = env.Int("REQUEST_RATE_LIMIT", 60)
	RequestRateLimitDuration int64 = 1 * 60
)

type RateLimitCookie struct {
	ExpirationTime time.Time // 过期时间
}

var (
	rateLimitCookies sync.Map // 使用 sync.Map 管理限速 Cookie
)

func AddRateLimitCookie(cookie string, expirationTime time.Time) {
	rateLimitCookies.Store(cookie, RateLimitCookie{
		ExpirationTime: expirationTime,
	})
	//fmt.Printf("Storing cookie: %s with value: %+v\n", cookie, RateLimitCookie{ExpirationTime: expirationTime})
}

type CookieManager struct {
	Cookies      []string
	currentIndex int
	mu           sync.Mutex
}

var (
	GSCookies    []string   // 存储所有的 cookies
	cookiesMutex sync.Mutex // 保护 GSCookies 的互斥锁
)

// InitGSCookies 初始化 GSCookies
func InitGSCookies() {
	cookiesMutex.Lock()
	defer cookiesMutex.Unlock()

	GSCookies = []string{}

	// 从环境变量中读取 GS_COOKIE 并拆分为切片
	cookieStr := os.Getenv("GS_COOKIE")
	if cookieStr != "" {

		for _, cookie := range strings.Split(cookieStr, ",") {
			// 如果 cookie 不包含 "session_id="，则添加前缀
			if !strings.Contains(cookie, "session_id=") {
				cookie = "session_id=" + cookie
			}
			GSCookies = append(GSCookies, cookie)
		}
	}
}

// RemoveCookie 删除指定的 cookie（支持并发）
func RemoveCookie(cookieToRemove string) {
	cookiesMutex.Lock()
	defer cookiesMutex.Unlock()

	// 创建一个新的切片，过滤掉需要删除的 cookie
	var newCookies []string
	for _, cookie := range GetGSCookies() {
		if cookie != cookieToRemove {
			newCookies = append(newCookies, cookie)
		}
	}

	// 更新 GSCookies
	GSCookies = newCookies
}

// GetGSCookies 获取 GSCookies 的副本
func GetGSCookies() []string {
	//cookiesMutex.Lock()
	//defer cookiesMutex.Unlock()

	// 返回 GSCookies 的副本，避免外部直接修改
	cookiesCopy := make([]string, len(GSCookies))
	copy(cookiesCopy, GSCookies)
	return cookiesCopy
}

// NewCookieManager 创建 CookieManager
func NewCookieManager() *CookieManager {
	var validCookies []string
	// 遍历 GSCookies
	for _, cookie := range GetGSCookies() {
		cookie = strings.TrimSpace(cookie)
		if cookie == "" {
			continue // 忽略空字符串
		}

		// 检查是否在 RateLimitCookies 中
		if value, ok := rateLimitCookies.Load(cookie); ok {
			rateLimitCookie, ok := value.(RateLimitCookie) // 正确转换为 RateLimitCookie
			if !ok {
				continue
			}
			if rateLimitCookie.ExpirationTime.After(time.Now()) {
				// 如果未过期，忽略该 cookie
				continue
			} else {
				// 如果已过期，从 RateLimitCookies 中删除
				rateLimitCookies.Delete(cookie)
			}
		}

		// 添加到有效 cookie 列表
		validCookies = append(validCookies, cookie)
	}

	return &CookieManager{
		Cookies:      validCookies,
		currentIndex: 0,
	}
}

func IsRateLimited(cookie string) bool {
	if value, ok := rateLimitCookies.Load(cookie); ok {
		rateLimitCookie := value.(RateLimitCookie)
		return rateLimitCookie.ExpirationTime.After(time.Now())
	}
	return false
}

func (cm *CookieManager) RemoveCookie(cookieToRemove string) error {
	cm.mu.Lock()
	defer cm.mu.Unlock()

	if len(cm.Cookies) == 0 {
		return errors.New("no cookies available")
	}

	// 查找要删除的cookie的索引
	index := -1
	for i, cookie := range cm.Cookies {
		if cookie == cookieToRemove {
			index = i
			break
		}
	}

	// 如果没找到要删除的cookie
	if index == -1 {
		return errors.New("RemoveCookie -> cookie not found")
	}

	// 从切片中删除cookie
	cm.Cookies = append(cm.Cookies[:index], cm.Cookies[index+1:]...)

	// 如果当前索引大于或等于删除后的切片长度，重置为0
	if cm.currentIndex >= len(cm.Cookies) {
		cm.currentIndex = 0
	}

	return nil
}

func (cm *CookieManager) GetNextCookie() (string, error) {
	cm.mu.Lock()
	defer cm.mu.Unlock()

	if len(cm.Cookies) == 0 {
		return "", errors.New("no cookies available")
	}

	cm.currentIndex = (cm.currentIndex + 1) % len(cm.Cookies)
	return cm.Cookies[cm.currentIndex], nil
}

func (cm *CookieManager) GetRandomCookie() (string, error) {
	cm.mu.Lock()
	defer cm.mu.Unlock()

	if len(cm.Cookies) == 0 {
		return "", errors.New("no cookies available")
	}

	// 生成随机索引
	randomIndex := rand.Intn(len(cm.Cookies))
	// 更新当前索引
	cm.currentIndex = randomIndex

	return cm.Cookies[randomIndex], nil
}

// SessionKey 定义复合键结构
type SessionKey struct {
	Cookie string
	Model  string
}

// SessionManager 会话管理器
type SessionManager struct {
	sessions map[SessionKey]string
	mutex    sync.RWMutex
}

// NewSessionManager 创建新的会话管理器
func NewSessionManager() *SessionManager {
	return &SessionManager{
		sessions: make(map[SessionKey]string),
	}
}

// AddSession 添加会话记录（写操作，需要写锁）
func (sm *SessionManager) AddSession(cookie string, model string, chatID string) {
	sm.mutex.Lock()
	defer sm.mutex.Unlock()

	key := SessionKey{
		Cookie: cookie,
		Model:  model,
	}
	sm.sessions[key] = chatID
}

// GetChatID 获取会话ID（读操作，使用读锁）
func (sm *SessionManager) GetChatID(cookie string, model string) (string, bool) {
	sm.mutex.RLock()
	defer sm.mutex.RUnlock()

	key := SessionKey{
		Cookie: cookie,
		Model:  model,
	}
	chatID, exists := sm.sessions[key]
	return chatID, exists
}

// DeleteSession 删除会话记录（写操作，需要写锁）
func (sm *SessionManager) DeleteSession(cookie string, model string) {
	sm.mutex.Lock()
	defer sm.mutex.Unlock()

	key := SessionKey{
		Cookie: cookie,
		Model:  model,
	}
	delete(sm.sessions, key)
}

// GetChatIDsByCookie 获取指定cookie关联的所有chatID列表(读操作,使用读锁)
func (sm *SessionManager) GetChatIDsByCookie(cookie string) []string {
	sm.mutex.RLock()
	defer sm.mutex.RUnlock()

	var chatIDs []string
	for key, chatID := range sm.sessions {
		if key.Cookie == cookie {
			chatIDs = append(chatIDs, chatID)
		}
	}
	return chatIDs
}

type SessionMapManager struct {
	sessionMap   map[string]string
	keys         []string
	currentIndex int
	mu           sync.Mutex
}

func NewSessionMapManager() *SessionMapManager {
	// 从初始map中提取所有的key
	keys := make([]string, 0, len(SessionImageChatMap))
	for k := range SessionImageChatMap {
		keys = append(keys, k)
	}

	return &SessionMapManager{
		sessionMap:   SessionImageChatMap,
		keys:         keys,
		currentIndex: 0,
	}
}

// GetCurrentKeyValue 获取当前索引对应的键值对
func (sm *SessionMapManager) GetCurrentKeyValue() (string, string, error) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	if len(sm.keys) == 0 {
		return "", "", errors.New("no sessions available")
	}

	currentKey := sm.keys[sm.currentIndex]
	return currentKey, sm.sessionMap[currentKey], nil
}

// GetNextKeyValue 获取下一个键值对
func (sm *SessionMapManager) GetNextKeyValue() (string, string, error) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	if len(sm.keys) == 0 {
		return "", "", errors.New("no sessions available")
	}

	sm.currentIndex = (sm.currentIndex + 1) % len(sm.keys)
	currentKey := sm.keys[sm.currentIndex]
	return currentKey, sm.sessionMap[currentKey], nil
}

// GetRandomKeyValue 随机获取一个键值对
func (sm *SessionMapManager) GetRandomKeyValue() (string, string, error) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	if len(sm.keys) == 0 {
		return "", "", errors.New("no sessions available")
	}

	randomIndex := rand.Intn(len(sm.keys))
	sm.currentIndex = randomIndex
	currentKey := sm.keys[randomIndex]
	return currentKey, sm.sessionMap[currentKey], nil
}

// AddKeyValue 添加新的键值对
func (sm *SessionMapManager) AddKeyValue(key, value string) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	// 如果key不存在，则添加到keys切片中
	if _, exists := sm.sessionMap[key]; !exists {
		sm.keys = append(sm.keys, key)
	}
	sm.sessionMap[key] = value
}

// RemoveKey 删除指定的键值对
func (sm *SessionMapManager) RemoveKey(key string) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	if _, exists := sm.sessionMap[key]; !exists {
		return
	}

	// 从map中删除
	delete(sm.sessionMap, key)

	// 从keys切片中删除
	for i, k := range sm.keys {
		if k == key {
			sm.keys = append(sm.keys[:i], sm.keys[i+1:]...)
			break
		}
	}

	// 调整currentIndex如果需要
	if sm.currentIndex >= len(sm.keys) && len(sm.keys) > 0 {
		sm.currentIndex = len(sm.keys) - 1
	}
}

// GetSize 获取当前map的大小
func (sm *SessionMapManager) GetSize() int {
	sm.mu.Lock()
	defer sm.mu.Unlock()
	return len(sm.keys)
}