new-api/controller/kling_aiping_native.go

package controller

import (
	"bytes"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"strings"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/dto"
	"github.com/QuantumNous/new-api/middleware"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/relay"
	klingaiping "github.com/QuantumNous/new-api/relay/channel/task/kling/aiping"
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/QuantumNous/new-api/service"
	"github.com/QuantumNous/new-api/types"
	"github.com/gin-gonic/gin"
)

func KlingAipingNativeTaskSubmit(c *gin.Context) {
	route, ok := klingaiping.FindRoute(c.Request.Method, c.FullPath(), klingaiping.RouteKindSubmit)
	if !ok {
		c.JSON(http.StatusNotFound, gin.H{"code": 404, "message": "route not found"})
		return
	}

	relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"code": 500, "message": err.Error()})
		return
	}
	payload, err := readJSONPayload(c)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"code": 400, "message": err.Error()})
		return
	}

	configureKlingAipingTaskRelayInfo(c, relayInfo, route, payload)

	relayTaskWithInfo(c, relayInfo)
}

func configureKlingAipingTaskRelayInfo(c *gin.Context, relayInfo *relaycommon.RelayInfo, route klingaiping.Route, payload map[string]any) {
	// Native Kling routes bypass Distribute(), so force relayTaskWithInfo to
	// select a channel instead of reading a preselected one from context.
	if relayInfo.ChannelMeta == nil {
		relayInfo.ChannelMeta = &relaycommon.ChannelMeta{}
	}
	modelName := resolveKlingAipingModel(payload, route)
	relayInfo.OriginModelName = modelName
	relayInfo.Action = route.Action
	relaycommon.StoreTaskRequest(c, relayInfo, route.Action, relaycommon.TaskSubmitReq{
		Model:    modelName,
		Prompt:   stringFromMap(payload, "prompt"),
		Seconds:  fmt.Sprint(payload["seconds"]),
		Metadata: payload,
	})
}

func KlingAipingNativeTaskFetch(c *gin.Context) {
	route, ok := klingaiping.FindRoute(c.Request.Method, c.FullPath(), klingaiping.RouteKindFetch)
	if !ok {
		c.JSON(http.StatusNotFound, gin.H{"code": 404, "message": "route not found"})
		return
	}
	task, ok := getKlingAipingUserTask(c, c.Param("task_id"), route.Action)
	if !ok {
		return
	}
	refreshKlingAipingTaskIfNeeded(task)
	c.JSON(http.StatusOK, buildKlingAipingTaskPayload(task))
}

func KlingAipingNativeTaskList(c *gin.Context) {
	route, ok := klingaiping.FindRoute(c.Request.Method, c.FullPath(), klingaiping.RouteKindList)
	if !ok {
		c.JSON(http.StatusNotFound, gin.H{"code": 404, "message": "route not found"})
		return
	}
	pageNum, pageSize, err := parseKlingAipingPage(c)
	if err != nil {
		c.JSON(http.StatusUnprocessableEntity, gin.H{"code": 422, "message": err.Error()})
		return
	}
	tasks := model.TaskGetAllUserTask(c.GetInt("id"), (pageNum-1)*pageSize, pageSize, model.SyncTaskQueryParams{
		Platform: constant.TaskPlatform(strconv.Itoa(constant.ChannelTypeKlingAiping)),
		Action:   route.Action,
	})
	data := make([]any, 0, len(tasks))
	for _, task := range tasks {
		data = append(data, taskDataObject(task))
	}
	c.JSON(http.StatusOK, gin.H{
		"code":       0,
		"message":    "success",
		"request_id": c.GetString(common.RequestIdKey),
		"data":       data,
	})
}

func KlingAipingNativeProxy(c *gin.Context) {
	route, ok := klingaiping.FindRoute(c.Request.Method, c.FullPath(), klingaiping.RouteKindProxy)
	if !ok {
		c.JSON(http.StatusNotFound, gin.H{"code": 404, "message": "route not found"})
		return
	}
	channel, _, err := service.CacheGetRandomSatisfiedChannel(&service.RetryParam{
		Ctx:                 c,
		TokenGroup:          common.GetContextKeyString(c, constant.ContextKeyUsingGroup),
		ModelName:           route.BillingModel,
		Retry:               common.GetPointer(0),
		RequiredChannelType: constant.ChannelTypeKlingAiping,
	})
	if err != nil {
		c.JSON(http.StatusServiceUnavailable, gin.H{"code": 503, "message": err.Error()})
		return
	}
	if setupErr := middleware.SetupContextForSelectedChannel(c, channel, route.BillingModel); setupErr != nil {
		c.JSON(setupErr.StatusCode, gin.H{"code": setupErr.GetErrorCode(), "message": setupErr.Error()})
		return
	}

	resp, err := doKlingAipingProxyRequest(c, route, channel)
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"code": 502, "message": err.Error()})
		return
	}
	defer resp.Body.Close()
	copyProxyResponse(c, resp)
}

