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"
	"github.com/sipeed/picoclaw/pkg/utils"
)

type ContextBuilder struct {
	workspace    string
	skillsLoader *skills.SkillsLoader
	memory       *MemoryStore
	tools        *tools.ToolRegistry // Direct reference to tool registry
}

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
}

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()

	return fmt.Sprintf(`# 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** - When remembering something, write to %s/memory/MEMORY.md

4. **Text Only Response** - If the user requests a text-only response (e.g. "text only", "no voice", "jangan pakai suara"), start your response with "[NO_VOICE]". The system will automatically strip this tag and suppress the voice generation.`,
		now, runtime, workspacePath, workspacePath, workspacePath, workspacePath, toolsSection, workspacePath)
}

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())

	// 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))
	}

	// 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")
}

func (cb *ContextBuilder) LoadBootstrapFiles() string {
	bootstrapFiles := []string{
		"AGENTS.md",
		"SOUL.md",
		"USER.md",
		"IDENTITY.md",
	}

	var result string
	for _, filename := range bootstrapFiles {
		filePath := filepath.Join(cb.workspace, filename)
		if data, err := os.ReadFile(filePath); err == nil {
			result += fmt.Sprintf("## %s\n\n%s\n\n", filename, string(data))
		}
	}

	return result
}

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

	systemPrompt := cb.BuildSystemPrompt()

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

	// Log system prompt summary for debugging (debug mode only)
	logger.DebugCF("agent", "System prompt built",
		map[string]interface{}{
			"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]interface{}{
			"preview": preview,
		})

	if summary != "" {
		systemPrompt += "\n\n## Summary of Previous Conversation\n\n" + summary
	}

	// Context Length Management
	// Approximate token limit (e.g. 128k tokens * 4 chars = 512k chars)
	// We use a safe limit of 500,000 characters to prevent 400 Bad Request errors
	const MaxContextChars = 500000

	// Safety: Ensure current message isn't larger than the limit
	if len(currentMessage) > MaxContextChars {
		logger.WarnCF("agent", "Current message too large, truncating", map[string]interface{}{
			"original": len(currentMessage),
			"limit":    MaxContextChars,
		})
		currentMessage = currentMessage[:MaxContextChars-1000] + "... (truncated)"
	}

	currentLen := len(systemPrompt) + len(currentMessage)
	availableChars := MaxContextChars - currentLen

	// Safety: If system prompt + current message exceeds limit, truncate system prompt
	if availableChars < 0 {
		logger.WarnCF("agent", "Context exceeded limit with just system prompt and current message", map[string]interface{}{
			"system_prompt_len": len(systemPrompt),
			"current_msg_len":   len(currentMessage),
		})

		// Reserve space for current message
		maxSystemPrompt := MaxContextChars - len(currentMessage)
		if maxSystemPrompt > 1000 { // Ensure we have at least some space for system prompt
			logger.WarnCF("agent", "Truncating system prompt to fit limit", map[string]interface{}{
				"original": len(systemPrompt),
				"new":      maxSystemPrompt,
			})
			systemPrompt = systemPrompt[:maxSystemPrompt-100] + "\n... (truncated)"
			availableChars = 0 // We fit exactly (with buffer)
		} else {
			// Extremely edge case: Current message takes up almost all space
			// We can't do much but warn, or truncate current message further
			logger.ErrorCF("agent", "Current message consumes almost all context window", nil)
		}
	}

	// Filter history to fit within context window
	var filteredHistory []providers.Message

	//This fix prevents the session memory from LLM failure due to elimination of toolu_IDs required from LLM
	// --- INICIO DEL FIX ---
	//Diegox-17
	for len(history) > 0 && (history[0].Role == "tool") {
		logger.DebugCF("agent", "Removing orphaned tool message from history to prevent LLM error",
			map[string]interface{}{"role": history[0].Role})
		history = history[1:]
	}
	//Diegox-17
	// --- FIN DEL FIX ---

	if availableChars > 0 && len(history) > 0 {
		// Add messages from end to start
		var keptMessages []providers.Message
		usedChars := 0

		for i := len(history) - 1; i >= 0; i-- {
			msg := history[i]
			msgLen := len(msg.Content)
			// Approximate length for multi-content
			for _, part := range msg.MultiContent {
				if part.Type == "text" {
					msgLen += len(part.Text)
				}
			}

			if usedChars+msgLen > availableChars {
				break
			}

			keptMessages = append(keptMessages, msg)
			usedChars += msgLen
		}

		// Reverse keptMessages to get correct order
		for i := len(keptMessages) - 1; i >= 0; i-- {
			filteredHistory = append(filteredHistory, keptMessages[i])
		}

		if len(filteredHistory) < len(history) {
			logger.WarnCF("agent", "Context truncated to fit limit", map[string]interface{}{
				"original_count": len(history),
				"kept_count":     len(filteredHistory),
				"used_chars":     usedChars,
				"limit_chars":    availableChars,
			})
		}
	} else if availableChars <= 0 {
		logger.WarnCF("agent", "Context exceeded limit with just system prompt and current message", map[string]interface{}{
			"system_prompt_len": len(systemPrompt),
			"current_msg_len":   len(currentMessage),
		})
	}

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

	messages = append(messages, filteredHistory...)

	// Construct user message
	userMsg := providers.Message{
		Role: "user",
	}

	if len(media) > 0 {
		// Multi-modal message
		userMsg.MultiContent = append(userMsg.MultiContent, providers.ContentPart{
			Type: "text",
			Text: currentMessage,
		})

		for _, path := range media {
			// Encode image
			dataURI, _, err := utils.EncodeImageToBase64(path)
			if err != nil {
				logger.ErrorCF("agent", "Failed to encode image", map[string]interface{}{"path": path, "error": err.Error()})
				continue
			}
			userMsg.MultiContent = append(userMsg.MultiContent, providers.ContentPart{
				Type: "image_url",
				ImageURL: &providers.ImageURL{
					URL: dataURI,
				},
			})
		}

		// Fallback content string for logging/compatibility
		userMsg.Content = currentMessage + fmt.Sprintf(" [Attached %d images]", len(media))
	} else {
		// Text-only message
		userMsg.Content = currentMessage
	}

	messages = append(messages, userMsg)

	return messages
}

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]interface{}) []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
}

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