common/websocket/model_api.go

package websocket

import (
	"context"
	"net/http"

	"github.com/Tencent/AI-Infra-Guard/common/utils/models"

	"github.com/Tencent/AI-Infra-Guard/pkg/database"
	"github.com/gin-gonic/gin"
	"trpc.group/trpc-go/trpc-go/log"
)

// ModelInfo 模型信息（用于创建）
type ModelInfo struct {
	Model   string `json:"model" binding:"required"`
	Token   string `json:"token" binding:"required"`
	BaseURL string `json:"base_url" binding:"required"`
	Limit   int    `json:"limit"`
	Note    string `json:"note"`
}

// CreateModelRequest 创建模型请求
type CreateModelRequest struct {
	ModelID string    `json:"model_id" binding:"required"`
	Model   ModelInfo `json:"model" binding:"required"`
}

// UpdateModelInfo 模型信息（用于更新）
// 这里不对 Token/BaseURL 使用 binding:"required"，以支持“只改名称等字段”的场景。
type UpdateModelInfo struct {
	Model   string `json:"model"`
	Token   string `json:"token"`
	BaseURL string `json:"base_url"`
	Limit   int    `json:"limit"`
	Note    string `json:"note"`
}

// UpdateModelRequest 更新模型请求
type UpdateModelRequest struct {
	Model UpdateModelInfo `json:"model" binding:"required"`
}

// DeleteModelRequest 删除模型请求
type DeleteModelRequest struct {
	ModelIDs []string `json:"model_ids" binding:"required"`
}

// ModelManager 模型管理器
type ModelManager struct {
	modelStore *database.ModelStore
}

const maskedToken = "********"

// maskToken 用于在对外返回模型信息时隐藏真实的 Token。
// 仅用于 JSON 返回，不影响数据库中实际存储的 Token。
func maskToken(token string) string {
	if token == "" {
		return ""
	}
	return maskedToken
}

// NewModelManager 创建新的ModelManager实例
func NewModelManager(modelStore *database.ModelStore) *ModelManager {
	return &ModelManager{
		modelStore: modelStore,
	}
}

// HandleGetModelList 获取模型列表接口
func HandleGetModelList(c *gin.Context, mm *ModelManager) {
	traceID := getTraceID(c)
	username := c.GetString("username")

	log.Debugf("用户请求获取模型列表: trace_id=%s, username=%s", traceID, username)

	var userModels []*database.Model
	var publicModels []*database.Model
	var err error

	// 获取public_user的模型（系统模型）
	publicModels, err = mm.modelStore.GetUserModels("public_user")
	if err != nil {
		log.Errorf("获取公共模型列表失败: trace_id=%s, username=%s, error=%v", traceID, username, err)
		// 公共模型获取失败不影响主流程，只记录警告
		log.Warnf("获取public_user模型失败，继续返回用户模型: %v", err)
		publicModels = []*database.Model{}
	}

	// 如果不是public_user，才获取用户自己的模型
	if username != "public_user" {
		userModels, err = mm.modelStore.GetUserModels(username)
		if err != nil {
			log.Errorf("获取用户模型列表失败: trace_id=%s, username=%s, error=%v", traceID, username, err)
			c.JSON(http.StatusOK, gin.H{
				"status":  1,
				"message": "获取模型列表失败: " + err.Error(),
				"data":    nil,
			})
			return
		}
	}

	// 转换为期望的返回格式
	var result []map[string]interface{}

	// 1. 首先添加系统模型（public_user的模型），永远排在前面
	for _, model := range publicModels {
		// 系统模型同样对外不暴露真实 token，但如果存在 token，
		// 通过掩码串告知“已配置密钥”，方便前端交互。
		item := map[string]interface{}{
			"model_id": model.ModelID,
			"model": map[string]interface{}{
				"model":    model.ModelName,
				"token":    maskToken(model.Token),
				"base_url": model.BaseURL,
				"note":     model.Note,
				"limit":    model.Limit,
			},
		}
		result = append(result, item)
	}

	// 2. 然后添加用户自己的模型
	for _, model := range userModels {
		item := map[string]interface{}{
			"model_id": model.ModelID,
			"model": map[string]interface{}{
				"model": model.ModelName,
				// 对外返回时也对用户模型的 token 进行掩码处理
				"token":    maskToken(model.Token),
				"base_url": model.BaseURL,
				"note":     model.Note,
				"limit":    model.Limit,
			},
		}
		result = append(result, item)
	}

	log.Debugf("获取模型列表成功: trace_id=%s, username=%s, userModels=%d, publicModels=%d, total=%d",
		traceID, username, len(userModels), len(publicModels), len(result))

	c.JSON(http.StatusOK, gin.H{
		"status":  0,
		"message": "获取模型列表成功",
		"data":    result,
	})
}

