gemini-banlancer/internal/handlers/apikey_handler.go

// Filename: internal/handlers/apikey_handler.go
package handlers

import (
	"gemini-balancer/internal/errors"
	"gemini-balancer/internal/models"
	"gemini-balancer/internal/response"
	"gemini-balancer/internal/service"
	"gemini-balancer/internal/task"
	"net/http"
	"strconv"
	"strings"

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

type APIKeyHandler struct {
	apiKeyService        *service.APIKeyService
	db                   *gorm.DB
	keyImportService     *service.KeyImportService
	keyValidationService *service.KeyValidationService
}

func NewAPIKeyHandler(apiKeyService *service.APIKeyService, db *gorm.DB, keyImportService *service.KeyImportService, keyValidationService *service.KeyValidationService) *APIKeyHandler {
	return &APIKeyHandler{
		apiKeyService:        apiKeyService,
		db:                   db,
		keyImportService:     keyImportService,
		keyValidationService: keyValidationService,
	}
}

// DTOs for API requests
type BulkAddKeysToGroupRequest struct {
	KeyGroupID       uint   `json:"key_group_id" binding:"required"`
	Keys             string `json:"keys"         binding:"required"`
	ValidateOnImport bool   `json:"validate_on_import"` // OmitEmpty/default is false
}

type BulkUnlinkKeysFromGroupRequest struct {
	KeyGroupID uint   `json:"key_group_id" binding:"required"`
	Keys       string `json:"keys"         binding:"required"`
}

type BulkHardDeleteKeysRequest struct {
	Keys string `json:"keys" binding:"required"`
}

type BulkRestoreKeysRequest struct {
	Keys string `json:"keys" binding:"required"`
}

type UpdateAPIKeyRequest struct {
	Status *string `json:"status" binding:"omitempty,oneof=ACTIVE,PENDING_VALIDATION,COOLDOWN,DISABLED,BANNED"`
}

type UpdateMappingRequest struct {
	Status models.APIKeyStatus `json:"status" binding:"required,oneof=ACTIVE PENDING_VALIDATION COOLDOWN DISABLED BANNED"`
}

type BulkTestKeysRequest struct {
	KeyGroupID uint   `json:"key_group_id" binding:"required"`
	Keys       string `json:"keys" binding:"required"`
}

type RestoreKeysRequest struct {
	KeyIDs []uint `json:"key_ids" binding:"required,gt=0"`
}
type BulkTestKeysForGroupRequest struct {
	Keys string `json:"keys" binding:"required"`
}

type BulkActionFilter struct {
	Status []string `json:"status"` // Changed to slice to accept multiple statuses
}
type BulkActionRequest struct {
	Action    string           `json:"action" binding:"required,oneof=revalidate set_status delete"`
	NewStatus string           `json:"new_status" binding:"omitempty,oneof=active disabled cooldown banned"` // For 'set_status' action
	Filter    BulkActionFilter `json:"filter" binding:"required"`
}

// --- Handler Methods ---

// AddMultipleKeysToGroup handles adding/linking multiple keys to a specific group.
func (h *APIKeyHandler) AddMultipleKeysToGroup(c *gin.Context) {
	var req BulkAddKeysToGroupRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrInvalidJSON, err.Error()))
		return
	}
	taskStatus, err := h.keyImportService.StartAddKeysTask(req.KeyGroupID, req.Keys, req.ValidateOnImport)
	if err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrTaskInProgress, err.Error()))
		return
	}
	response.Success(c, taskStatus)
}

// UnlinkMultipleKeysFromGroup handles unlinking multiple keys from a specific group.
func (h *APIKeyHandler) UnlinkMultipleKeysFromGroup(c *gin.Context) {
	var req BulkUnlinkKeysFromGroupRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrInvalidJSON, err.Error()))
		return
	}
	taskStatus, err := h.keyImportService.StartUnlinkKeysTask(req.KeyGroupID, req.Keys)
	if err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrTaskInProgress, err.Error()))
		return
	}
	response.Success(c, taskStatus)
}

// HardDeleteMultipleKeys handles globally deleting multiple key entities.
func (h *APIKeyHandler) HardDeleteMultipleKeys(c *gin.Context) {
	var req BulkHardDeleteKeysRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrInvalidJSON, err.Error()))
		return
	}
	taskStatus, err := h.keyImportService.StartHardDeleteKeysTask(req.Keys)
	if err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrTaskInProgress, err.Error()))
		return
	}
	response.Success(c, taskStatus)
}

// RestoreMultipleKeys handles restoring multiple keys to ACTIVE status globally.
func (h *APIKeyHandler) RestoreMultipleKeys(c *gin.Context) {
	var req BulkRestoreKeysRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrInvalidJSON, err.Error()))
		return
	}
	taskStatus, err := h.keyImportService.StartRestoreKeysTask(req.Keys)
	if err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrTaskInProgress, err.Error()))
		return
	}
	response.Success(c, taskStatus)
}

