agent-team/internal/api/server.go

package api

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/labstack/echo/v4"
	"github.com/sdaduanbilei/agent-team/internal/hub"
	"github.com/sdaduanbilei/agent-team/internal/room"
	"github.com/sdaduanbilei/agent-team/internal/user"
)

type Server struct {
	e         *echo.Echo
	agentsDir string
	skillsDir string
	roomsDir  string
	usersDir  string
	teamsDir  string
	rooms     map[string]*room.Room
	user      *user.User
	mu        sync.RWMutex
	clients   map[string]map[*websocket.Conn]bool // roomID -> conns
	clientsMu sync.Mutex
	upgrader  websocket.Upgrader
}

func New(agentsDir, skillsDir, roomsDir, usersDir, teamsDir string) *Server {
	s := &Server{
		e:         echo.New(),
		agentsDir: agentsDir,
		skillsDir: skillsDir,
		roomsDir:  roomsDir,
		usersDir:  usersDir,
		teamsDir:  teamsDir,
		rooms:     make(map[string]*room.Room),
		clients:   make(map[string]map[*websocket.Conn]bool),
		upgrader:  websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }},
	}
	s.loadUser()
	s.loadRooms()
	s.routes()
	return s
}

func (s *Server) Start(addr string) error {
	return s.e.Start(addr)
}

func (s *Server) routes() {
	s.e.Static("/", "web/dist")

	g := s.e.Group("/api")
	g.GET("/rooms", s.listRooms)
	g.POST("/rooms", s.createRoom)
	g.GET("/agents", s.listAgents)
	g.GET("/agents/:name/files/:file", s.readAgentFile)
	g.PUT("/agents/:name/files/:file", s.writeAgentFile)
	g.POST("/agents", s.createAgent)
	g.DELETE("/agents/:name", s.deleteAgent)
	g.POST("/hub/install", s.hubInstall)
	g.POST("/hub/install-zip", s.hubInstallZIP)
	g.GET("/teams", s.listTeams)
	g.GET("/teams/:name", s.getTeam)
	g.GET("/teams/:name/agents/:agent/files/:file", s.getTeamAgentFile)
	g.PUT("/teams/:name/agents/:agent/files/:file", s.saveTeamAgentFile)
	g.GET("/teams/:name/knowledge", s.listTeamKnowledge)
	g.GET("/teams/:name/knowledge/:file", s.getTeamKnowledgeFile)
	g.PUT("/teams/:name/knowledge/:file", s.saveTeamKnowledgeFile)
	g.DELETE("/teams/:name", s.deleteTeam)

	g.GET("/rooms/:id/workspace", s.listWorkspace)
	g.GET("/rooms/:id/workspace/:file", s.getWorkspaceFile)
	g.GET("/rooms/:id/tasks", s.getTasks)
	g.GET("/rooms/:id/history", s.listHistory)
	g.GET("/rooms/:id/messages", s.getMessages)
	g.GET("/user", s.getUser)
	g.GET("/user/profile", s.getUserProfile)
	g.PUT("/user/profile", s.saveUserProfile)
	g.PUT("/user/config", s.saveUserConfig)

	s.e.GET("/ws/:roomID", s.wsHandler)
}

func (s *Server) loadRooms() {
	entries, _ := os.ReadDir(s.roomsDir)
	for _, e := range entries {
		if !e.IsDir() {
			continue
		}
		r, err := room.Load(filepath.Join(s.roomsDir, e.Name()), s.agentsDir, s.skillsDir)
		if err != nil {
			continue
		}
		r.User = s.user
		r.Broadcast = func(ev room.Event) { s.broadcast(ev.RoomID, ev) }
		s.rooms[e.Name()] = r
	}
}

func (s *Server) loadUser() {
	u, err := user.Load(s.usersDir)
	if err != nil {
		u = &user.User{
			Config: user.Config{
				Name:        "用户",
				Provider:    "deepseek",
				Model:       "deepseek-chat",
				APIKeyEnv:   "DEEPSEEK_API_KEY",
				AvatarColor: "#5865F2",
			},
			Profile: "# 我的简介\n\n## 我是谁\n（介绍你的身份、职业背景）\n\n## 工作风格\n- 喜欢的沟通方式\n- 重视的点\n- 不喜欢的内容\n",
			Dir:     filepath.Join(s.usersDir, user.DefaultUser),
		}
		os.MkdirAll(u.Dir, 0755)
		os.WriteFile(filepath.Join(u.Dir, "USER.md"), []byte("---\nname: 用户\ndescription: \nprovider: deepseek\nmodel: deepseek-chat\napi_key_env: DEEPSEEK_API_KEY\navatar_color: \"#5865F2\"\n---\n"), 0644)
		os.WriteFile(filepath.Join(u.Dir, "PROFILE.md"), []byte(u.Profile), 0644)
	}
	s.user = u
}

