easyai-ai-gateway/apps/api/internal/httpapi/gemini_compat.go

package httpapi

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"mime"
	"net/http"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/runner"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

type geminiImageTaskMapping struct {
	Kind  string
	Body  map[string]any
	Model string
}

type geminiImageInput struct {
	URI      string
	MimeType string
}

type geminiUploadSession struct {
	Version     string
	DisplayName string
	MimeType    string
	SizeBytes   int64
	CreatedAt   time.Time
}

var geminiGenerateContentRoutePrefixes = []string{
	"/v1beta/models/",
	"/v1/models/",
	"/models/",
}

func (s *Server) registerGeminiGenerateContentRoutes(mux *http.ServeMux) {
	handler := s.auth.Require(auth.PermissionBasic, http.HandlerFunc(s.geminiGenerateContent))
	for _, prefix := range geminiGenerateContentRoutePrefixes {
		mux.Handle("POST "+prefix, geminiGenerateContentRouteHandler(prefix, handler))
	}
}

func geminiGenerateContentRouteHandler(prefix string, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		model, ok := geminiGenerateContentModelFromPath(prefix, r.URL.Path)
		if !ok {
			http.NotFound(w, r)
			return
		}
		r.SetPathValue("model", model)
		next.ServeHTTP(w, r)
	})
}

func geminiGenerateContentModelFromPath(prefix string, requestPath string) (string, bool) {
	if !strings.HasPrefix(requestPath, prefix) {
		return "", false
	}
	model, ok := strings.CutSuffix(strings.TrimPrefix(requestPath, prefix), ":generateContent")
	if !ok || strings.TrimSpace(model) == "" || strings.Contains(model, "/") {
		return "", false
	}
	return model, true
}

func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
	user, ok := auth.UserFromContext(r.Context())
	if !ok {
		writeError(w, http.StatusUnauthorized, "unauthorized")
		return
	}
	var native map[string]any
	if err := json.NewDecoder(r.Body).Decode(&native); err != nil {
		writeError(w, http.StatusBadRequest, "invalid json body", "invalid_json_body")
		return
	}
	mapping, err := geminiImageTaskBody(r.PathValue("model"), native)
	if err != nil {
		writeError(w, http.StatusBadRequest, err.Error(), clients.ErrorCode(err))
		return
	}
	if !apiKeyScopeAllowed(user, mapping.Kind) {
		writeError(w, http.StatusForbidden, "api key scope does not allow this capability")
		return
	}
	prepared, err := s.prepareTaskRequest(r.Context(), r, user, mapping.Body)
	if err != nil {
		s.logger.Warn("prepare gemini task request failed", "kind", mapping.Kind, "error", err)
		status := http.StatusBadRequest
		if code := clients.ErrorCode(err); strings.HasPrefix(code, "upload_") || code == "request_asset_upload_failed" {
			status = http.StatusBadGateway
		}
		writeError(w, status, err.Error(), clients.ErrorCode(err))
		return
	}
	task, err := s.store.CreateTask(r.Context(), store.CreateTaskInput{
		Kind:            mapping.Kind,
		Model:           mapping.Model,
		RunMode:         runModeFromRequest(prepared.Body),
		Async:           false,
		Request:         prepared.Body,
		ConversationID:  prepared.ConversationID,
		NewMessageCount: prepared.NewMessageCount,
		MessageRefs:     prepared.MessageRefs,
	}, user)
	if err != nil {
		s.logger.Error("create gemini task failed", "kind", mapping.Kind, "error", err)
		writeError(w, http.StatusInternalServerError, "create task failed")
		return
	}
	runCtx, cancelRun := s.requestExecutionContext(r)
	defer cancelRun()
	result, runErr := s.runner.Execute(runCtx, task, user)
	if runErr != nil {
		if !requestStillConnected(r) {
			return
		}
		writeErrorWithDetails(w, statusFromRunError(runErr), runErrorMessage(runErr), runErrorDetails(runErr), runErrorCode(runErr))
		return
	}
	if !requestStillConnected(r) {
		return
	}
	writeJSON(w, http.StatusOK, geminiGenerateContentResponse(result.Output, mapping.Model))
}

