controller/admin.go

package controller

import (
	"encoding/json"
	"fmt"
	"genspark2api/common"
	"genspark2api/common/config"
	logger "genspark2api/common/loggger"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
)

// 系统启动时间
var startTime = time.Now()

// GSCookie 定义cookie的结构
type GSCookie struct {
	Value        string    // cookie值
	RequestCount int       // 请求次数
	FailCount    int       // 失败请求数
	LastUsedTime time.Time // 最后使用时间
	Disabled     bool      // 是否禁用
}

// 模拟的Cookie统计
var cookieStats = make(map[string]*GSCookie)

// 模型信息
type ModelInfo struct {
	ID       string `json:"id"`
	Name     string `json:"name"`
	Category string `json:"category"`
	Enabled  bool   `json:"enabled"`
}

// ModelsListRequest 请求体
type ModelsListRequest struct {
	Models []ModelInfo `json:"models"`
}

// 全局变量，存储模型列表
var modelsList []ModelInfo
var modelsListFile = "./data/models_list.json"

// 初始化模型列表
func init() {
	// 加载模型列表（如果存在）
	loadModelsList()
}

// AdminIndex 返回管理界面
func AdminIndex(c *gin.Context) {
	c.File("./public/index.html")
}

// GetConfig 获取当前配置
func GetConfig(c *gin.Context) {
	// 从环境变量获取配置
	configData := map[string]string{
		"GS_COOKIE":                       os.Getenv("GS_COOKIE"),
		"API_SECRET":                      os.Getenv("API_SECRET"),
		"AUTO_MODEL_CHAT_MAP_TYPE":        os.Getenv("AUTO_MODEL_CHAT_MAP_TYPE"),
		"AUTO_DEL_CHAT":                   os.Getenv("AUTO_DEL_CHAT"),
		"MODEL_CHAT_MAP":                  os.Getenv("MODEL_CHAT_MAP"),
		"PROXY_URL":                       os.Getenv("PROXY_URL"),
		"RECAPTCHA_PROXY_URL":             os.Getenv("RECAPTCHA_PROXY_URL"),
		"REQUEST_RATE_LIMIT":              os.Getenv("REQUEST_RATE_LIMIT"),
		"RATE_LIMIT_COOKIE_LOCK_DURATION": os.Getenv("RATE_LIMIT_COOKIE_LOCK_DURATION"),
		"ROUTE_PREFIX":                    os.Getenv("ROUTE_PREFIX"),
		"DEBUG":                           os.Getenv("DEBUG"),
		"REASONING_HIDE":                  os.Getenv("REASONING_HIDE"),
	}

	c.JSON(http.StatusOK, configData)
}

// SaveConfig 保存配置到.env文件
func SaveConfig(c *gin.Context) {
	// 解析请求体
	var configData map[string]string
	if err := c.ShouldBindJSON(&configData); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{
			"success": false,
			"message": "无效的请求数据",
		})
		return
	}

	// 读取现有的.env文件（如果存在）
	envFilePath := "./data/.env"
	envLines := []string{}
	envMap := make(map[string]string)

	// 确保data目录存在
	if err := os.MkdirAll("./data", 0755); err != nil {
		logger.SysLog(fmt.Sprintf("创建data目录失败: %s", err.Error()))
	}

	if fileData, err := ioutil.ReadFile(envFilePath); err == nil {
		content := string(fileData)
		lines := strings.Split(content, "\n")
		for _, line := range lines {
			line = strings.TrimSpace(line)
			if line != "" && !strings.HasPrefix(line, "#") {
				parts := strings.SplitN(line, "=", 2)
				if len(parts) == 2 {
					key := strings.TrimSpace(parts[0])
					value := strings.TrimSpace(parts[1])
					envMap[key] = value
				} else {
					envLines = append(envLines, line) // 保留非键值对行
				}
			} else {
				envLines = append(envLines, line) // 保留注释和空行
			}
		}
	}

	// 更新配置
	for key, value := range configData {
		if value != "" {
			envMap[key] = value
			// 同时更新环境变量
			os.Setenv(key, value)
		}
	}

	// 重建.env文件内容
	var newContent strings.Builder

	// 先写入非键值对的行
	for _, line := range envLines {
		if !strings.Contains(line, "=") {
			newContent.WriteString(line)
			newContent.WriteString("\n")
		}
	}

	// 写入所有配置
	for key, value := range envMap {
		newContent.WriteString(fmt.Sprintf("%s=%s\n", key, value))
	}

	// 写入文件
	if err := ioutil.WriteFile(envFilePath, []byte(newContent.String()), 0644); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{
			"success": false,
			"message": "保存配置失败: " + err.Error(),
		})
		return
	}

	// 特殊处理 - 如果更新了GS_COOKIE，则重新初始化cookie池
	if _, ok := configData["GS_COOKIE"]; ok {
		config.InitGSCookies()
		// 更新Cookie统计信息
		updateCookieStats()
	}

	c.JSON(http.StatusOK, gin.H{
		"success": true,
		"message": "配置已保存",
	})
}

// updateCookieStats 更新Cookie统计信息
func updateCookieStats() {
	// 获取当前所有Cookie
	currentCookies := config.GetGSCookies()

	// 为新的Cookie创建统计条目
	for _, cookie := range currentCookies {
		if _, exists := cookieStats[cookie]; !exists {
			cookieStats[cookie] = &GSCookie{
				Value:        cookie,
				RequestCount: 0,
				FailCount:    0,
				LastUsedTime: time.Now(),
				Disabled:     false,
			}
		}
	}

	// 移除不再存在的Cookie
	for cookie := range cookieStats {
		exists := false
		for _, currentCookie := range currentCookies {
			if cookie == currentCookie {
				exists = true
				break
			}
		}
		if !exists {
			delete(cookieStats, cookie)
		}
	}
}