func (s *Server) broadcast(roomID string, ev room.Event) {
	s.clientsMu.Lock()
	defer s.clientsMu.Unlock()
	data, _ := json.Marshal(ev)
	for conn := range s.clients[roomID] {
		conn.WriteMessage(websocket.TextMessage, data)
	}
}

func (s *Server) wsHandler(c echo.Context) error {
	roomID := c.Param("roomID")
	conn, err := s.upgrader.Upgrade(c.Response(), c.Request(), nil)
	if err != nil {
		return err
	}
	s.clientsMu.Lock()
	if s.clients[roomID] == nil {
		s.clients[roomID] = make(map[*websocket.Conn]bool)
	}
	s.clients[roomID][conn] = true
	s.clientsMu.Unlock()

	defer func() {
		s.clientsMu.Lock()
		delete(s.clients[roomID], conn)
		s.clientsMu.Unlock()
		conn.Close()
	}()

	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			break
		}
		var ev struct {
			Type     string `json:"type"`
			Content  string `json:"content"`
			UserName string `json:"user_name"`
		}
		if json.Unmarshal(msg, &ev) != nil || ev.Type != "user_message" {
			continue
		}
		s.mu.RLock()
		r := s.rooms[roomID]
		s.mu.RUnlock()
		if r == nil {
			continue
		}
		userName := ev.UserName
		if userName == "" {
			userName = s.user.GetName()
		}
		go r.HandleUserMessage(context.Background(), userName, ev.Content)
	}
	return nil
}

// --- REST handlers ---

func (s *Server) listRooms(c echo.Context) error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	type roomInfo struct {
		ID     string      `json:"id"`
		Name   string      `json:"name"`
		Type   string      `json:"type"`
		Status room.Status `json:"status"`
		Master string      `json:"master"`
	}
	var list []roomInfo
	for id, r := range s.rooms {
		list = append(list, roomInfo{ID: id, Name: r.Config.Name, Type: string(r.Config.Type), Status: r.Status, Master: r.Config.Master})
	}
	return c.JSON(200, list)
}

func (s *Server) createRoom(c echo.Context) error {
	var cfg room.Config
	if err := c.Bind(&cfg); err != nil {
		return err
	}
	dir := filepath.Join(s.roomsDir, cfg.Name)
	os.MkdirAll(filepath.Join(dir, "workspace"), 0755)
	os.MkdirAll(filepath.Join(dir, "history"), 0755)

	content := "---\nname: " + cfg.Name + "\ntype: " + string(cfg.Type) + "\nmaster: " + cfg.Master + "\nmembers:\n"
	for _, m := range cfg.Members {
		content += "  - " + m + "\n"
	}
	content += "---\n"
	os.WriteFile(filepath.Join(dir, "room.md"), []byte(content), 0644)

	r, err := room.Load(dir, s.agentsDir, s.skillsDir)
	if err != nil {
		return c.JSON(500, map[string]string{"error": err.Error()})
	}
	r.Broadcast = func(ev room.Event) { s.broadcast(ev.RoomID, ev) }
	s.mu.Lock()
	s.rooms[cfg.Name] = r
	s.mu.Unlock()
	return c.JSON(201, map[string]string{"id": cfg.Name})
}

func (s *Server) listAgents(c echo.Context) error {
	entries, _ := os.ReadDir(s.agentsDir)
	type agentInfo struct {
		Name string `json:"name"`
	}
	var list []agentInfo
	for _, e := range entries {
		if e.IsDir() {
			list = append(list, agentInfo{Name: e.Name()})
		}
	}
	return c.JSON(200, list)
}

func (s *Server) readAgentFile(c echo.Context) error {
	name := c.Param("name")
	file := c.Param("file") // AGENT.md or SOUL.md
	data, err := os.ReadFile(filepath.Join(s.agentsDir, name, file))
	if err != nil {
		return c.JSON(404, map[string]string{"error": "not found"})
	}
	return c.JSON(200, map[string]string{"content": string(data)})
}

func (s *Server) writeAgentFile(c echo.Context) error {
	name := c.Param("name")
	file := c.Param("file")
	var body struct {
		Content string `json:"content"`
	}
	if err := c.Bind(&body); err != nil {
		return err
	}
	dir := filepath.Join(s.agentsDir, name)
	os.MkdirAll(dir, 0755)
	return os.WriteFile(filepath.Join(dir, file), []byte(body.Content), 0644)
}