func geminiImageTaskBody(model string, native map[string]any) (geminiImageTaskMapping, error) {
	model = strings.TrimSpace(model)
	if model == "" {
		return geminiImageTaskMapping{}, &clients.ClientError{Code: "bad_request", Message: "model is required", Retryable: false}
	}
	if native == nil {
		native = map[string]any{}
	}
	prompt := geminiPromptFromNative(native)
	imageInputs := geminiImageInputsFromNative(native)
	body := cloneMap(native)
	body["model"] = model
	body["prompt"] = prompt
	if count := geminiCandidateCount(native); count > 0 {
		body["n"] = count
	}
	if aspectRatio, imageSize := geminiImageConfig(native); aspectRatio != "" {
		body["aspect_ratio"] = aspectRatio
	} else if imageSize != "" {
		body["resolution"] = imageSize
	}
	if _, imageSize := geminiImageConfig(native); imageSize != "" {
		body["resolution"] = imageSize
	}
	if len(imageInputs) == 0 {
		body["modelType"] = "image_generate"
		return geminiImageTaskMapping{Kind: "images.generations", Body: body, Model: model}, nil
	}
	images := make([]any, 0, len(imageInputs))
	for _, input := range imageInputs {
		images = append(images, input.URI)
	}
	if len(images) == 1 {
		body["image"] = images[0]
	} else {
		body["image"] = images
	}
	body["modelType"] = "image_edit"
	return geminiImageTaskMapping{Kind: "images.edits", Body: body, Model: model}, nil
}

func geminiPromptFromNative(body map[string]any) string {
	parts := []string{}
	for _, part := range geminiOrderedNativeParts(body) {
		if text := strings.TrimSpace(stringFromRequestAny(part["text"])); text != "" {
			parts = append(parts, text)
		}
	}
	if len(parts) > 0 {
		return strings.Join(parts, "\n")
	}
	return firstNonEmptyRequestString(body, "prompt", "positive", "input")
}

func geminiImageInputsFromNative(body map[string]any) []geminiImageInput {
	inputs := []geminiImageInput{}
	for _, part := range geminiOrderedNativeParts(body) {
		inline := mapFromGeminiAny(firstPresentGemini(part["inlineData"], part["inline_data"]))
		if data := strings.TrimSpace(stringFromRequestAny(inline["data"])); data != "" {
			mimeType := firstNonEmptyGeminiString(inline["mimeType"], inline["mime_type"])
			if mimeType == "" {
				mimeType = "image/png"
			}
			if !strings.HasPrefix(strings.ToLower(data), "data:") {
				data = "data:" + mimeType + ";base64," + data
			}
			inputs = append(inputs, geminiImageInput{URI: data, MimeType: mimeType})
			continue
		}
		fileData := mapFromGeminiAny(firstPresentGemini(part["fileData"], part["file_data"]))
		if fileURI := firstNonEmptyGeminiString(fileData["fileUri"], fileData["file_uri"], fileData["uri"]); fileURI != "" {
			inputs = append(inputs, geminiImageInput{
				URI:      fileURI,
				MimeType: firstNonEmptyGeminiString(fileData["mimeType"], fileData["mime_type"]),
			})
		}
	}
	return inputs
}

func geminiOrderedNativeParts(body map[string]any) []map[string]any {
	contents := normalizeGeminiContents(body["contents"])
	parts := []map[string]any{}
	for _, content := range contents {
		contentMap := mapFromGeminiAny(content)
		rawParts, _ := contentMap["parts"].([]any)
		if len(rawParts) == 0 {
			if text := stringFromRequestAny(contentMap["text"]); text != "" {
				rawParts = []any{map[string]any{"text": text}}
			}
		}
		for _, rawPart := range rawParts {
			switch typed := rawPart.(type) {
			case string:
				parts = append(parts, map[string]any{"text": typed})
			case map[string]any:
				parts = append(parts, typed)
			}
		}
	}
	return parts
}

func normalizeGeminiContents(value any) []any {
	switch typed := value.(type) {
	case []any:
		out := make([]any, 0, len(typed))
		for _, item := range typed {
			if text, ok := item.(string); ok {
				out = append(out, map[string]any{"parts": []any{map[string]any{"text": text}}})
			} else {
				out = append(out, item)
			}
		}
		return out
	case string:
		return []any{map[string]any{"parts": []any{map[string]any{"text": typed}}}}
	case map[string]any:
		return []any{typed}
	default:
		return nil
	}
}

