cc-connect/core/command.go

package core

import (
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
)

// CustomCommand represents a registered slash command (from config or agent command files).
type CustomCommand struct {
	Name        string // command name without leading "/"
	Description string
	Prompt      string // template with {{1}}, {{2}}, {{2*}}, {{args}} placeholders
	Source      string // "config" or "agent" (for display)
}

// CommandRegistry holds all available custom commands and resolves agent command files.
type CommandRegistry struct {
	mu        sync.RWMutex
	commands  map[string]*CustomCommand // from config.toml or runtime add
	agentDirs []string                  // directories to scan for *.md command files
}

func NewCommandRegistry() *CommandRegistry {
	return &CommandRegistry{
		commands: make(map[string]*CustomCommand),
	}
}

// Add registers a custom command.
func (r *CommandRegistry) Add(name, description, prompt, source string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.commands[strings.ToLower(name)] = &CustomCommand{
		Name:        name,
		Description: description,
		Prompt:      prompt,
		Source:      source,
	}
}

// Remove deletes a config-defined custom command by name. Returns false if not found.
func (r *CommandRegistry) Remove(name string) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	lower := strings.ToLower(name)
	if _, ok := r.commands[lower]; ok {
		delete(r.commands, lower)
		return true
	}
	return false
}

// SetAgentDirs sets the directories to scan for agent command files.
func (r *CommandRegistry) SetAgentDirs(dirs []string) {
	r.agentDirs = dirs
}

// Resolve looks up a command by name. Config commands take priority, then
// agent command directories are scanned for a matching .md file.
func (r *CommandRegistry) Resolve(name string) (*CustomCommand, bool) {
	lower := strings.ToLower(name)

	r.mu.RLock()
	if c, ok := r.commands[lower]; ok {
		r.mu.RUnlock()
		return c, true
	}
	r.mu.RUnlock()

	for _, dir := range r.agentDirs {
		mdPath := filepath.Join(dir, name+".md")
		absDir, err1 := filepath.Abs(dir)
		absPath, err2 := filepath.Abs(mdPath)
		if err1 != nil || err2 != nil || !strings.HasPrefix(absPath, absDir+string(filepath.Separator)) {
			continue
		}
		data, err := os.ReadFile(mdPath)
		if err != nil {
			continue
		}
		content := strings.TrimSpace(string(data))
		if content == "" {
			continue
		}
		slog.Debug("command: loaded agent command file", "path", mdPath)
		return &CustomCommand{
			Name:   name,
			Prompt: content,
			Source: "agent",
		}, true
	}

	return nil, false
}

// ListAll returns all registered commands (config + agent command files).
func (r *CommandRegistry) ListAll() []*CustomCommand {
	r.mu.RLock()
	defer r.mu.RUnlock()

	seen := make(map[string]bool)
	var result []*CustomCommand

	for _, c := range r.commands {
		result = append(result, c)
		seen[strings.ToLower(c.Name)] = true
	}

	for _, dir := range r.agentDirs {
		entries, err := os.ReadDir(dir)
		if err != nil {
			continue
		}
		for _, entry := range entries {
			if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".md") {
				continue
			}
			name := strings.TrimSuffix(entry.Name(), ".md")
			if seen[strings.ToLower(name)] {
				continue
			}
			seen[strings.ToLower(name)] = true

			desc := ""
			data, err := os.ReadFile(filepath.Join(dir, entry.Name()))
			if err == nil {
				first, _, _ := strings.Cut(strings.TrimSpace(string(data)), "\n")
				if len([]rune(first)) > 60 {
					first = string([]rune(first)[:60]) + "..."
				}
				desc = first
			}

			result = append(result, &CustomCommand{
				Name:        name,
				Description: desc,
				Source:      "agent",
			})
		}
	}

	return result
}

// placeholderRe matches {{1}}, {{2*}}, {{args}}, and variants with defaults like {{1:foo}}.
var placeholderRe = regexp.MustCompile(`\{\{(\d+\*?|args)(:[^}]*)?\}\}`)

// ExpandPrompt replaces template placeholders with the provided arguments.
//
// Supported placeholders:
//   - {{1}}, {{2}}, ...       — positional argument (1-based)
//   - {{1:default}}           — positional with default value if arg not provided
//   - {{2*}}                  — argument N and everything after it
//   - {{2*:default}}          — same, with default
//   - {{args}}                — all arguments joined by space
//   - {{args:default}}        — all arguments, with default if none provided
//
// If the template has no placeholders, arguments are appended to the end.
func ExpandPrompt(template string, args []string) string {
	if !placeholderRe.MatchString(template) {
		if len(args) > 0 {
			return template + "\n\n" + strings.Join(args, " ")
		}
		return template
	}

	result := placeholderRe.ReplaceAllStringFunc(template, func(match string) string {
		inner := match[2 : len(match)-2] // strip {{ and }}
		key, defaultVal, hasDefault := strings.Cut(inner, ":")

		if key == "args" {
			if len(args) > 0 {
				return strings.Join(args, " ")
			}
			if hasDefault {
				return defaultVal
			}
			return ""
		}
		if strings.HasSuffix(key, "*") {
			idx := 0
			fmt.Sscanf(key, "%d", &idx)
			if idx >= 1 && idx-1 < len(args) {
				return strings.Join(args[idx-1:], " ")
			}
			if hasDefault {
				return defaultVal
			}
			return ""
		}
		idx := 0
		fmt.Sscanf(key, "%d", &idx)
		if idx >= 1 && idx-1 < len(args) {
			return args[idx-1]
		}
		if hasDefault {
			return defaultVal
		}
		return ""
	})

	return result
}