// GetStatus 获取系统状态信息
func GetStatus(c *gin.Context) {
	// 计算运行时间
	uptime := time.Since(startTime).String()

	// 确保Cookie统计已更新
	updateCookieStats()

	// Cookie池信息
	cookieCount := len(config.GetGSCookies())
	cookies := make([]map[string]interface{}, 0, cookieCount)

	// 获取所有Cookie信息
	for cookie, stats := range cookieStats {
		cookieInfo := map[string]interface{}{
			"value":        cookie,
			"active":       !stats.Disabled,
			"requestCount": stats.RequestCount,
			"lastUsed":     stats.LastUsedTime.Format("2006-01-02 15:04:05"),
		}
		cookies = append(cookies, cookieInfo)
	}

	// 计算请求统计
	totalRequests := 0
	successRequests := 0
	failedRequests := 0

	for _, stats := range cookieStats {
		totalRequests += stats.RequestCount
		successRequests += (stats.RequestCount - stats.FailCount)
		failedRequests += stats.FailCount
	}

	// 返回状态信息
	c.JSON(http.StatusOK, gin.H{
		"version":         common.Version,
		"uptime":          uptime,
		"cookieCount":     cookieCount,
		"totalRequests":   totalRequests,
		"successRequests": successRequests,
		"failedRequests":  failedRequests,
		"cookies":         cookies,
	})
}

// GetModels 获取模型列表
func GetModels(c *gin.Context) {
	// 如果模型列表为空，从默认OpenAI模型列表初始化
	if len(modelsList) == 0 {
		initializeModelsFromDefaults()
	}

	c.JSON(http.StatusOK, gin.H{
		"models": modelsList,
	})
}

// SaveModels 保存模型列表
func SaveModels(c *gin.Context) {
	var request ModelsListRequest
	if err := c.ShouldBindJSON(&request); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{
			"success": false,
			"message": "无效的请求数据",
		})
		return
	}

	// 更新模型列表
	modelsList = request.Models

	// 保存到文件
	if err := saveModelsToFile(); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{
			"success": false,
			"message": "保存模型列表失败: " + err.Error(),
		})
		return
	}

	// 更新DefaultOpenaiModelList
	updateDefaultOpenaiModelList()

	c.JSON(http.StatusOK, gin.H{
		"success": true,
		"message": "模型列表已保存",
	})
}

// 从默认值初始化模型列表
func initializeModelsFromDefaults() {
	// 清空现有列表
	modelsList = []ModelInfo{}

	// 文本模型
	for _, model := range common.TextModelList {
		modelsList = append(modelsList, ModelInfo{
			ID:       model,
			Name:     getModelDisplayName(model),
			Category: "text",
			Enabled:  true,
		})
	}

	// 图像模型
	for _, model := range common.ImageModelList {
		modelsList = append(modelsList, ModelInfo{
			ID:       model,
			Name:     getModelDisplayName(model),
			Category: "image",
			Enabled:  true,
		})
	}

	// 视频模型
	for _, model := range common.VideoModelList {
		modelsList = append(modelsList, ModelInfo{
			ID:       model,
			Name:     getModelDisplayName(model),
			Category: "video",
			Enabled:  true,
		})
	}

	// 保存到文件
	saveModelsToFile()
}

// 获取模型显示名称
func getModelDisplayName(modelID string) string {
	// 生成一个友好的显示名称
	displayName := modelID

	// 转换为首字母大写
	parts := strings.Split(displayName, "-")
	for i, part := range parts {
		if len(part) > 0 {
			// 对一些常见缩写进行特殊处理
			if part == "gpt" || part == "dall" {
				parts[i] = strings.ToUpper(part)
			} else {
				parts[i] = strings.ToUpper(part[:1]) + part[1:]
			}
		}
	}

	displayName = strings.Join(parts, " ")
	displayName = strings.Replace(displayName, "/", " ", -1)

	return displayName
}

// 保存模型列表到文件
func saveModelsToFile() error {
	// 确保data目录存在
	if err := os.MkdirAll("./data", 0755); err != nil {
		return err
	}

	// 序列化数据
	data, err := json.MarshalIndent(modelsList, "", "  ")
	if err != nil {
		return err
	}

	// 写入文件
	return ioutil.WriteFile(modelsListFile, data, 0644)
}

// 从文件加载模型列表
func loadModelsList() {
	// 如果文件不存在，则初始化默认值
	if _, err := os.Stat(modelsListFile); os.IsNotExist(err) {
		initializeModelsFromDefaults()
		return
	}

	// 读取文件
	data, err := ioutil.ReadFile(modelsListFile)
	if err != nil {
		logger.SysLog(fmt.Sprintf("无法读取模型列表文件: %s", err.Error()))
		initializeModelsFromDefaults()
		return
	}

	// 解析JSON
	if err := json.Unmarshal(data, &modelsList); err != nil {
		logger.SysLog(fmt.Sprintf("解析模型列表文件失败: %s", err.Error()))
		initializeModelsFromDefaults()
		return
	}

	// 如果列表为空，初始化默认值
	if len(modelsList) == 0 {
		initializeModelsFromDefaults()
	}

	// 更新DefaultOpenaiModelList
	updateDefaultOpenaiModelList()
}

// 更新DefaultOpenaiModelList
func updateDefaultOpenaiModelList() {
	// 清空现有列表
	common.DefaultOpenaiModelList = []string{}

	// 遍历模型列表，将启用的模型添加到DefaultOpenaiModelList
	for _, model := range modelsList {
		if model.Enabled {
			common.DefaultOpenaiModelList = append(common.DefaultOpenaiModelList, model.ID)
		}
	}
}