func geminiCandidateCount(body map[string]any) int {
	generationConfig := mapFromGeminiAny(firstPresentGemini(body["generationConfig"], body["generation_config"]))
	for _, value := range []any{body["n"], body["candidateCount"], body["candidate_count"], generationConfig["n"], generationConfig["candidateCount"], generationConfig["candidate_count"]} {
		if count := positiveIntFromGeminiAny(value); count > 0 {
			return count
		}
	}
	return 0
}

func geminiImageConfig(body map[string]any) (string, string) {
	generationConfig := mapFromGeminiAny(firstPresentGemini(body["generationConfig"], body["generation_config"]))
	imageConfig := mapFromGeminiAny(firstPresentGemini(generationConfig["imageConfig"], generationConfig["image_config"]))
	return firstNonEmptyGeminiString(imageConfig["aspectRatio"], imageConfig["aspect_ratio"]),
		firstNonEmptyGeminiString(imageConfig["imageSize"], imageConfig["image_size"])
}

func geminiGenerateContentResponse(output map[string]any, model string) map[string]any {
	parts := geminiPartsFromOutput(output)
	response := map[string]any{
		"candidates": []any{map[string]any{
			"index": 0,
			"content": map[string]any{
				"role":  "model",
				"parts": parts,
			},
			"finishReason": "STOP",
		}},
		"modelVersion": model,
	}
	if usage := geminiUsageMetadataFromOutput(output); len(usage) > 0 {
		response["usageMetadata"] = usage
	}
	return response
}

func geminiPartsFromOutput(output map[string]any) []any {
	parts := []any{}
	data, _ := output["data"].([]any)
	for _, rawItem := range data {
		item, _ := rawItem.(map[string]any)
		if item == nil {
			continue
		}
		if text := firstNonEmptyGeminiString(item["text"], item["content"]); text != "" {
			parts = append(parts, map[string]any{"text": text})
			continue
		}
		mimeType := firstNonEmptyGeminiString(item["mime_type"], item["mimeType"])
		if b64 := strings.TrimSpace(stringFromRequestAny(item["b64_json"])); b64 != "" {
			if parsed := parseGeminiDataURL(b64); parsed != nil {
				mimeType = firstNonEmptyGeminiString(mimeType, parsed.MimeType)
				b64 = parsed.Data
			}
			parts = append(parts, map[string]any{"inlineData": map[string]any{
				"mimeType": firstNonEmptyGeminiString(mimeType, "image/png"),
				"data":     b64,
			}})
			continue
		}
		upload := mapFromGeminiAny(item["upload"])
		if uri := firstNonEmptyGeminiString(item["url"], item["image_url"], upload["url"]); uri != "" {
			parts = append(parts, map[string]any{"fileData": map[string]any{
				"mimeType": firstNonEmptyGeminiString(mimeType, mimeTypeFromGeminiURI(uri)),
				"fileUri":  uri,
			}})
		}
	}
	if len(parts) == 0 {
		if outputs, ok := output["output"].([]any); ok {
			for _, raw := range outputs {
				if uri := stringFromRequestAny(raw); uri != "" {
					parts = append(parts, map[string]any{"fileData": map[string]any{
						"mimeType": mimeTypeFromGeminiURI(uri),
						"fileUri":  uri,
					}})
				}
			}
		}
	}
	if len(parts) == 0 {
		parts = append(parts, map[string]any{"text": firstNonEmptyGeminiString(output["message"])})
	}
	return parts
}

func geminiUsageMetadataFromOutput(output map[string]any) map[string]any {
	usage := mapFromGeminiAny(output["usage"])
	prompt := positiveIntFromGeminiAny(firstPresentGemini(usage["prompt_tokens"], usage["input_tokens"]))
	candidates := positiveIntFromGeminiAny(firstPresentGemini(usage["completion_tokens"], usage["output_tokens"]))
	total := positiveIntFromGeminiAny(usage["total_tokens"])
	if total == 0 && (prompt > 0 || candidates > 0) {
		total = prompt + candidates
	}
	meta := map[string]any{}
	if prompt > 0 {
		meta["promptTokenCount"] = prompt
	}
	if candidates > 0 {
		meta["candidatesTokenCount"] = candidates
	}
	if total > 0 {
		meta["totalTokenCount"] = total
	}
	return meta
}

