agent-team/internal/room/room.go

package room

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/sdaduanbilei/agent-team/internal/agent"
	"github.com/sdaduanbilei/agent-team/internal/llm"
	"github.com/sdaduanbilei/agent-team/internal/skill"
	"github.com/sdaduanbilei/agent-team/internal/user"
	"gopkg.in/yaml.v3"
)

type RoomType string

const (
	TypeDept   RoomType = "dept"
	TypeLeader RoomType = "leader"
)

type Status string

const (
	StatusPending  Status = "pending"
	StatusThinking Status = "thinking"
	StatusWorking  Status = "working"
)

type Config struct {
	Name    string   `yaml:"name"`
	Type    RoomType `yaml:"type"`
	Master  string   `yaml:"master"`  // agent name
	Members []string `yaml:"members"` // agent names
	Color   string   `yaml:"color"`   // avatar color
	Team    string   `yaml:"team"`    // installed team name
}

type Room struct {
	Config    Config
	Dir       string
	master    *agent.Agent
	members   map[string]*agent.Agent
	skillMeta []skill.Meta
	User      *user.User
	Status    Status
	ActiveAgent string // for working status display
	Broadcast func(Event) // set by api layer
}

type EventType string

const (
	EvtAgentMessage EventType = "agent_message"
	EvtTaskAssign   EventType = "task_assign"
	EvtReview       EventType = "review"
	EvtRoomStatus   EventType = "room_status"
	EvtTasksUpdate  EventType = "tasks_update"
	EvtWorkspaceFile EventType = "workspace_file"
)

type Event struct {
	Type        EventType `json:"type"`
	RoomID      string    `json:"room_id"`
	Agent       string    `json:"agent,omitempty"`
	Role        string    `json:"role,omitempty"` // master | member | challenge
	Content     string    `json:"content,omitempty"`
	Streaming   bool      `json:"streaming,omitempty"`
	From        string    `json:"from,omitempty"`
	To          string    `json:"to,omitempty"`
	Task        string    `json:"task,omitempty"`
	Feedback    string    `json:"feedback,omitempty"`
	Status      Status    `json:"status,omitempty"`
	ActiveAgent string    `json:"active_agent,omitempty"`
	Action      string    `json:"action,omitempty"`
	Filename    string    `json:"filename,omitempty"`
}

type BoardEntry struct {
	Author  string
	Content string
	Type    string // "draft" | "challenge"
}

type SharedBoard struct {
	mu      sync.RWMutex
	entries []BoardEntry
}

func (b *SharedBoard) Add(author, content, typ string) {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.entries = append(b.entries, BoardEntry{Author: author, Content: content, Type: typ})
}

func (b *SharedBoard) ToContext() string {
	b.mu.RLock()
	defer b.mu.RUnlock()
	if len(b.entries) == 0 {
		return ""
	}
	var sb strings.Builder
	sb.WriteString("<team_board>\n")
	for _, e := range b.entries {
		fmt.Fprintf(&sb, "  <entry type=\"%s\" author=\"%s\">\n%s\n  </entry>\n", e.Type, e.Author, e.Content)
	}
	sb.WriteString("</team_board>")
	return sb.String()
}

func Load(roomDir string, agentsDir string, skillsDir string) (*Room, error) {
	data, err := os.ReadFile(filepath.Join(roomDir, "room.md"))
	if err != nil {
		return nil, err
	}
	cfg, err := parseRoomConfig(data)
	if err != nil {
		return nil, err
	}

	r := &Room{Config: cfg, Dir: roomDir, members: make(map[string]*agent.Agent)}

	if cfg.Master != "" {
		agentPath := resolveAgentPath(agentsDir, cfg.Team, cfg.Master)
		r.master, err = agent.Load(agentPath)
		if err != nil {
			return nil, fmt.Errorf("load master %s: %w", cfg.Master, err)
		}
		for _, name := range cfg.Members {
			a, err := agent.Load(resolveAgentPath(agentsDir, cfg.Team, name))
			if err != nil {
				return nil, fmt.Errorf("load member %s: %w", name, err)
			}
			r.members[name] = a
		}
	}

	r.skillMeta, _ = skill.Discover(skillsDir)
	return r, nil
}