func readJSONPayload(c *gin.Context) (map[string]any, error) {
	body, err := common.GetBodyStorage(c)
	if err != nil {
		return nil, err
	}
	data, err := body.Bytes()
	if err != nil {
		return nil, err
	}
	_, _ = body.Seek(0, io.SeekStart)
	c.Request.Body = io.NopCloser(body)
	payload := map[string]any{}
	if strings.TrimSpace(string(data)) == "" {
		return payload, nil
	}
	if err := common.Unmarshal(data, &payload); err != nil {
		return nil, err
	}
	return payload, nil
}

func resolveKlingAipingModel(payload map[string]any, route klingaiping.Route) string {
	if modelName := stringFromMap(payload, "model_name"); modelName != "" {
		return modelName
	}
	if modelName := stringFromMap(payload, "model"); modelName != "" {
		return modelName
	}
	if route.BillingModel != "" {
		return route.BillingModel
	}
	return "kling-v3"
}

func stringFromMap(payload map[string]any, key string) string {
	if value, ok := payload[key].(string); ok {
		return strings.TrimSpace(value)
	}
	return ""
}

func getKlingAipingUserTask(c *gin.Context, taskID string, action string) (*model.Task, bool) {
	task, exist, err := model.GetByTaskId(c.GetInt("id"), taskID)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"code": 500, "message": err.Error()})
		return nil, false
	}
	if !exist || task.Platform != constant.TaskPlatform(strconv.Itoa(constant.ChannelTypeKlingAiping)) || task.Action != action {
		c.JSON(http.StatusBadRequest, gin.H{"code": 400, "message": "task not found"})
		return nil, false
	}
	return task, true
}

func refreshKlingAipingTaskIfNeeded(task *model.Task) {
	if task.Status == model.TaskStatusSuccess || task.Status == model.TaskStatusFailure {
		return
	}
	channelModel, err := model.GetChannelById(task.ChannelId, true)
	if err != nil || channelModel == nil {
		return
	}
	adaptor := relay.GetTaskAdaptor(task.Platform)
	if adaptor == nil {
		return
	}
	resp, err := adaptor.FetchTask(channelModel.GetBaseURL(), channelModel.Key, map[string]any{
		"task_id": task.GetUpstreamTaskID(),
		"action":  task.Action,
	}, channelModel.GetSetting().Proxy)
	if err != nil || resp == nil {
		return
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil || resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return
	}
	taskInfo, err := adaptor.ParseTaskResult(body)
	if err != nil || taskInfo == nil {
		return
	}
	snap := task.Snapshot()
	task.Data = body
	if taskInfo.Status != "" {
		task.Status = model.TaskStatus(taskInfo.Status)
	}
	if taskInfo.Progress != "" {
		task.Progress = taskInfo.Progress
	}
	if taskInfo.Url != "" {
		task.PrivateData.ResultURL = taskInfo.Url
	}
	if !snap.Equal(task.Snapshot()) {
		_, _ = task.UpdateWithStatus(snap.Status)
	}
}

func buildKlingAipingTaskPayload(task *model.Task) map[string]any {
	payload := map[string]any{
		"code":    0,
		"message": "success",
		"data":    taskDataObject(task),
	}
	return payload
}

func taskDataObject(task *model.Task) map[string]any {
	payload := map[string]any{}
	_ = common.Unmarshal(task.Data, &payload)
	delete(payload, "aiping_id")
	data, _ := payload["data"].(map[string]any)
	if data == nil {
		data = map[string]any{}
	}
	data["task_id"] = task.TaskID
	if _, ok := data["task_status"]; !ok {
		data["task_status"] = mapKlingAipingTaskStatus(task.Status)
	}
	if _, ok := data["task_status_msg"]; !ok {
		data["task_status_msg"] = task.FailReason
	}
	if _, ok := data["created_at"]; !ok && task.CreatedAt != 0 {
		data["created_at"] = task.CreatedAt
	}
	if _, ok := data["updated_at"]; !ok && task.UpdatedAt != 0 {
		data["updated_at"] = task.UpdatedAt
	}
	ensureKlingAipingWatermarkURL(data)
	return data
}

func mapKlingAipingTaskStatus(status model.TaskStatus) string {
	switch status {
	case model.TaskStatusSubmitted, model.TaskStatusQueued:
		return "submitted"
	case model.TaskStatusInProgress:
		return "processing"
	case model.TaskStatusSuccess:
		return "succeed"
	case model.TaskStatusFailure:
		return "failed"
	default:
		return "processing"
	}
}

func ensureKlingAipingWatermarkURL(data map[string]any) {
	taskResult, _ := data["task_result"].(map[string]any)
	if taskResult == nil {
		return
	}
	videos, _ := taskResult["videos"].([]any)
	for _, videoAny := range videos {
		video, _ := videoAny.(map[string]any)
		if video == nil {
			continue
		}
		if _, ok := video["watermark_url"]; !ok {
			video["watermark_url"] = ""
		}
	}
}