func (s *Server) geminiFilesUpload(w http.ResponseWriter, r *http.Request) {
	if err := validateGeminiFilesVersion(r.PathValue("version")); err != nil {
		writeError(w, http.StatusBadRequest, err.Error(), clients.ErrorCode(err))
		return
	}
	if geminiUploadCommandHas(r, "start") {
		s.startGeminiFilesUpload(w, r)
		return
	}
	s.finalizeGeminiFilesUpload(w, r, "", geminiUploadSession{
		Version:     r.PathValue("version"),
		DisplayName: "gemini-upload-" + newGeminiUploadID(),
		MimeType:    r.Header.Get("Content-Type"),
		CreatedAt:   time.Now(),
	})
}

func (s *Server) geminiFilesUploadFinalize(w http.ResponseWriter, r *http.Request) {
	if err := validateGeminiFilesVersion(r.PathValue("version")); err != nil {
		writeError(w, http.StatusBadRequest, err.Error(), clients.ErrorCode(err))
		return
	}
	uploadID := strings.TrimSpace(r.PathValue("uploadID"))
	value, ok := s.geminiUploadSessions.LoadAndDelete(uploadID)
	if !ok {
		writeError(w, http.StatusBadRequest, "Gemini upload session not found", "bad_request")
		return
	}
	session, _ := value.(geminiUploadSession)
	s.finalizeGeminiFilesUpload(w, r, uploadID, session)
}

func (s *Server) startGeminiFilesUpload(w http.ResponseWriter, r *http.Request) {
	var body map[string]any
	_ = json.NewDecoder(r.Body).Decode(&body)
	fileMeta := mapFromGeminiAny(body["file"])
	uploadID := newGeminiUploadID()
	session := geminiUploadSession{
		Version:     r.PathValue("version"),
		DisplayName: firstNonEmptyGeminiString(fileMeta["display_name"], fileMeta["displayName"], "gemini-upload-"+uploadID),
		MimeType:    firstNonEmptyGeminiString(r.Header.Get("X-Goog-Upload-Header-Content-Type"), "application/octet-stream"),
		SizeBytes:   int64(positiveIntFromGeminiAny(r.Header.Get("X-Goog-Upload-Header-Content-Length"))),
		CreatedAt:   time.Now(),
	}
	s.geminiUploadSessions.Store(uploadID, session)
	w.Header().Set("X-Goog-Upload-URL", absoluteRequestURL(r, "/upload/"+session.Version+"/files/"+uploadID))
	w.Header().Set("X-Goog-Upload-Status", "active")
	writeJSON(w, http.StatusOK, map[string]any{})
}

func (s *Server) finalizeGeminiFilesUpload(w http.ResponseWriter, r *http.Request, uploadID string, session geminiUploadSession) {
	if uploadID == "" {
		uploadID = newGeminiUploadID()
	}
	r.Body = http.MaxBytesReader(w, r.Body, maxGatewayUploadBytes)
	payload, err := io.ReadAll(r.Body)
	if err != nil || len(payload) == 0 {
		writeError(w, http.StatusBadRequest, "file payload is required", "bad_request")
		return
	}
	mimeType := firstNonEmptyGeminiString(session.MimeType, http.DetectContentType(payload), "application/octet-stream")
	upload, err := s.runner.UploadFile(r.Context(), runner.FileUploadPayload{
		Bytes:       payload,
		ContentType: mimeType,
		FileName:    geminiUploadFileName(uploadID, session.DisplayName, mimeType),
		Source:      "gemini-files-api",
		Scene:       store.FileStorageSceneUpload,
	})
	if err != nil {
		s.logger.Error("Gemini files upload failed", "error", err)
		status := http.StatusBadGateway
		if clients.ErrorCode(err) == "upload_no_channel" {
			status = http.StatusServiceUnavailable
		}
		writeError(w, status, err.Error(), clients.ErrorCode(err))
		return
	}
	sizeBytes := session.SizeBytes
	if sizeBytes <= 0 {
		sizeBytes = int64(len(payload))
	}
	writeJSON(w, http.StatusOK, geminiFileResponse("files/"+uploadID, stringFromRequestAny(upload["url"]), mimeType, sizeBytes))
}