func (r *Room) emit(e Event) {
	e.RoomID = r.Config.Name
	if r.Broadcast != nil {
		r.Broadcast(e)
	}
}

func (r *Room) setStatus(s Status, activeAgent, action string) {
	r.Status = s
	r.ActiveAgent = activeAgent
	r.emit(Event{Type: EvtRoomStatus, Status: s, ActiveAgent: activeAgent, Action: action})
}

// AppendHistory persists a message to today's history file.
func (r *Room) AppendHistory(role, agentName, content string) {
	dir := filepath.Join(r.Dir, "history")
	os.MkdirAll(dir, 0755)
	filename := filepath.Join(dir, time.Now().Format("2006-01-02")+".md")
	line := fmt.Sprintf("\n**[%s] %s** (%s)\n\n%s\n", time.Now().Format("15:04:05"), agentName, role, content)
	f, err := os.OpenFile(filename, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return
	}
	defer f.Close()
	f.WriteString(line)
}

func (r *Room) runMembersParallel(ctx context.Context, assignments map[string]string, board *SharedBoard, skillXML string) map[string]string {
	results := make(map[string]string)
	var mu sync.Mutex
	var wg sync.WaitGroup

	for memberName, task := range assignments {
		wg.Add(1)
		go func(name, t string) {
			defer wg.Done()
			member, ok := r.members[name]
			if !ok {
				return
			}
			r.setStatus(StatusWorking, member.Config.Name, t)
			r.emit(Event{Type: EvtTaskAssign, From: r.master.Config.Name, To: name, Task: t})

			boardCtx := board.ToContext()
			extraCtx := skillXML
			if boardCtx != "" {
				extraCtx = boardCtx + "\n\n" + skillXML
			}
			memberSystem := member.BuildSystemPrompt(extraCtx)
			memberMsgs := []llm.Message{
				llm.NewMsg("system", memberSystem),
				llm.NewMsg("user", t),
			}
			var memberReply strings.Builder
			_, err := member.Chat(ctx, memberMsgs, func(token string) {
				memberReply.WriteString(token)
				r.emit(Event{Type: EvtAgentMessage, Agent: name, Role: "member", Content: token, Streaming: true})
			})
			if err != nil {
				mu.Lock()
				results[name] = fmt.Sprintf("[error] %v", err)
				mu.Unlock()
				return
			}
			result := memberReply.String()
			mu.Lock()
			results[name] = result
			mu.Unlock()
			r.AppendHistory("member", name, result)
			board.Add(name, result, "draft")

			if strings.Contains(result, "# ") {
				filename := fmt.Sprintf("%s-%s.md", name, time.Now().Format("20060102-150405"))
				r.saveWorkspace(filename, result)
				r.emit(Event{Type: EvtWorkspaceFile, Filename: filename, Content: result})
			}
		}(memberName, task)
	}
	wg.Wait()
	return results
}

