Atlas/internal/api/user_handlers.go

package api

import (
	"encoding/json"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/gin-gonic/gin"

	"github.com/rose_cat707/Atlas/internal/domain"
	"github.com/rose_cat707/Atlas/internal/provider"
	"github.com/rose_cat707/Atlas/internal/service/asset"
	"github.com/rose_cat707/Atlas/internal/service/chat"
	"github.com/rose_cat707/Atlas/internal/store"
)

func (s *Server) registerUserRoutes(r gin.IRoutes) {
	r.Use(s.sessionMiddleware())
	r.POST("/user/session", s.userSession)
	r.GET("/user/meta", s.userMeta)
	r.GET("/user/conversations", s.listConversations)
	r.POST("/user/conversations", s.createConversation)
	r.POST("/user/conversations/messages", s.postNewConversationMessage)
	r.GET("/user/conversations/:id", s.getConversation)
	r.PUT("/user/conversations/:id", s.updateConversation)
	r.DELETE("/user/conversations/:id", s.deleteConversation)
	r.GET("/user/conversations/:id/messages", s.listMessages)
	r.POST("/user/conversations/:id/messages", s.postMessage)
	r.POST("/user/generate", s.userGenerate)
	r.POST("/user/upload", s.userUpload)
	r.GET("/user/files/:id", s.serveFile)
	r.GET("/user/generations", s.listGenerations)
	r.GET("/user/generations/:id", s.getGeneration)
	r.DELETE("/user/generations/:id", s.deleteGeneration)
	r.GET("/user/presets", s.listPresets)
	r.POST("/user/presets", s.createPreset)
	r.PUT("/user/presets/:id", s.updatePreset)
	r.DELETE("/user/presets/:id", s.deletePreset)
	r.GET("/user/tasks/active", s.listActiveTasks)
	r.GET("/user/tasks/:id/events", s.streamTaskEvents)
	r.GET("/user/tasks/:id", s.getTask)
	r.POST("/user/tasks/:id/cancel", s.cancelTask)
}

func (s *Server) userSession(c *gin.Context) {
	c.JSON(http.StatusOK, gin.H{"session_id": sessionID(c)})
}

func (s *Server) userMeta(c *gin.Context) {
	ctx := c.Request.Context()
	specs := s.app.Invoke.Specs()
	imageModels, _ := s.app.Models.ListImageModelsForMeta(ctx)
	out := gin.H{}
	for st, spec := range specs {
		providers, err := s.app.Store.ListProviders(ctx, string(st), true)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		}
		var provList []gin.H
		var allModels any
		switch st {
		case domain.ServiceImageGeneration:
			allModels = imageModels
		default:
			var entries []provider.ModelEntry
			for _, p := range providers {
				cfg, err := s.app.Router.Resolve(c.Request.Context(), string(st), p.ID)
				if err != nil {
					continue
				}
				models, err := s.app.Models.ListEffective(ctx, cfg)
				if err != nil {
					continue
				}
				entries = append(entries, models...)
			}
			allModels = entries
		}
		for _, p := range providers {
			provList = append(provList, gin.H{"id": p.ID, "name": p.Name, "is_default": p.IsDefault})
		}
		caps := gin.H{}
		if spec.Capabilities.Stream != nil {
			caps["stream"] = true
		}
		if st == domain.ServiceImageGeneration {
			if spec.Capabilities.Async != nil {
				caps["async"] = true
			}
		} else {
			if spec.Capabilities.Sync != nil {
				caps["sync"] = true
			}
			if spec.Capabilities.Async != nil {
				caps["async"] = true
			}
		}
		defaultModel := ""
		switch models := allModels.(type) {
		case []provider.ModelDetail:
			defaultModel = provider.DefaultModelDetail(models)
		case []provider.ModelEntry:
			defaultModel = provider.DefaultModel(models)
		}
		out[string(st)] = gin.H{
			"capabilities":  caps,
			"providers":     provList,
			"models":        allModels,
			"default_model": defaultModel,
		}
	}
	c.JSON(http.StatusOK, gin.H{"service_types": out})
}

func (s *Server) listConversations(c *gin.Context) {
	items, err := s.app.Store.ListConversations(c.Request.Context(), sessionID(c))
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"items": items})
}

func (s *Server) createConversation(c *gin.Context) {
	c.JSON(http.StatusBadRequest, gin.H{"error": "conversation is created when sending the first message"})
}

func (s *Server) postNewConversationMessage(c *gin.Context) {
	s.streamConversationMessage(c, 0)
}

func (s *Server) getConversation(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	conv, err := s.app.Store.GetConversation(c.Request.Context(), sessionID(c), id)
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}
	c.JSON(http.StatusOK, conv)
}

func (s *Server) updateConversation(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	var in store.Conversation
	if err := c.ShouldBindJSON(&in); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	in.ID = id
	in.SessionID = sessionID(c)
	if err := s.app.Store.UpdateConversation(c.Request.Context(), &in); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, in)
}

func (s *Server) deleteConversation(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	if err := s.app.Store.DeleteConversation(c.Request.Context(), sessionID(c), id); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"ok": true})
}

func (s *Server) listMessages(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	items, err := s.app.Store.ListMessages(c.Request.Context(), id, 80)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"items": items})
}

func (s *Server) postMessage(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	s.streamConversationMessage(c, id)
}

