picoclaw/pkg/agent/context.go

package agent

import (
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/sipeed/picoclaw/pkg/logger"
	"github.com/sipeed/picoclaw/pkg/providers"
	"github.com/sipeed/picoclaw/pkg/skills"
	"github.com/sipeed/picoclaw/pkg/tools"
)

const orchestrationGuidance = `## Orchestration

You are the conductor, not the performer. Prefer delegation over doing everything inline.

Use **spawn** (non-blocking) when:
- Tasks can run in parallel or in the background
- Multiple independent tasks can run simultaneously — spawn each one
- You don't need the result to decide the next step
- The operation is long-running (builds, fetches, analysis, file processing)

Use **subagent** (blocking) when:
- You need the result before you can continue
- Correctness of the next step depends on the outcome

Do inline only when:
- It's a single fast tool call (read a file, quick search)
- Delegation overhead clearly outweighs the benefit

Default bias: if a task involves more than 2-3 tool calls or can run independently, delegate it.
When you spawn, immediately plan what comes next — blocking means you've stopped thinking.
Fork aggressively: explore multiple directions simultaneously.

After spawning, record the assignment in ## Orchestration > Delegated in MEMORY.md.
When results come back, synthesize and decide the next fork.`

type ContextBuilder struct {
	workspace              string
	workDir                string              // session-specific working directory (worktree or project subdir)
	skillsLoader           *skills.SkillsLoader
	memory                 *MemoryStore
	tools                  *tools.ToolRegistry // Direct reference to tool registry
	peerNote               string              // set per-call from loop.go for peer session awareness
	orchestrationEnabled   bool                // set from AgentLoop when --orchestration flag is used
}

func getGlobalConfigDir() string {
	home, err := os.UserHomeDir()
	if err != nil {
		return ""
	}
	return filepath.Join(home, ".picoclaw")
}

func NewContextBuilder(workspace string) *ContextBuilder {
	// builtin skills: skills directory in current project
	// Use the skills/ directory under the current working directory
	wd, _ := os.Getwd()
	builtinSkillsDir := filepath.Join(wd, "skills")
	globalSkillsDir := filepath.Join(getGlobalConfigDir(), "skills")

	return &ContextBuilder{
		workspace:    workspace,
		skillsLoader: skills.NewSkillsLoader(workspace, globalSkillsDir, builtinSkillsDir),
		memory:       NewMemoryStore(workspace),
	}
}

// SetToolsRegistry sets the tools registry for dynamic tool summary generation.
func (cb *ContextBuilder) SetToolsRegistry(registry *tools.ToolRegistry) {
	cb.tools = registry
}

// SetWorkDir sets the session-specific working directory (e.g., worktree path
// or project subdirectory). Bootstrap files found here take priority over workspace.
func (cb *ContextBuilder) SetWorkDir(dir string) {
	cb.workDir = dir
}

// SetPeerNote sets the peer session awareness note for the current call.
func (cb *ContextBuilder) SetPeerNote(note string) {
	cb.peerNote = note
}

// SetOrchestrationEnabled sets whether orchestration is enabled.
func (cb *ContextBuilder) SetOrchestrationEnabled(enabled bool) {
	cb.orchestrationEnabled = enabled
}