func (r *Room) runChallengeRound(ctx context.Context, board *SharedBoard, skillXML string) {
	var challengers []string
	for name, member := range r.members {
		if member.Config.CanChallenge {
			challengers = append(challengers, name)
		}
	}
	if len(challengers) == 0 {
		return
	}

	boardCtx := board.ToContext()
	if boardCtx == "" {
		return
	}

	var wg sync.WaitGroup
	for _, name := range challengers {
		wg.Add(1)
		go func(n string) {
			defer wg.Done()
			member := r.members[n]
			extraCtx := boardCtx + "\n\n" + skillXML
			memberSystem := member.BuildSystemPrompt(extraCtx)
			memberMsgs := []llm.Message{
				llm.NewMsg("system", memberSystem+"\n\nReview the team board above. If you see issues or want to challenge any draft, output CHALLENGE:<your concern>. Otherwise output AGREE."),
				llm.NewMsg("user", "Please review the team board and provide your feedback."),
			}
			var reply strings.Builder
			_, err := member.Chat(ctx, memberMsgs, func(token string) {
				reply.WriteString(token)
				r.emit(Event{Type: EvtAgentMessage, Agent: n, Role: "challenge", Content: token, Streaming: true})
			})
			if err != nil {
				return
			}
			result := reply.String()
			if strings.Contains(result, "CHALLENGE:") {
				board.Add(n, result, "challenge")
				r.AppendHistory("challenge", n, result)
			}
		}(name)
	}
	wg.Wait()
}

// Handle processes a user message through master orchestration.
func (r *Room) Handle(ctx context.Context, userMsg string) error {
	return r.HandleUserMessage(ctx, "user", userMsg)
}

// HandleUserMessage processes a user message with a specific user name.
func (r *Room) HandleUserMessage(ctx context.Context, userName, userMsg string) error {
	if r.master == nil {
		return fmt.Errorf("room has no master agent configured")
	}
	r.AppendHistory("user", userName, userMsg)
	r.setStatus(StatusThinking, "", "")

	// Build master context
	teamXML := r.buildTeamXML()
	skillXML := skill.ToXML(r.skillMeta)

	// Build user info XML
	var userXML string
	if r.User != nil {
		userXML = r.User.BuildUserXML()
	}

	extraContext := userXML + "\n\n" + teamXML + "\n\n" + skillXML
	systemPrompt := r.master.BuildSystemPrompt(extraContext)

	masterMsgs := []llm.Message{
		llm.NewMsg("system", systemPrompt+"\n\nYou are the master of this team. When you need a team member to do something, output a line like: ASSIGN:<member_name>:<task description>. When you are done reviewing and satisfied, output DONE:<summary>."),
		llm.NewMsg("user", userMsg),
	}

	// Master planning loop
	for iteration := 0; iteration < 5; iteration++ {
		var masterReply strings.Builder
		_, err := r.master.Chat(ctx, masterMsgs, func(token string) {
			masterReply.WriteString(token)
			r.emit(Event{Type: EvtAgentMessage, Agent: r.master.Config.Name, Role: "master", Content: token, Streaming: true})
		})
		if err != nil {
			return err
		}
		reply := masterReply.String()
		masterMsgs = append(masterMsgs, llm.NewMsg("assistant", reply))
		r.AppendHistory("master", r.master.Config.Name, reply)

		// Parse assignments
		assignments := parseAssignments(reply)
		if len(assignments) == 0 {
			// No assignments, master is done
			break
		}

		// Execute assignments in parallel
		board := &SharedBoard{}
		results := r.runMembersParallel(ctx, assignments, board, skillXML)

		// Run challenge round
		r.runChallengeRound(ctx, board, skillXML)

		// Feed results back to master for review
		r.setStatus(StatusThinking, "", "")
		var resultsStr strings.Builder
		for memberName, result := range results {
			resultsStr.WriteString(fmt.Sprintf("[%s] %s\n", memberName, result))
		}
		boardCtx := board.ToContext()
		feedbackMsg := "Team results:\n" + resultsStr.String()
		if boardCtx != "" {
			feedbackMsg += "\n\nTeam board:\n" + boardCtx
		}
		feedbackMsg += "\n\nPlease review. If satisfied output DONE:<summary>, otherwise output ASSIGN instructions for revisions."
		masterMsgs = append(masterMsgs, llm.NewMsg("user", feedbackMsg))

		// Update tasks
		r.updateTasks(masterMsgs)

		if strings.Contains(reply, "DONE:") {
			break
		}
	}

	r.setStatus(StatusPending, "", "")

	// Auto-update master memory after task completion
	go r.updateMasterMemory(context.Background(), userMsg, masterMsgs)

	return nil
}