func parseKlingAipingPage(c *gin.Context) (int, int, error) {
	pageNum := parseIntDefault(c.Query("pageNum"), 1)
	pageSize := parseIntDefault(c.Query("pageSize"), 30)
	if pageNum < 1 || pageNum > 1000 {
		return 0, 0, fmt.Errorf("pageNum must be in [1, 1000]")
	}
	if pageSize < 1 || pageSize > 500 {
		return 0, 0, fmt.Errorf("pageSize must be in [1, 500]")
	}
	return pageNum, pageSize, nil
}

func parseIntDefault(raw string, fallback int) int {
	if strings.TrimSpace(raw) == "" {
		return fallback
	}
	v, err := strconv.Atoi(raw)
	if err != nil {
		return -1
	}
	return v
}

func doKlingAipingProxyRequest(c *gin.Context, route klingaiping.Route, channelModel *model.Channel) (*http.Response, error) {
	baseURL := strings.TrimRight(channelModel.GetBaseURL(), "/")
	upstreamPath := strings.Replace(route.UpstreamPath, ":id", c.Param("id"), 1)
	url := baseURL + upstreamPath
	if c.Request.URL.RawQuery != "" {
		url += "?" + c.Request.URL.RawQuery
	}
	var body io.Reader
	if c.Request.Method != http.MethodGet {
		data, err := proxyBodyBytes(c, route)
		if err != nil {
			return nil, err
		}
		body = bytes.NewReader(data)
	}
	req, err := http.NewRequest(c.Request.Method, url, body)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Content-Type", "application/json")
	key := common.GetContextKeyString(c, constant.ContextKeyChannelKey)
	if key == "" {
		key = channelModel.Key
	}
	req.Header.Set("Authorization", "Bearer "+key)
	client, err := service.GetHttpClientWithProxy(channelModel.GetSetting().Proxy)
	if err != nil {
		return nil, err
	}
	return client.Do(req)
}

func proxyBodyBytes(c *gin.Context, route klingaiping.Route) ([]byte, error) {
	storage, err := common.GetBodyStorage(c)
	if err != nil {
		return nil, err
	}
	data, err := storage.Bytes()
	if err != nil {
		return nil, err
	}
	if route.PathPattern == "/v1/elements/delete" {
		payload := map[string]any{}
		if err := common.Unmarshal(data, &payload); err == nil {
			if _, ok := payload["element_id"]; !ok {
				if id := stringFromMap(payload, "id"); id != "" {
					payload["element_id"] = id
					delete(payload, "id")
					data, _ = common.Marshal(payload)
				}
			}
		}
	}
	return data, nil
}

func copyProxyResponse(c *gin.Context, resp *http.Response) {
	if resp.StatusCode >= http.StatusBadRequest {
		copyNormalizedProxyError(c, resp)
		return
	}
	for key, values := range resp.Header {
		for _, value := range values {
			c.Writer.Header().Add(key, value)
		}
	}
	c.Status(resp.StatusCode)
	_, _ = io.Copy(c.Writer, resp.Body)
}

func copyNormalizedProxyError(c *gin.Context, resp *http.Response) {
	body, _ := io.ReadAll(resp.Body)
	message := strings.TrimSpace(string(body))
	payload := map[string]any{}
	if len(body) > 0 && common.Unmarshal(body, &payload) == nil {
		if msg := stringFromMap(payload, "message"); msg != "" {
			message = msg
		} else if msg := stringFromMap(payload, "msg"); msg != "" {
			message = msg
		} else if detail, ok := payload["detail"]; ok {
			if detailMap, ok := detail.(map[string]any); ok {
				if msg := stringFromMap(detailMap, "message"); msg != "" {
					message = msg
				} else if msg := stringFromMap(detailMap, "msg"); msg != "" {
					message = msg
				} else {
					message = fmt.Sprint(detail)
				}
			} else {
				message = fmt.Sprint(detail)
			}
		}
		delete(payload, "msg")
	} else {
		payload = map[string]any{}
	}
	if message == "" {
		message = resp.Status
	}
	if _, ok := payload["code"]; !ok {
		payload["code"] = resp.StatusCode
	}
	payload["message"] = message
	payload["request_id"] = c.GetString(common.RequestIdKey)
	c.JSON(resp.StatusCode, payload)
}

func normalizeKlingAipingTaskError(taskErr *dto.TaskError) {
	if taskErr == nil || strings.TrimSpace(taskErr.Message) == "" {
		return
	}
	payload := map[string]any{}
	if common.Unmarshal([]byte(taskErr.Message), &payload) != nil {
		return
	}
	if msg := stringFromMap(payload, "message"); msg != "" {
		taskErr.Message = msg
		return
	}
	if msg := stringFromMap(payload, "msg"); msg != "" {
		taskErr.Message = msg
		return
	}
	if detail, ok := payload["detail"].(map[string]any); ok {
		if msg := stringFromMap(detail, "message"); msg != "" {
			taskErr.Message = msg
		}
	}
}