func (cb *ContextBuilder) getIdentity() string {
	now := time.Now().Format("2006-01-02 15:04 (Monday)")
	workspacePath, _ := filepath.Abs(filepath.Join(cb.workspace))
	runtime := fmt.Sprintf("%s %s, Go %s", runtime.GOOS, runtime.GOARCH, runtime.Version())

	// Build tools section dynamically
	toolsSection := cb.buildToolsSection()

	// Build prompt with optional orchestration banner
	var prompt string
	if cb.orchestrationEnabled {
		prompt = `/_/_/_/_/_/_/_/_/_/_/_/_/_/_/_/_/_/
O R C H E S T R A  M O D E
\_\_\_\_\_\_\_\_\_\_\_\_\_\_\_\_\_\_\_\_\_

`
	}
	return fmt.Sprintf(prompt+`# picoclaw 🦞

You are picoclaw, a helpful AI assistant.

## Current Time
%s

## Runtime
%s

## Workspace
Your workspace is at: %s
- Memory: %s/memory/MEMORY.md
- Daily Notes: %s/memory/YYYYMM/YYYYMMDD.md
- Skills: %s/skills/{skill-name}/SKILL.md

%s

## Important Rules

1. **ALWAYS use tools** - When you need to perform an action (schedule reminders, send messages, execute commands, etc.), you MUST call the appropriate tool. Do NOT just say you'll do it or pretend to do it.

2. **Be helpful and accurate** - When using tools, briefly explain what you're doing.

3. **Memory & Plans**
   - Use memory/MEMORY.md for structured plans.
   - NEVER remove or overwrite the header block (# Active Plan, > Task:, > Status:, > Phase:). The system parses these lines to track plan state.
   - If Status is "interviewing": Ask clarifying questions.
     After each answer, use edit_file to save findings to ## Context in memory/MEMORY.md.
     When you have enough information, add ## Phase sections with "- [ ]" checkbox steps, and ## Commands section below the header. Then change > Status: to "review".
   - If Status is "review": The plan is awaiting user approval. Do NOT change Status yourself.
   - If Status is "executing": Work through the current Phase's steps.
     Mark each "- [x]" via edit_file. The system will auto-advance phases.
   - Plan format (header is written by the system — do NOT delete it):
     # Active Plan
     > Task: <description>
     > Status: interviewing | review | executing
     > Phase: <current phase number>
     ## Phase 1: <title>
     - [ ] Step 1
     - [ ] Step 2
     ## Phase 2: <title>
     - [ ] Step 1
     ## Commands
     build: <build command>
     test: <test command>
     lint: <lint command>
     ## Context
     <requirements, decisions, environment>
   - Keep each phase to 3-5 steps. Do NOT create plans without /plan.
   - Always ask about build/test/lint commands during interview.

4. **Response Formatting**
   - NEVER use ASCII box-drawing characters (┌─┐│└─┘╔═╗║╚═╝ etc.) or ASCII art diagrams.
   - Use markdown headings, bold, lists, and indentation for structure.
   - Keep lines short — most users read on mobile.
   - For architecture/flow, use arrow text: CLI → Pipeline → Adapters`,
		now, runtime, workspacePath, workspacePath, workspacePath, workspacePath, toolsSection)
}

func (cb *ContextBuilder) buildToolsSection() string {
	if cb.tools == nil {
		return ""
	}

	summaries := cb.tools.GetSummaries()
	if len(summaries) == 0 {
		return ""
	}

	var sb strings.Builder
	sb.WriteString("## Available Tools\n\n")
	sb.WriteString(
		"**CRITICAL**: You MUST use tools to perform actions. Do NOT pretend to execute commands or schedule tasks.\n\n",
	)
	sb.WriteString("You have access to the following tools:\n\n")
	for _, s := range summaries {
		sb.WriteString(s)
		sb.WriteString("\n")
	}

	return sb.String()
}

func (cb *ContextBuilder) BuildSystemPrompt() string {
	parts := []string{}

	// Core identity section
	parts = append(parts, cb.getIdentity())

	// Orchestration guidance — injected only when spawn tool is registered
	if cb.tools != nil {
		if _, hasSpawn := cb.tools.Get("spawn"); hasSpawn {
			parts = append(parts, orchestrationGuidance)
		}
	}

	// Bootstrap files
	bootstrapContent := cb.LoadBootstrapFiles()
	if bootstrapContent != "" {
		parts = append(parts, bootstrapContent)
	}

	// Skills - show summary, AI can read full content with read_file tool
	skillsSummary := cb.skillsLoader.BuildSkillsSummary()
	if skillsSummary != "" {
		parts = append(parts, fmt.Sprintf(`# Skills

The following skills extend your capabilities. To use a skill, read its SKILL.md file using the read_file tool.

%s`, skillsSummary))
	}

	// Runtime status from tools (e.g., background processes)
	if cb.tools != nil {
		if status := cb.tools.GetRuntimeStatus(); status != "" {
			parts = append(parts, status)
		}
	}

	// Peer session coordination
	if cb.peerNote != "" {
		parts = append(parts, "## Active Sessions\n\n"+cb.peerNote)
	}

	// Memory context
	memoryContext := cb.memory.GetMemoryContext()
	if memoryContext != "" {
		parts = append(parts, "# Memory\n\n"+memoryContext)
	}

	// Join with "---" separator
	return strings.Join(parts, "\n\n---\n\n")
}