func (r *Room) updateMasterMemory(ctx context.Context, task string, msgs []llm.Message) {
	summaryPrompt := fmt.Sprintf("Based on this task: %q\nSummarize key learnings and patterns in 3-5 bullet points for future reference. Be concise.", task)
	memMsgs := append(msgs, llm.NewMsg("user", summaryPrompt))
	summary, err := r.master.Chat(ctx, memMsgs, nil)
	if err != nil || summary == "" {
		return
	}
	filename := time.Now().Format("2006-01") + ".md"
	existing, _ := os.ReadFile(filepath.Join(r.master.Dir, "memory", filename))
	content := string(existing) + fmt.Sprintf("\n## %s — %s\n\n%s\n", time.Now().Format("2006-01-02"), task[:min(50, len(task))], summary)
	r.master.SaveMemory(filename, content)
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func parseAssignments(text string) map[string]string {
	result := make(map[string]string)
	for _, line := range strings.Split(text, "\n") {
		if strings.HasPrefix(line, "ASSIGN:") {
			parts := strings.SplitN(strings.TrimPrefix(line, "ASSIGN:"), ":", 2)
			if len(parts) == 2 {
				result[strings.TrimSpace(parts[0])] = strings.TrimSpace(parts[1])
			}
		}
	}
	return result
}

func (r *Room) buildTeamXML() string {
	var sb strings.Builder
	sb.WriteString("<team_members>\n")
	for name, a := range r.members {
		fmt.Fprintf(&sb, "  <member>\n    <name>%s</name>\n    <description>%s</description>\n  </member>\n", name, a.Config.Description)
	}
	sb.WriteString("</team_members>")
	return sb.String()
}

func (r *Room) saveWorkspace(filename, content string) {
	dir := filepath.Join(r.Dir, "workspace")
	os.MkdirAll(dir, 0755)
	os.WriteFile(filepath.Join(dir, filename), []byte(content), 0644)
}

func (r *Room) updateTasks(msgs []llm.Message) {
	// Extract task list from conversation and save
	var tasks strings.Builder
	tasks.WriteString("# Tasks\n\n")
	for _, m := range msgs {
		if m.Role == "assistant" && strings.Contains(m.Content, "ASSIGN:") {
			for _, line := range strings.Split(m.Content, "\n") {
				if strings.HasPrefix(line, "ASSIGN:") {
					parts := strings.SplitN(strings.TrimPrefix(line, "ASSIGN:"), ":", 2)
					if len(parts) == 2 {
						tasks.WriteString(fmt.Sprintf("- [ ] [%s] %s\n", strings.TrimSpace(parts[0]), strings.TrimSpace(parts[1])))
					}
				}
			}
		}
	}
	content := tasks.String()
	os.WriteFile(filepath.Join(r.Dir, "tasks.md"), []byte(content), 0644)
	r.emit(Event{Type: EvtTasksUpdate, Content: content})
}

func parseRoomConfig(data []byte) (Config, error) {
	var cfg Config
	if !bytes.HasPrefix(data, []byte("---")) {
		return cfg, fmt.Errorf("missing frontmatter")
	}
	parts := bytes.SplitN(data, []byte("---"), 3)
	if len(parts) < 3 {
		return cfg, fmt.Errorf("invalid frontmatter")
	}
	return cfg, yaml.Unmarshal(parts[1], &cfg)
}

// resolveAgentPath finds agent dir: prefers agentsDir/team/name, falls back to agentsDir/name
func resolveAgentPath(agentsDir, team, name string) string {
	if team != "" {
		p := filepath.Join(agentsDir, team, name)
		if _, err := os.Stat(p); err == nil {
			return p
		}
	}
	return filepath.Join(agentsDir, name)
}