func (s *Server) streamConversationMessage(c *gin.Context, convID int64) {
	var in struct {
		Content    string `json:"content"`
		Model      string `json:"model"`
		ProviderID int64  `json:"provider_id"`
	}
	if err := c.ShouldBindJSON(&in); err != nil || in.Content == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "content required"})
		return
	}

	ctx := c.Request.Context()
	sid := sessionID(c)
	if convID == 0 {
		conv := &store.Conversation{
			SessionID:  sid,
			Title:      "新对话",
			Model:      in.Model,
			ParamsJSON: "{}",
		}
		id, err := s.app.Store.CreateConversation(ctx, conv)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		}
		convID = id
	} else if _, err := s.app.Store.GetConversation(ctx, sid, convID); err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}

	c.Header("X-Conversation-Id", strconv.FormatInt(convID, 10))
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Status(http.StatusOK)
	err := s.app.Chat.StreamMessage(ctx, chat.StreamRequest{
		SessionID:      sid,
		ConversationID: convID,
		Content:        in.Content,
		Model:          in.Model,
		ProviderID:     in.ProviderID,
	}, c.Writer, func() { c.Writer.Flush() })
	if err != nil {
		_, _ = c.Writer.Write([]byte("data: " + `{"error":"` + err.Error() + `"}` + "\n\n"))
		c.Writer.Flush()
	}
}

func (s *Server) userGenerate(c *gin.Context) {
	var in struct {
		ServiceType string          `json:"service_type"`
		Params      json.RawMessage `json:"params"`
		Model       string          `json:"model"`
		ProviderID  int64           `json:"provider_id"`
	}
	if err := c.ShouldBindJSON(&in); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	st := domain.ServiceType(in.ServiceType)
	if !domain.IsImageInvokeService(st) {
		c.JSON(http.StatusBadRequest, gin.H{"error": "unsupported service_type"})
		return
	}
	result, err := s.app.Invoke.Run(c.Request.Context(), domain.InvokeRequest{
		SessionID:   sessionID(c),
		ServiceType: st,
		ProviderID:  in.ProviderID,
		Model:       in.Model,
		Params:      in.Params,
		Mode:        domain.ModeAsync,
	})
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, result)
}

func (s *Server) userUpload(c *gin.Context) {
	file, err := c.FormFile("file")
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "file required"})
		return
	}
	f, err := file.Open()
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	defer f.Close()
	path, size, mimeType, err := asset.SaveUpload(s.app.DataDir, sessionID(c), file.Filename, f, 10<<20)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	exp := time.Now().Add(7 * 24 * time.Hour).UTC().Format(time.RFC3339)
	rec := &store.FileRecord{
		SessionID:   sessionID(c),
		Kind:        "upload",
		Mime:        mimeType,
		StoragePath: path,
		SizeBytes:   size,
		ExpiresAt:   &exp,
	}
	id, err := s.app.Store.InsertFile(c.Request.Context(), rec)
	if err != nil {
		os.Remove(path)
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	signed := s.app.Assets.SignedURL("/api/v1", id, 7*24*time.Hour)
	c.JSON(http.StatusOK, gin.H{"file_id": id, "signed_url": signed, "mime": mimeType})
}

func (s *Server) serveFile(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	expStr := c.Query("exp")
	sig := c.Query("sig")
	exp, sig, err := asset.ParseSignedQuery(expStr, sig)
	if err != nil || !s.app.Assets.Verify(id, exp, sig) {
		c.JSON(http.StatusForbidden, gin.H{"error": "invalid signature"})
		return
	}
	rec, err := s.app.Store.GetFile(c.Request.Context(), sessionID(c), id)
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}
	c.Header("Content-Type", rec.Mime)
	c.File(rec.StoragePath)
}

func (s *Server) listGenerations(c *gin.Context) {
	limit, offset := parseLimitOffset(c, 20, 100)
	st := c.Query("service_type")
	items, total, err := s.app.Store.ListGenerations(c.Request.Context(), sessionID(c), st, limit, offset)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"items": items, "total": total})
}

func (s *Server) getGeneration(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	g, err := s.app.Store.GetGeneration(c.Request.Context(), sessionID(c), id)
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}
	c.JSON(http.StatusOK, g)
}

func (s *Server) deleteGeneration(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	if err := s.app.Store.DeleteGeneration(c.Request.Context(), sessionID(c), id); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"ok": true})
}

func (s *Server) listPresets(c *gin.Context) {
	items, err := s.app.Store.ListPresets(c.Request.Context(), sessionID(c), c.Query("service_type"))
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"items": items})
}

func (s *Server) createPreset(c *gin.Context) {
	var in store.Preset
	if err := c.ShouldBindJSON(&in); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	in.SessionID = sessionID(c)
	id, err := s.app.Store.CreatePreset(c.Request.Context(), &in)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	in.ID = id
	c.JSON(http.StatusOK, in)
}

func (s *Server) updatePreset(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	var in store.Preset
	if err := c.ShouldBindJSON(&in); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	in.ID = id
	in.SessionID = sessionID(c)
	if err := s.app.Store.UpdatePreset(c.Request.Context(), &in); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, in)
}

func (s *Server) deletePreset(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	if err := s.app.Store.DeletePreset(c.Request.Context(), sessionID(c), id); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"ok": true})
}

func (s *Server) listActiveTasks(c *gin.Context) {
	items, err := s.app.Store.ListActiveTasks(c.Request.Context(), sessionID(c), c.Query("service_type"))
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"items": items})
}

func (s *Server) getTask(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	if v, ok := s.app.Store.Cache.Get("task:" + strconv.FormatInt(id, 10)); ok {
		c.JSON(http.StatusOK, v)
		return
	}
	t, err := s.app.Store.GetTask(c.Request.Context(), sessionID(c), id)
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}
	c.JSON(http.StatusOK, t)
}

func (s *Server) cancelTask(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	if err := s.app.Invoke.CancelTask(c.Request.Context(), sessionID(c), id); err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"ok": true})
}