// BootstrapFileInfo describes a resolved bootstrap file.
type BootstrapFileInfo struct {
	Name  string `json:"name"`
	Path  string `json:"path"`  // empty = not found
	Scope string `json:"scope"` // "project" or "global"
}

// bootstrapFileSpec defines the search scope for each bootstrap file.
type bootstrapFileSpec struct {
	Name  string
	Scope string // "project" = workDir→planWorkDir→workspace, "global" = workspace only
}

var bootstrapSpecs = []bootstrapFileSpec{
	{Name: "AGENTS.md", Scope: "project"},
	{Name: "IDENTITY.md", Scope: "project"},
	{Name: "SOUL.md", Scope: "global"},
	{Name: "USER.md", Scope: "global"},
}

// bootstrapProjectDirs returns de-duplicated search directories for project-scoped files.
func (cb *ContextBuilder) bootstrapProjectDirs() []string {
	seen := map[string]bool{}
	var dirs []string
	for _, d := range []string{cb.workDir, cb.memory.GetPlanWorkDir(), cb.workspace} {
		if d != "" && !seen[d] {
			seen[d] = true
			dirs = append(dirs, d)
		}
	}
	return dirs
}

func (cb *ContextBuilder) LoadBootstrapFiles() string {
	projectDirs := cb.bootstrapProjectDirs()

	var sb strings.Builder
	for _, spec := range bootstrapSpecs {
		var dirs []string
		if spec.Scope == "global" {
			dirs = []string{cb.workspace}
		} else {
			dirs = projectDirs
		}
		for _, dir := range dirs {
			filePath := filepath.Join(dir, spec.Name)
			if data, err := os.ReadFile(filePath); err == nil {
				fmt.Fprintf(&sb, "## %s\n\n%s\n\n", spec.Name, data)
				break
			}
		}
	}

	return sb.String()
}

// ResolveBootstrapPaths returns path resolution info for each bootstrap file
// using the same search logic as LoadBootstrapFiles.
func (cb *ContextBuilder) ResolveBootstrapPaths() []BootstrapFileInfo {
	projectDirs := cb.bootstrapProjectDirs()

	result := make([]BootstrapFileInfo, 0, len(bootstrapSpecs))
	for _, spec := range bootstrapSpecs {
		info := BootstrapFileInfo{Name: spec.Name, Scope: spec.Scope}
		var dirs []string
		if spec.Scope == "global" {
			dirs = []string{cb.workspace}
		} else {
			dirs = projectDirs
		}
		for _, dir := range dirs {
			filePath := filepath.Join(dir, spec.Name)
			if _, err := os.Stat(filePath); err == nil {
				info.Path = filePath
				break
			}
		}
		result = append(result, info)
	}
	return result
}