func (s *Server) createAgent(c echo.Context) error {
	var body struct {
		Name string `json:"name"`
	}
	if err := c.Bind(&body); err != nil {
		return err
	}
	dir := filepath.Join(s.agentsDir, body.Name)
	os.MkdirAll(filepath.Join(dir, "memory"), 0755)
	agentMD := "---\nname: " + body.Name + "\ndescription: \nprovider: deepseek\nmodel: deepseek-chat\napi_key_env: DEEPSEEK_API_KEY\nskills: []\n---\n"
	os.WriteFile(filepath.Join(dir, "AGENT.md"), []byte(agentMD), 0644)
	os.WriteFile(filepath.Join(dir, "SOUL.md"), []byte("You are "+body.Name+"."), 0644)
	return c.JSON(201, map[string]string{"name": body.Name})
}

func (s *Server) deleteAgent(c echo.Context) error {
	name := c.Param("name")
	return os.RemoveAll(filepath.Join(s.agentsDir, name))
}

func (s *Server) hubInstall(c echo.Context) error {
	var body struct {
		Repo string `json:"repo"`
	}
	if err := c.Bind(&body); err != nil {
		return err
	}
	team, err := hub.Install(body.Repo, s.agentsDir, s.skillsDir, s.teamsDir)
	if err != nil {
		return c.JSON(500, map[string]string{"error": err.Error()})
	}
	s.loadRooms()
	return c.JSON(200, team)
}

func (s *Server) hubInstallZIP(c echo.Context) error {
	file, err := c.FormFile("file")
	if err != nil {
		return c.JSON(400, map[string]string{"error": "no file uploaded"})
	}
	src, err := file.Open()
	if err != nil {
		return err
	}
	defer src.Close()

	tmpFile := filepath.Join(os.TempDir(), file.Filename)
	dst, err := os.Create(tmpFile)
	if err != nil {
		return err
	}
	defer os.Remove(tmpFile)
	defer dst.Close()

	if _, err := io.Copy(dst, src); err != nil {
		return err
	}

	team, err := hub.InstallZIP(tmpFile, s.agentsDir, s.skillsDir, s.teamsDir)
	if err != nil {
		return c.JSON(500, map[string]string{"error": err.Error()})
	}
	s.loadRooms()
	return c.JSON(200, team)
}

func (s *Server) listWorkspace(c echo.Context) error {
	id := c.Param("id")
	dir := filepath.Join(s.roomsDir, id, "workspace")
	entries, _ := os.ReadDir(dir)
	var files []string
	for _, e := range entries {
		if !e.IsDir() {
			files = append(files, e.Name())
		}
	}
	return c.JSON(200, files)
}

func (s *Server) getTasks(c echo.Context) error {
	id := c.Param("id")
	data, _ := os.ReadFile(filepath.Join(s.roomsDir, id, "tasks.md"))
	return c.JSON(200, map[string]string{"content": string(data)})
}

func (s *Server) listHistory(c echo.Context) error {
	id := c.Param("id")
	entries, _ := os.ReadDir(filepath.Join(s.roomsDir, id, "history"))
	var files []string
	for _, e := range entries {
		files = append(files, e.Name())
	}
	return c.JSON(200, files)
}

func (s *Server) createSkill(c echo.Context) error {
	var body struct {
		Name    string `json:"name"`
		Content string `json:"content"`
	}
	if err := c.Bind(&body); err != nil {
		return err
	}
	dir := filepath.Join(s.skillsDir, body.Name)
	os.MkdirAll(dir, 0755)
	content := body.Content
	if content == "" {
		content = "---\nname: " + body.Name + "\ndescription: \n---\n\n# " + body.Name + "\n\n描述这个 skill 的用途和使用步骤。\n"
	}
	return os.WriteFile(filepath.Join(dir, "SKILL.md"), []byte(content), 0644)
}

func (s *Server) getWorkspaceFile(c echo.Context) error {
	id := c.Param("id")
	file := c.Param("file")
	data, err := os.ReadFile(filepath.Join(s.roomsDir, id, "workspace", file))
	if err != nil {
		return c.JSON(404, map[string]string{"error": "not found"})
	}
	return c.JSON(200, map[string]string{"content": string(data)})
}