// HandleGetModelDetail 获取模型详情接口
func HandleGetModelDetail(c *gin.Context, mm *ModelManager) {
	traceID := getTraceID(c)
	modelID := c.Param("modelId")
	username := c.GetString("username")

	// 1. 字段校验
	if modelID == "" {
		log.Errorf("模型ID为空: trace_id=%s, username=%s", traceID, username)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "模型ID不能为空",
			"data":    nil,
		})
		return
	}

	log.Debugf("用户请求获取模型详情: trace_id=%s, modelID=%s, username=%s", traceID, modelID, username)

	// 2. 获取模型信息
	model, err := mm.modelStore.GetModel(modelID)
	if err != nil {
		log.Errorf("获取模型详情失败: trace_id=%s, modelID=%s, username=%s, error=%v", traceID, modelID, username, err)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "模型不存在",
			"data":    nil,
		})
		return
	}

	// 3. 身份校验（只有创建者可以查看）
	if model.Username != username {
		log.Errorf("无权限查看模型: trace_id=%s, modelID=%s, username=%s, owner=%s", traceID, modelID, username, model.Username)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "无权限查看此模型",
			"data":    nil,
		})
		return
	}

	log.Debugf("获取模型详情成功: trace_id=%s, modelID=%s, username=%s", traceID, modelID, username)

	// 转换为期望的返回格式
	result := map[string]interface{}{
		"model_id": model.ModelID,
		"model": map[string]interface{}{
			"model": model.ModelName,
			// 对外隐藏真实 token，前端如需修改，只能输入新 token
			"token":    maskToken(model.Token),
			"base_url": model.BaseURL,
			"note":     model.Note,
			"limit":    model.Limit,
		},
	}

	c.JSON(http.StatusOK, gin.H{
		"status":  0,
		"message": "获取模型详情成功",
		"data":    result,
	})
}

// HandleCreateModel 创建模型接口
func HandleCreateModel(c *gin.Context, mm *ModelManager) {
	traceID := getTraceID(c)
	username := c.GetString("username")

	// 1. 字段校验
	var req CreateModelRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		log.Errorf("请求参数解析失败: trace_id=%s, username=%s, error=%v", traceID, username, err)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "请求参数错误: " + err.Error(),
			"data":    nil,
		})
		return
	}

	// 2. 验证必填字段
	if req.ModelID == "" {
		log.Errorf("模型ID为空: trace_id=%s, username=%s", traceID, username)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "模型ID不能为空",
			"data":    nil,
		})
		return
	}

	if req.Model.Model == "" {
		log.Errorf("模型名称为空: trace_id=%s, username=%s", traceID, username)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "模型名称不能为空",
			"data":    nil,
		})
		return
	}

	if req.Model.Token == "" {
		log.Errorf("API Token为空: trace_id=%s, username=%s", traceID, username)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "API Token不能为空",
			"data":    nil,
		})
		return
	}

	if req.Model.BaseURL == "" {
		log.Errorf("基础URL为空: trace_id=%s, username=%s", traceID, username)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "基础URL不能为空",
			"data":    nil,
		})
		return
	}
	if req.Model.Limit == 0 {
		req.Model.Limit = 1000
	}

	log.Debugf("用户请求创建模型: trace_id=%s, modelID=%s, modelName=%s, username=%s", traceID, req.ModelID, req.Model.Model, username)

	// 3. 检查模型是否已存在
	exists, err := mm.modelStore.CheckModelExists(req.ModelID)
	if err != nil {
		log.Errorf("检查模型是否存在失败: trace_id=%s, modelID=%s, username=%s, error=%v", traceID, req.ModelID, username, err)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "检查模型失败: " + err.Error(),
			"data":    nil,
		})
		return
	}

	if exists {
		log.Errorf("模型已存在: trace_id=%s, modelID=%s, username=%s", traceID, req.ModelID, username)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "模型ID已存在",
			"data":    nil,
		})
		return
	}
	// 校验模型 token base_url
	ai := models.NewOpenAI(req.Model.Token, req.Model.Model, req.Model.BaseURL)
	err = ai.Vaild(context.Background())
	if err != nil {
		log.Errorf("模型校验失败: trace_id=%s, modelID=%s, username=%s, error=%v", traceID, req.ModelID, username, err)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "模型校验失败: " + err.Error(),
			"data":    nil,
		})
		return
	}

	// 4. 创建模型
	model := &database.Model{
		ModelID:   req.ModelID,
		Username:  username,
		ModelName: req.Model.Model,
		Token:     req.Model.Token,
		BaseURL:   req.Model.BaseURL,
		Note:      req.Model.Note,
		Limit:     req.Model.Limit,
	}

	err = mm.modelStore.CreateModel(model)
	if err != nil {
		log.Errorf("创建模型失败: trace_id=%s, modelID=%s, username=%s, error=%v", traceID, req.ModelID, username, err)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "创建模型失败: " + err.Error(),
			"data":    nil,
		})
		return
	}

	log.Debugf("创建模型成功: trace_id=%s, modelID=%s, modelName=%s, username=%s", traceID, req.ModelID, req.Model.Model, username)

	c.JSON(http.StatusOK, gin.H{
		"status":  0,
		"message": "模型创建成功",
		"data":    nil,
	})
}