func (cb *ContextBuilder) BuildMessages(
	history []providers.Message,
	summary string,
	currentMessage string,
	media []string,
	channel, chatID string,
) []providers.Message {
	messages := []providers.Message{}

	var sysBuilder strings.Builder
	sysBuilder.WriteString(cb.BuildSystemPrompt())

	// Add Current Session info if provided
	if channel != "" && chatID != "" {
		fmt.Fprintf(&sysBuilder, "\n\n## Current Session\nChannel: %s\nChat ID: %s", channel, chatID)
	}

	// Log system prompt summary for debugging (debug mode only)
	systemPrompt := sysBuilder.String()
	logger.DebugCF("agent", "System prompt built",
		map[string]any{
			"total_chars":   len(systemPrompt),
			"total_lines":   strings.Count(systemPrompt, "\n") + 1,
			"section_count": strings.Count(systemPrompt, "\n\n---\n\n") + 1,
		})

	// Log preview of system prompt (avoid logging huge content)
	preview := systemPrompt
	if len(preview) > 500 {
		preview = preview[:500] + "... (truncated)"
	}
	logger.DebugCF("agent", "System prompt preview",
		map[string]any{
			"preview": preview,
		})

	if summary != "" {
		sysBuilder.WriteString("\n\n## Summary of Previous Conversation\n\n")
		sysBuilder.WriteString(summary)
		systemPrompt = sysBuilder.String()
	}

	history = sanitizeHistoryForProvider(history)

	messages = append(messages, providers.Message{
		Role:    "system",
		Content: systemPrompt,
	})

	messages = append(messages, history...)

	if strings.TrimSpace(currentMessage) != "" {
		messages = append(messages, providers.Message{
			Role:    "user",
			Content: currentMessage,
		})
	}

	return messages
}

func sanitizeHistoryForProvider(history []providers.Message) []providers.Message {
	if len(history) == 0 {
		return history
	}

	sanitized := make([]providers.Message, 0, len(history))
	for _, msg := range history {
		switch msg.Role {
		case "tool":
			if len(sanitized) == 0 {
				logger.DebugCF("agent", "Dropping orphaned leading tool message", map[string]any{})
				continue
			}
			last := sanitized[len(sanitized)-1]
			// Allow tool results after their assistant or after sibling tool results
			if last.Role == "tool" || (last.Role == "assistant" && len(last.ToolCalls) > 0) {
				sanitized = append(sanitized, msg)
			} else {
				logger.DebugCF("agent", "Dropping orphaned tool message", map[string]any{})
			}

		case "assistant":
			if len(msg.ToolCalls) > 0 {
				if len(sanitized) == 0 {
					logger.DebugCF("agent", "Dropping assistant tool-call turn at history start", map[string]any{})
					continue
				}
				prev := sanitized[len(sanitized)-1]
				if prev.Role != "user" && prev.Role != "tool" {
					logger.DebugCF(
						"agent",
						"Dropping assistant tool-call turn with invalid predecessor",
						map[string]any{"prev_role": prev.Role},
					)
					continue
				}
			}
			sanitized = append(sanitized, msg)

		default:
			sanitized = append(sanitized, msg)
		}
	}

	return sanitized
}

func (cb *ContextBuilder) AddToolResult(
	messages []providers.Message,
	toolCallID, toolName, result string,
) []providers.Message {
	messages = append(messages, providers.Message{
		Role:       "tool",
		Content:    result,
		ToolCallID: toolCallID,
	})
	return messages
}

func (cb *ContextBuilder) AddAssistantMessage(
	messages []providers.Message,
	content string,
	toolCalls []map[string]any,
) []providers.Message {
	msg := providers.Message{
		Role:    "assistant",
		Content: content,
	}
	// Always add assistant message, whether or not it has tool calls
	messages = append(messages, msg)
	return messages
}

func (cb *ContextBuilder) loadSkills() string {
	allSkills := cb.skillsLoader.ListSkills()
	if len(allSkills) == 0 {
		return ""
	}

	var skillNames []string
	for _, s := range allSkills {
		skillNames = append(skillNames, s.Name)
	}

	content := cb.skillsLoader.LoadSkillsForContext(skillNames)
	if content == "" {
		return ""
	}

	return "# Skill Definitions\n\n" + content
}

// LoadSkill loads a skill by name, returning its content (with frontmatter stripped) and whether it was found.
func (cb *ContextBuilder) LoadSkill(name string) (string, bool) {
	return cb.skillsLoader.LoadSkill(name)
}