func (s *Server) getMessages(c echo.Context) error {
	id := c.Param("id")
	// Read today's history file and parse into messages
	historyFile := filepath.Join(s.roomsDir, id, "history", time.Now().Format("2006-01-02")+".md")
	data, err := os.ReadFile(historyFile)
	if err != nil {
		return c.JSON(200, []interface{}{})
	}
	type msg struct {
		ID      string `json:"id"`
		Agent   string `json:"agent"`
		Role    string `json:"role"`
		Content string `json:"content"`
	}
	var msgs []msg
	lines := strings.Split(string(data), "\n## ")
	for i, block := range lines {
		if block == "" {
			continue
		}
		// Parse "**[HH:MM:SS] agentName** (role)\n\ncontent"
		parts := strings.SplitN(block, "\n\n", 2)
		if len(parts) < 2 {
			continue
		}
		header := parts[0]
		content := strings.TrimSpace(parts[1])
		// Extract role and agent from header like "**[15:04:05] agentName** (role)"
		var agentName, role string
		fmt.Sscanf(header, "**[%*s] %s** (%s)", &agentName, &role)
		agentName = strings.TrimSuffix(agentName, "**")
		role = strings.TrimSuffix(role, ")")
		if agentName == "" {
			agentName = "unknown"
		}
		if role == "" {
			role = "member"
		}
		msgs = append(msgs, msg{ID: fmt.Sprintf("%d", i), Agent: agentName, Role: role, Content: content})
	}
	return c.JSON(200, msgs)
}

func (s *Server) getUser(c echo.Context) error {
	if s.user == nil {
		return c.JSON(404, map[string]string{"error": "user not found"})
	}
	return c.JSON(200, map[string]interface{}{
		"name":         s.user.GetName(),
		"provider":     s.user.GetProvider(),
		"model":        s.user.GetModel(),
		"api_key_env":  s.user.GetAPIKeyEnv(),
		"avatar_color": s.user.GetAvatarColor(),
		"description":  s.user.Config.Description,
		"has_profile":  s.user.Profile != "",
	})
}

func (s *Server) getUserProfile(c echo.Context) error {
	if s.user == nil {
		return c.JSON(404, map[string]string{"error": "user not found"})
	}
	return c.JSON(200, map[string]string{"content": s.user.Profile})
}

func (s *Server) saveUserProfile(c echo.Context) error {
	if s.user == nil {
		return c.JSON(404, map[string]string{"error": "user not found"})
	}
	var body struct {
		Content string `json:"content"`
	}
	if err := c.Bind(&body); err != nil {
		return err
	}
	if err := s.user.SaveProfile(body.Content); err != nil {
		return err
	}
	return c.JSON(200, map[string]string{"status": "ok"})
}

func (s *Server) saveUserConfig(c echo.Context) error {
	if s.user == nil {
		return c.JSON(404, map[string]string{"error": "user not found"})
	}
	var body struct {
		Name        string `json:"name"`
		Description string `json:"description"`
		Provider    string `json:"provider"`
		Model       string `json:"model"`
		APIKeyEnv   string `json:"api_key_env"`
		AvatarColor string `json:"avatar_color"`
	}
	if err := c.Bind(&body); err != nil {
		return err
	}
	cfg := s.user.Config
	if body.Name != "" {
		cfg.Name = body.Name
	}
	if body.Description != "" {
		cfg.Description = body.Description
	}
	if body.Provider != "" {
		cfg.Provider = body.Provider
	}
	if body.Model != "" {
		cfg.Model = body.Model
	}
	if body.APIKeyEnv != "" {
		cfg.APIKeyEnv = body.APIKeyEnv
	}
	if body.AvatarColor != "" {
		cfg.AvatarColor = body.AvatarColor
	}
	if err := s.user.SaveConfig(cfg); err != nil {
		return err
	}
	return c.JSON(200, map[string]string{"status": "ok"})
}

// --- Teams API ---

func (s *Server) listTeams(c echo.Context) error {
	entries, err := os.ReadDir(s.teamsDir)
	if err != nil {
		return c.JSON(200, []interface{}{})
	}
	var teams []interface{}
	for _, e := range entries {
		if !e.IsDir() {
			continue
		}
		team, err := loadTeam(filepath.Join(s.teamsDir, e.Name(), "team.yaml"))
		if err != nil {
			continue
		}
		teams = append(teams, map[string]interface{}{
			"name":         team.Name,
			"description":  team.Description,
			"author":       team.Author,
			"repo_url":     team.RepoURL,
			"agents":       team.Agents,
			"skills":       team.Skills,
			"installed_at": team.InstalledAt,
		})
	}
	return c.JSON(200, teams)
}

