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.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()
	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 []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
			}
			allModels = append(allModels, models...)
			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 spec.Capabilities.Sync != nil {
			caps["sync"] = true
		}
		if spec.Capabilities.Async != nil {
			caps["async"] = true
		}
		out[string(st)] = gin.H{
			"capabilities":  caps,
			"providers":     provList,
			"models":        allModels,
			"default_model": provider.DefaultModel(allModels),
		}
	}
	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) {
	var in struct {
		Title string `json:"title"`
		Model string `json:"model"`
	}
	_ = c.ShouldBindJSON(&in)
	if in.Title == "" {
		in.Title = "新对话"
	}
	conv := &store.Conversation{SessionID: sessionID(c), Title: in.Title, Model: in.Model, ParamsJSON: "{}"}
	id, err := s.app.Store.CreateConversation(c.Request.Context(), conv)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	conv.ID = id
	c.JSON(http.StatusOK, conv)
}

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)
	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
	}
	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(c.Request.Context(), chat.StreamRequest{
		SessionID:      sessionID(c),
		ConversationID: id,
		Content:        in.Content,
		Model:          in.Model,
		ProviderID:     in.ProviderID,
	}, c.Writer, func() { c.Writer.Flush() })
	if err != nil {
		// headers already sent for SSE
		_, _ = 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"`
		Mode        string          `json:"mode"`
		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
	}
	mode := domain.InvokeMode(in.Mode)
	if mode == "" {
		mode = domain.ModeSync
	}
	result, err := s.app.Invoke.Run(c.Request.Context(), domain.InvokeRequest{
		SessionID:   sessionID(c),
		ServiceType: domain.ServiceType(in.ServiceType),
		ProviderID:  in.ProviderID,
		Model:       in.Model,
		Params:      in.Params,
		Mode:        mode,
	})
	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})
}