// ListSkills returns all available skills from all tiers.
func (cb *ContextBuilder) ListSkills() []skills.SkillInfo {
	return cb.skillsLoader.ListSkills()
}

// Memory returns the underlying MemoryStore for direct plan queries.
func (cb *ContextBuilder) Memory() *MemoryStore {
	return cb.memory
}

// ---------- Plan passthrough methods ----------

// ReadMemory reads the long-term memory (MEMORY.md).
func (cb *ContextBuilder) ReadMemory() string {
	return cb.memory.ReadLongTerm()
}

// WriteMemory writes content to the long-term memory file.
func (cb *ContextBuilder) WriteMemory(content string) error {
	return cb.memory.WriteLongTerm(content)
}

// ClearMemory removes the long-term memory file.
func (cb *ContextBuilder) ClearMemory() error {
	return cb.memory.ClearLongTerm()
}

// HasActivePlan returns true if MEMORY.md contains an active plan.
func (cb *ContextBuilder) HasActivePlan() bool {
	return cb.memory.HasActivePlan()
}

// GetPlanStatus returns the plan status: "interviewing", "executing", or "".
func (cb *ContextBuilder) GetPlanStatus() string {
	return cb.memory.GetPlanStatus()
}

// IsPlanComplete returns true if all steps in all phases are [x].
func (cb *ContextBuilder) IsPlanComplete() bool {
	return cb.memory.IsPlanComplete()
}

// IsCurrentPhaseComplete returns true if all steps in the current phase are [x].
func (cb *ContextBuilder) IsCurrentPhaseComplete() bool {
	return cb.memory.IsCurrentPhaseComplete()
}

// AdvancePhase increments the current phase number by 1.
func (cb *ContextBuilder) AdvancePhase() error {
	return cb.memory.AdvancePhase()
}

// SetCurrentPhase sets the current phase number to n.
func (cb *ContextBuilder) SetCurrentPhase(n int) error {
	return cb.memory.SetPhase(n)
}

// GetCurrentPhase returns the current phase number.
func (cb *ContextBuilder) GetCurrentPhase() int {
	return cb.memory.GetCurrentPhase()
}

// GetTotalPhases returns the total number of phases in the plan.
func (cb *ContextBuilder) GetTotalPhases() int {
	return cb.memory.GetTotalPhases()
}

// FormatPlanDisplay returns a user-facing display of the full plan.
func (cb *ContextBuilder) FormatPlanDisplay() string {
	return cb.memory.FormatPlanDisplay()
}

// MarkStep marks a step as done in the specified phase.
func (cb *ContextBuilder) MarkStep(phase, step int) error {
	return cb.memory.MarkStep(phase, step)
}

// AddStep appends a new step to the given phase.
func (cb *ContextBuilder) AddStep(phase int, desc string) error {
	return cb.memory.AddStep(phase, desc)
}

// ValidatePlanStructure validates plan structure for interview→review transition.
func (cb *ContextBuilder) ValidatePlanStructure() error {
	return cb.memory.ValidatePlanStructure()
}

// SetPlanStatus sets the plan status.
func (cb *ContextBuilder) SetPlanStatus(status string) error {
	return cb.memory.SetStatus(status)
}

// GetPlanWorkDir returns the WorkDir from the plan metadata, or "".
func (cb *ContextBuilder) GetPlanWorkDir() string {
	return cb.memory.GetPlanWorkDir()
}

// GetPlanTaskName returns the task description from the plan metadata, or "".
func (cb *ContextBuilder) GetPlanTaskName() string {
	return cb.memory.GetPlanTaskName()
}

// GetSkillsInfo returns information about loaded skills.
func (cb *ContextBuilder) GetSkillsInfo() map[string]any {
	allSkills := cb.skillsLoader.ListSkills()
	skillNames := make([]string, 0, len(allSkills))
	for _, s := range allSkills {
		skillNames = append(skillNames, s.Name)
	}
	return map[string]any{
		"total":     len(allSkills),
		"available": len(allSkills),
		"names":     skillNames,
	}
}