// HandleUpdateModel 更新模型接口
func HandleUpdateModel(c *gin.Context, mm *ModelManager) {
	traceID := getTraceID(c)
	modelID := c.Param("modelId")
	username := c.GetString("username")

	// 1. 字段校验
	if modelID == "" {
		log.Errorf("模型ID为空: trace_id=%s, username=%s", traceID, username)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "模型ID不能为空",
			"data":    nil,
		})
		return
	}

	var req UpdateModelRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		log.Errorf("请求参数解析失败: trace_id=%s, modelID=%s, username=%s, error=%v", traceID, modelID, username, err)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "请求参数错误: " + err.Error(),
			"data":    nil,
		})
		return
	}

	log.Infof("用户请求更新模型: trace_id=%s, modelID=%s, username=%s", traceID, modelID, username)

	// 2. 身份校验（检查模型是否存在且属于该用户）
	exists, err := mm.modelStore.CheckModelExistsByUser(modelID, username)
	if err != nil {
		log.Errorf("检查模型权限失败: trace_id=%s, modelID=%s, username=%s, error=%v", traceID, modelID, username, err)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "检查模型权限失败: " + err.Error(),
			"data":    nil,
		})
		return
	}

	if !exists {
		log.Errorf("模型不存在或无权限: trace_id=%s, modelID=%s, username=%s", traceID, modelID, username)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "模型不存在或无权限",
			"data":    nil,
		})
		return
	}

	// 3. 构造更新字段
	// 支持“只改模型名称/备注，不改 key/base_url”的场景：
	// - 前端在编辑时如果不填写 token/base_url，则保持数据库中的原值不变；
	// - 只有在显式传入新 token/base_url 且不等于掩码串时才会更新。
	updates := map[string]interface{}{
		"model_name": req.Model.Model,
		"note":       req.Model.Note,
		"limit":      req.Model.Limit,
	}
	if req.Model.Token != "" && req.Model.Token != maskedToken {
		updates["token"] = req.Model.Token
	}
	if req.Model.BaseURL != "" {
		updates["base_url"] = req.Model.BaseURL
	}

	err = mm.modelStore.UpdateModel(modelID, username, updates)
	if err != nil {
		log.Errorf("更新模型失败: trace_id=%s, modelID=%s, username=%s, error=%v", traceID, modelID, username, err)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "更新模型失败: " + err.Error(),
			"data":    nil,
		})
		return
	}

	log.Infof("更新模型成功: trace_id=%s, modelID=%s, username=%s", traceID, modelID, username)

	c.JSON(http.StatusOK, gin.H{
		"status":  0,
		"message": "模型更新成功",
		"data":    nil,
	})
}

// HandleDeleteModel 删除模型接口（支持单个和批量）
func HandleDeleteModel(c *gin.Context, mm *ModelManager) {
	traceID := getTraceID(c)
	username := c.GetString("username")

	// 1. 字段校验
	var req DeleteModelRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		log.Errorf("请求参数解析失败: trace_id=%s, username=%s, error=%v", traceID, username, err)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "请求参数错误: " + err.Error(),
			"data":    nil,
		})
		return
	}

	if len(req.ModelIDs) == 0 {
		log.Errorf("模型ID列表为空: trace_id=%s, username=%s", traceID, username)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "模型ID列表不能为空",
			"data":    nil,
		})
		return
	}

	log.Infof("用户请求删除模型: trace_id=%s, modelIDs=%v, username=%s", traceID, req.ModelIDs, username)

	// 2. 身份校验（检查所有模型是否属于该用户）
	for _, modelID := range req.ModelIDs {
		exists, err := mm.modelStore.CheckModelExistsByUser(modelID, username)
		if err != nil {
			log.Errorf("检查模型权限失败: trace_id=%s, modelID=%s, username=%s, error=%v", traceID, modelID, username, err)
			c.JSON(http.StatusOK, gin.H{
				"status":  1,
				"message": "检查模型权限失败: " + err.Error(),
				"data":    nil,
			})
			return
		}

		if !exists {
			log.Errorf("模型不存在或无权限: trace_id=%s, modelID=%s, username=%s", traceID, modelID, username)
			c.JSON(http.StatusOK, gin.H{
				"status":  1,
				"message": "模型不存在或无权限",
				"data":    nil,
			})
			return
		}
	}

	// 3. 批量删除模型
	deletedCount, err := mm.modelStore.BatchDeleteModels(req.ModelIDs, username)
	if err != nil {
		log.Errorf("删除模型失败: trace_id=%s, modelIDs=%v, username=%s, error=%v", traceID, req.ModelIDs, username, err)
		c.JSON(http.StatusOK, gin.H{
			"status":  1,
			"message": "删除模型失败: " + err.Error(),
			"data":    nil,
		})
		return
	}

	log.Infof("删除模型成功: trace_id=%s, modelIDs=%v, username=%s, deletedCount=%d", traceID, req.ModelIDs, username, deletedCount)

	c.JSON(http.StatusOK, gin.H{
		"status":  0,
		"message": "删除成功",
		"data":    nil,
	})
}