func (h *APIKeyHandler) TestMultipleKeys(c *gin.Context) {
	var req BulkTestKeysRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrInvalidJSON, err.Error()))
		return
	}
	taskStatus, err := h.keyValidationService.StartTestKeysTask(req.KeyGroupID, req.Keys)
	if err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrTaskInProgress, err.Error()))
		return
	}
	response.Success(c, taskStatus)
}

func (h *APIKeyHandler) ListAPIKeys(c *gin.Context) {
	var params models.APIKeyQueryParams
	if err := c.ShouldBindQuery(&params); err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrBadRequest, err.Error()))
		return
	}
	if params.IDs != "" {
		idStrs := strings.Split(params.IDs, ",")
		ids := make([]uint, 0, len(idStrs))
		for _, s := range idStrs {
			id, err := strconv.ParseUint(s, 10, 64)
			if err == nil {
				ids = append(ids, uint(id))
			}
		}
		if len(ids) > 0 {
			keys, err := h.apiKeyService.GetKeysByIds(ids)
			if err != nil {
				response.Error(c, &errors.APIError{
					HTTPStatus: http.StatusInternalServerError,
					Code:       "DATA_FETCH_ERROR",
					Message:    err.Error(),
				})
				return
			}
			response.Success(c, keys)
			return
		}
	}
	if params.Page <= 0 {
		params.Page = 1
	}
	if params.PageSize <= 0 {
		params.PageSize = 20
	}
	result, err := h.apiKeyService.ListAPIKeys(&params)
	if err != nil {
		response.Error(c, errors.ParseDBError(err))
		return
	}
	response.Success(c, result)
}

// ListKeysForGroup handles the GET /keygroups/:id/keys request.
func (h *APIKeyHandler) ListKeysForGroup(c *gin.Context) {
	// 1. Manually handle the path parameter.
	groupID, err := strconv.ParseUint(c.Param("id"), 10, 32)
	if err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrBadRequest, "Invalid group ID format"))
		return
	}
	// 2. Bind query parameters using the correctly tagged struct.
	var params models.APIKeyQueryParams
	if err := c.ShouldBindQuery(&params); err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrBadRequest, err.Error()))
		return
	}
	// 3. Set server-side defaults and the path parameter.
	if params.Page <= 0 {
		params.Page = 1
	}
	if params.PageSize <= 0 {
		params.PageSize = 20
	}
	params.KeyGroupID = uint(groupID)
	// 4. Call the service layer.
	paginatedResult, err := h.apiKeyService.ListAPIKeys(&params)
	if err != nil {
		response.Error(c, errors.ParseDBError(err))
		return
	}

	// 5. [THE FIX] Return a successful response using the standard `response.Success`
	// and a gin.H map, as confirmed to exist in your project.
	response.Success(c, gin.H{
		"items": paginatedResult.Items,
		"total": paginatedResult.Total,
		"page":  paginatedResult.Page,
		"pages": paginatedResult.TotalPages,
	})
}

func (h *APIKeyHandler) TestKeysForGroup(c *gin.Context) {
	// Group ID is now correctly sourced from the URL path.
	groupID, err := strconv.ParseUint(c.Param("id"), 10, 32)
	if err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrBadRequest, "Invalid group ID format"))
		return
	}
	// The request body is now simpler, only needing the keys.
	var req BulkTestKeysForGroupRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrInvalidJSON, err.Error()))
		return
	}
	// Call the same underlying service, but with unambiguous context.
	taskStatus, err := h.keyValidationService.StartTestKeysTask(uint(groupID), req.Keys)
	if err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrTaskInProgress, err.Error()))
		return
	}
	response.Success(c, taskStatus)
}

// UpdateAPIKey is DEPRECATED. Status is now contextual to a group.
func (h *APIKeyHandler) UpdateAPIKey(c *gin.Context) {
	err := errors.NewAPIError(errors.ErrBadRequest, "This endpoint is deprecated. Use 'PUT /keygroups/:id/apikeys/:keyId' to update key status within a group context.")
	response.Error(c, err)
}

// UpdateGroupAPIKeyMapping handles updating a key's status within a specific group.
// Route: PUT /keygroups/:id/apikeys/:keyId
func (h *APIKeyHandler) UpdateGroupAPIKeyMapping(c *gin.Context) {
	groupID, err := strconv.ParseUint(c.Param("id"), 10, 32)
	if err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrBadRequest, "Invalid Group ID format"))
		return
	}
	keyID, err := strconv.ParseUint(c.Param("keyId"), 10, 32)
	if err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrBadRequest, "Invalid Key ID format"))
		return
	}
	var req UpdateMappingRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrInvalidJSON, err.Error()))
		return
	}
	// Directly use the service to handle the logic
	updatedMapping, err := h.apiKeyService.UpdateMappingStatus(uint(groupID), uint(keyID), req.Status)
	if err != nil {
		var apiErr *errors.APIError
		if errors.As(err, &apiErr) {
			response.Error(c, apiErr)
		} else {
			response.Error(c, errors.ParseDBError(err))
		}
		return
	}
	response.Success(c, updatedMapping)
}