func (s *Server) getTeam(c echo.Context) error {
	name := c.Param("name")
	team, err := loadTeam(filepath.Join(s.teamsDir, name, "team.yaml"))
	if err != nil {
		return c.JSON(404, map[string]string{"error": "team not found"})
	}
	return c.JSON(200, map[string]interface{}{
		"name":         team.Name,
		"description":  team.Description,
		"author":       team.Author,
		"repo_url":     team.RepoURL,
		"agents":       team.Agents,
		"skills":       team.Skills,
		"installed_at": team.InstalledAt,
	})
}

func (s *Server) getTeamAgentFile(c echo.Context) error {
	teamName := c.Param("name")
	agentName := c.Param("agent")
	fileName := c.Param("file")
	path := filepath.Join(s.agentsDir, teamName, agentName, fileName)
	data, err := os.ReadFile(path)
	if err != nil {
		return c.JSON(404, map[string]string{"error": "file not found"})
	}
	return c.JSON(200, map[string]string{"content": string(data)})
}

func (s *Server) saveTeamAgentFile(c echo.Context) error {
	teamName := c.Param("name")
	agentName := c.Param("agent")
	fileName := c.Param("file")
	path := filepath.Join(s.agentsDir, teamName, agentName, fileName)
	var body struct {
		Content string `json:"content"`
	}
	if err := c.Bind(&body); err != nil {
		return err
	}
	os.MkdirAll(filepath.Dir(path), 0755)
	if err := os.WriteFile(path, []byte(body.Content), 0644); err != nil {
		return err
	}
	return c.JSON(200, map[string]string{"status": "ok"})
}

func (s *Server) listTeamKnowledge(c echo.Context) error {
	name := c.Param("name")
	dir := filepath.Join(s.agentsDir, name, "knowledge")
	entries, err := os.ReadDir(dir)
	if err != nil {
		return c.JSON(200, []string{})
	}
	var files []string
	for _, e := range entries {
		if !e.IsDir() {
			files = append(files, e.Name())
		}
	}
	return c.JSON(200, files)
}

func (s *Server) getTeamKnowledgeFile(c echo.Context) error {
	name := c.Param("name")
	fileName := c.Param("file")
	path := filepath.Join(s.agentsDir, name, "knowledge", fileName)
	data, err := os.ReadFile(path)
	if err != nil {
		return c.JSON(404, map[string]string{"error": "file not found"})
	}
	return c.JSON(200, map[string]string{"content": string(data)})
}

func (s *Server) saveTeamKnowledgeFile(c echo.Context) error {
	name := c.Param("name")
	fileName := c.Param("file")
	path := filepath.Join(s.agentsDir, name, "knowledge", fileName)
	var body struct {
		Content string `json:"content"`
	}
	if err := c.Bind(&body); err != nil {
		return err
	}
	os.MkdirAll(filepath.Dir(path), 0755)
	if err := os.WriteFile(path, []byte(body.Content), 0644); err != nil {
		return err
	}
	return c.JSON(200, map[string]string{"status": "ok"})
}

func (s *Server) deleteTeam(c echo.Context) error {
	name := c.Param("name")
	// Delete team directory
	if err := os.RemoveAll(filepath.Join(s.teamsDir, name)); err != nil {
		return err
	}
	// Delete agents
	if err := os.RemoveAll(filepath.Join(s.agentsDir, name)); err != nil {
		// Continue even if error
	}
	return c.JSON(200, map[string]string{"status": "ok"})
}

func loadTeam(path string) (*hub.Team, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	// Simple YAML parse
	team := &hub.Team{}
	lines := strings.Split(string(data), "\n")
	for _, line := range lines {
		if strings.HasPrefix(line, "name: ") {
			team.Name = strings.TrimSpace(strings.TrimPrefix(line, "name: "))
		} else if strings.HasPrefix(line, "description: ") {
			team.Description = strings.TrimSpace(strings.TrimPrefix(line, "description: "))
		} else if strings.HasPrefix(line, "author: ") {
			team.Author = strings.TrimSpace(strings.TrimPrefix(line, "author: "))
		} else if strings.HasPrefix(line, "repo_url: ") {
			team.RepoURL = strings.TrimSpace(strings.TrimPrefix(line, "repo_url: "))
		} else if strings.HasPrefix(line, "installed_at: ") {
			team.InstalledAt = strings.TrimSpace(strings.TrimPrefix(line, "installed_at: "))
		} else if strings.HasPrefix(line, "  - ") {
			team.Agents = append(team.Agents, strings.TrimSpace(strings.TrimPrefix(line, "  - ")))
		}
	}
	return team, nil
}