func geminiFileResponse(name string, uri string, mimeType string, sizeBytes int64) map[string]any {
	return map[string]any{"file": map[string]any{
		"name":      name,
		"uri":       uri,
		"mimeType":  firstNonEmptyGeminiString(mimeType, "application/octet-stream"),
		"sizeBytes": strconv.FormatInt(sizeBytes, 10),
		"state":     "ACTIVE",
	}}
}

func validateGeminiFilesVersion(version string) error {
	if version == "v1" || version == "v1beta" {
		return nil
	}
	return &clients.ClientError{Code: "bad_request", Message: "unsupported Gemini Files API version", Retryable: false}
}

func geminiUploadCommandHas(r *http.Request, command string) bool {
	for _, item := range strings.Split(r.Header.Get("X-Goog-Upload-Command"), ",") {
		if strings.EqualFold(strings.TrimSpace(item), command) {
			return true
		}
	}
	return false
}

func newGeminiUploadID() string {
	var payload [12]byte
	if _, err := rand.Read(payload[:]); err != nil {
		return strconv.FormatInt(time.Now().UnixNano(), 36)
	}
	return hex.EncodeToString(payload[:])
}

func geminiUploadFileName(uploadID string, displayName string, mimeType string) string {
	name := strings.TrimSpace(strings.ReplaceAll(displayName, " ", ""))
	name = path.Base(name)
	if name == "." || name == "/" {
		name = ""
	}
	if strings.Contains(name, ".") {
		return name
	}
	if name == "" {
		name = "gemini-upload-" + uploadID
	}
	return name + extensionFromGeminiMime(mimeType)
}

func extensionFromGeminiMime(mimeType string) string {
	extensions, err := mime.ExtensionsByType(strings.ToLower(strings.TrimSpace(mimeType)))
	if err == nil && len(extensions) > 0 {
		return extensions[0]
	}
	return ".bin"
}

func absoluteRequestURL(r *http.Request, targetPath string) string {
	proto := strings.TrimSpace(r.Header.Get("X-Forwarded-Proto"))
	if proto == "" {
		proto = "http"
	}
	host := strings.TrimSpace(r.Header.Get("X-Forwarded-Host"))
	if host == "" {
		host = r.Host
	}
	return fmt.Sprintf("%s://%s%s", proto, host, targetPath)
}

type parsedGeminiDataURL struct {
	MimeType string
	Data     string
}

func parseGeminiDataURL(value string) *parsedGeminiDataURL {
	if !strings.HasPrefix(strings.ToLower(value), "data:") {
		return nil
	}
	prefix, data, ok := strings.Cut(value, ",")
	if !ok || !strings.Contains(strings.ToLower(prefix), ";base64") {
		return nil
	}
	mimeType := strings.TrimPrefix(strings.Split(prefix, ";")[0], "data:")
	if mimeType == "" {
		mimeType = "image/png"
	}
	return &parsedGeminiDataURL{MimeType: mimeType, Data: data}
}

func mimeTypeFromGeminiURI(value string) string {
	extension := strings.ToLower(path.Ext(strings.Split(value, "?")[0]))
	if extension == "" {
		return "image/png"
	}
	if mimeType := mime.TypeByExtension(extension); mimeType != "" {
		return mimeType
	}
	return "image/png"
}

func mapFromGeminiAny(value any) map[string]any {
	if typed, ok := value.(map[string]any); ok {
		return typed
	}
	return map[string]any{}
}

func firstPresentGemini(values ...any) any {
	for _, value := range values {
		if value != nil {
			return value
		}
	}
	return nil
}

func firstNonEmptyGeminiString(values ...any) string {
	for _, value := range values {
		if text := strings.TrimSpace(stringFromRequestAny(value)); text != "" {
			return text
		}
	}
	return ""
}

func positiveIntFromGeminiAny(value any) int {
	switch typed := value.(type) {
	case int:
		if typed > 0 {
			return typed
		}
	case int64:
		if typed > 0 {
			return int(typed)
		}
	case float64:
		if typed > 0 {
			return int(typed)
		}
	case json.Number:
		if parsed, err := typed.Int64(); err == nil && parsed > 0 {
			return int(parsed)
		}
	case string:
		if parsed, err := strconv.Atoi(strings.TrimSpace(typed)); err == nil && parsed > 0 {
			return parsed
		}
	}
	return 0
}