// HardDeleteAPIKey handles globally deleting a single key entity.
func (h *APIKeyHandler) HardDeleteAPIKey(c *gin.Context) {
	id, err := strconv.Atoi(c.Param("id"))
	if err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrBadRequest, "Invalid ID format"))
		return
	}
	if err := h.apiKeyService.HardDeleteAPIKeyByID(uint(id)); err != nil {
		response.Error(c, errors.ParseDBError(err))
		return
	}
	response.Success(c, gin.H{"message": "API key globally deleted successfully"})
}

// RestoreKeysInGroup 恢复指定Key的接口
// POST /keygroups/:id/apikeys/restore
func (h *APIKeyHandler) RestoreKeysInGroup(c *gin.Context) {
	groupID, err := strconv.ParseUint(c.Param("id"), 10, 32)
	if err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrBadRequest, "Invalid Group ID format"))
		return
	}
	var req RestoreKeysRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrInvalidJSON, err.Error()))
		return
	}
	taskStatus, err := h.apiKeyService.StartRestoreKeysTask(uint(groupID), req.KeyIDs)
	if err != nil {
		var apiErr *errors.APIError
		if errors.As(err, &apiErr) {
			response.Error(c, apiErr)
		} else {
			response.Error(c, errors.NewAPIError(errors.ErrInternalServer, err.Error()))
		}
		return
	}
	response.Success(c, taskStatus)
}

// RestoreAllBannedInGroup 一键恢复所有Banned Key的接口
// POST /keygroups/:id/apikeys/restore-all-banned
func (h *APIKeyHandler) RestoreAllBannedInGroup(c *gin.Context) {
	groupID, err := strconv.ParseUint(c.Param("groupId"), 10, 32)
	if err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrBadRequest, "Invalid Group ID format"))
		return
	}
	taskStatus, err := h.apiKeyService.StartRestoreAllBannedTask(uint(groupID))
	if err != nil {
		var apiErr *errors.APIError
		if errors.As(err, &apiErr) {
			response.Error(c, apiErr)
		} else {
			response.Error(c, errors.NewAPIError(errors.ErrInternalServer, err.Error()))
		}
		return
	}
	response.Success(c, taskStatus)
}

// HandleBulkAction handles generic bulk actions on a key group based on server-side filters.
// Route: POST /keygroups/:id/bulk-actions
func (h *APIKeyHandler) HandleBulkAction(c *gin.Context) {
	// 1. Parse GroupID from URL
	groupID, err := strconv.ParseUint(c.Param("id"), 10, 32)
	if err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrBadRequest, "Invalid Group ID format"))
		return
	}
	// 2. Bind the JSON payload to our new DTO
	var req BulkActionRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrInvalidJSON, err.Error()))
		return
	}
	// 3. Central logic: based on the action, call the appropriate service method.
	var task *task.Status
	var apiErr *errors.APIError
	switch req.Action {
	case "revalidate":
		// Assume keyValidationService has a method that accepts a filter
		task, err = h.keyValidationService.StartTestKeysByFilterTask(uint(groupID), req.Filter.Status)
	case "set_status":
		if req.NewStatus == "" {
			apiErr = errors.NewAPIError(errors.ErrBadRequest, "new_status is required for set_status action")
			break
		}
		// Assume apiKeyService has a method to update status by filter
		targetStatus := models.APIKeyStatus(req.NewStatus) // Convert string to your model's type
		task, err = h.apiKeyService.StartUpdateStatusByFilterTask(uint(groupID), req.Filter.Status, targetStatus)
	case "delete":
		// Assume keyImportService has a method to unlink by filter
		task, err = h.keyImportService.StartUnlinkKeysByFilterTask(uint(groupID), req.Filter.Status)
	default:
		apiErr = errors.NewAPIError(errors.ErrBadRequest, "Unsupported action: "+req.Action)
	}
	// 4. Handle errors from the switch block
	if apiErr != nil {
		response.Error(c, apiErr)
		return
	}
	if err != nil {
		// Attempt to parse it as a known APIError, otherwise, wrap it.
		var parsedErr *errors.APIError
		if errors.As(err, &parsedErr) {
			response.Error(c, parsedErr)
		} else {
			response.Error(c, errors.NewAPIError(errors.ErrInternalServer, err.Error()))
		}
		return
	}
	// 5. Return the task status on success
	response.Success(c, task)
}

// ExportKeysForGroup handles requests to export all keys for a group based on status filters.
// Route: GET /keygroups/:id/apikeys/export
func (h *APIKeyHandler) ExportKeysForGroup(c *gin.Context) {
	groupID, err := strconv.ParseUint(c.Param("id"), 10, 32)
	if err != nil {
		response.Error(c, errors.NewAPIError(errors.ErrBadRequest, "Invalid Group ID format"))
		return
	}
	// Use QueryArray to correctly parse `status[]=active&status[]=cooldown`
	statuses := c.QueryArray("status")
	keyStrings, err := h.apiKeyService.GetAPIKeyStringsForExport(uint(groupID), statuses)
	if err != nil {
		response.Error(c, errors.ParseDBError(err))
		return
	}
	response.Success(c, keyStrings)
}