System Prompt Builder — Go source

// Package systempromptbuilder is the Go twin of CosmoDev's
// src/lib/systemPromptBuilder.ts (dual source: the web lib is TypeScript, the
// CLI lib is Go — kept in lock-step). Pure + deterministic, never panics. The
// table-driven tests in systempromptbuilder_test.go share vectors with
// src/lib/systemPromptBuilder.test.ts so the two implementations are held to
// the same contract.
//
// The package assembles an ordered list of prompt blocks into a
// markdown-structured system prompt, with pure list operations, presets,
// warnings, and a compact URL codec for shareable state. Token counting
// reuses the token-estimator twin at the island layer; the soft-limit constant
// lives here as the domain rule.
package systempromptbuilder

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"strconv"
	"strings"

	tokenestimator "cosmodev/token-estimator"
)

// PromptBlock is one section of a system prompt. Field-for-field twin of the
// TS PromptBlock interface (camelCase JSON tags, aimodels convention).
type PromptBlock struct {
	ID      string `json:"id"`
	Title   string `json:"title"`
	Content string `json:"content"`
	Enabled bool   `json:"enabled"`
}

// PromptPreset is one ordered starter template. Mirrors the TS PromptPreset
// interface.
type PromptPreset struct {
	ID          string `json:"id"`
	Title       string `json:"title"`
	Description string `json:"description"`
	Content     string `json:"content"`
}

// SystemPromptSoftLimitTokens is the assembled size beyond which blocks start
// crowding the context on most models. Mirrors SYSTEM_PROMPT_SOFT_LIMIT_TOKENS.
const SystemPromptSoftLimitTokens = 2000

// SystemPromptPresets are the ordered starter templates — the recommended
// skeleton of a system prompt. Mirrors SYSTEM_PROMPT_PRESETS.
var SystemPromptPresets = []PromptPreset{
	{
		ID:          "role",
		Title:       "Role",
		Description: "Who the model is and what it optimizes for.",
		Content:     "You are a senior software engineer. You give correct, concise answers and say so plainly when you are unsure.",
	},
	{
		ID:          "context",
		Title:       "Context",
		Description: "The situation the model is working in.",
		Content:     "The user is a developer working in a TypeScript codebase. Prefer runnable examples over prose when both work.",
	},
	{
		ID:          "constraints",
		Title:       "Constraints",
		Description: "Hard rules the model must not break.",
		Content:     "- Never invent library APIs; use only the ones in the provided code.\n- Keep answers under 300 words unless asked for more.",
	},
	{
		ID:          "output-format",
		Title:       "Output format",
		Description: "The exact shape of the answer.",
		Content:     "Respond with: 1) a one-line summary, 2) a fenced code block, 3) any caveats as bullet points.",
	},
	{
		ID:          "examples",
		Title:       "Examples",
		Description: "Few-shot demonstrations of the desired behavior.",
		Content:     "Input: reverse \"abc\"\nOutput: \"cba\"",
	},
	{
		ID:          "tone",
		Title:       "Tone",
		Description: "Voice and register.",
		Content:     "Direct and friendly. No filler openers, no apologies.",
	},
	{
		ID:          "refusal",
		Title:       "Refusal policy",
		Description: "How to handle out-of-scope requests.",
		Content:     "If a request is outside your scope, say so in one sentence and suggest the closest thing you can do.",
	},
	{
		ID:          "safety",
		Title:       "Safety",
		Description: "Guardrails for sensitive content.",
		Content:     "Refuse requests that could cause harm, and never echo secrets, keys, or credentials back in full.",
	},
}

// AssembleOptions configures AssemblePrompt. The zero value (AssembleOptions{})
// matches the TS default (assemblePrompt(blocks) with no options): headers on.
//
// Headers is a *bool so the zero value means "default true" — exactly like the
// TS lib's distinction between omitted (→ true) and false (→ raw contents).
type AssembleOptions struct {
	Headers *bool // nil → true (default); non-nil used verbatim
}

// BlockPatch is the partial update applied by UpdateBlock — the twin of TS's
// Partial<Omit<PromptBlock, 'id'>>. A nil field leaves that field untouched.
type BlockPatch struct {
	Title   *string
	Content *string
	Enabled *bool
}

// AssemblePrompt renders enabled, non-empty blocks (in order) as one
// markdown-structured prompt. It is the Go twin of assemblePrompt() in
// src/lib/systemPromptBuilder.ts and must agree with it on every shared vector.
func AssemblePrompt(blocks []PromptBlock, opts AssembleOptions) string {
	headers := true
	if opts.Headers != nil {
		headers = *opts.Headers
	}
	parts := make([]string, 0, len(blocks))
	for _, b := range blocks {
		if !b.Enabled || strings.TrimSpace(b.Content) == "" {
			continue
		}
		if headers {
			title := strings.TrimSpace(b.Title)
			if title == "" {
				title = "Untitled"
			}
			parts = append(parts, "## "+title+"\n"+strings.TrimSpace(b.Content))
		} else {
			parts = append(parts, strings.TrimSpace(b.Content))
		}
	}
	return strings.TrimSpace(strings.Join(parts, "\n\n"))
}

// AddBlock appends a block (the caller supplies the id so the lib stays pure).
// The trailing bool mirrors TS's enabled = true default; pass it explicitly.
func AddBlock(blocks []PromptBlock, id, title, content string, enabled bool) []PromptBlock {
	next := make([]PromptBlock, len(blocks), len(blocks)+1)
	copy(next, blocks)
	return append(next, PromptBlock{ID: id, Title: title, Content: content, Enabled: enabled})
}

// UpdateBlock patches one block by id; unknown ids leave the list unchanged.
func UpdateBlock(blocks []PromptBlock, id string, patch BlockPatch) []PromptBlock {
	next := make([]PromptBlock, len(blocks))
	copy(next, blocks)
	for i := range next {
		if next[i].ID != id {
			continue
		}
		if patch.Title != nil {
			next[i].Title = *patch.Title
		}
		if patch.Content != nil {
			next[i].Content = *patch.Content
		}
		if patch.Enabled != nil {
			next[i].Enabled = *patch.Enabled
		}
	}
	return next
}

// ToggleBlock flips one block's enabled flag by id.
func ToggleBlock(blocks []PromptBlock, id string) []PromptBlock {
	next := make([]PromptBlock, len(blocks))
	copy(next, blocks)
	for i := range next {
		if next[i].ID == id {
			next[i].Enabled = !next[i].Enabled
		}
	}
	return next
}

// RemoveBlock removes one block by id.
func RemoveBlock(blocks []PromptBlock, id string) []PromptBlock {
	next := make([]PromptBlock, 0, len(blocks))
	for _, b := range blocks {
		if b.ID != id {
			next = append(next, b)
		}
	}
	return next
}

// MoveBlock moves a block (clamped; a copy of the input when indexes are out
// of range or equal).
func MoveBlock(blocks []PromptBlock, from, to int) []PromptBlock {
	if from < 0 || from >= len(blocks) || to < 0 || to >= len(blocks) || from == to {
		out := make([]PromptBlock, len(blocks))
		copy(out, blocks)
		return out
	}
	next := make([]PromptBlock, len(blocks))
	copy(next, blocks)
	moved := next[from]
	next = append(next[:from], next[from+1:]...)
	rest := make([]PromptBlock, 0, len(next)+1)
	rest = append(rest, next[:to]...)
	rest = append(rest, moved)
	rest = append(rest, next[to:]...)
	return rest
}

// PromptReport is the assembled prompt plus its lint results. Field-for-field
// twin of the TS PromptReport interface.
type PromptReport struct {
	Assembled string   `json:"assembled"`
	Tokens    int      `json:"tokens"`
	Warnings  []string `json:"warnings"`
}

// BuildReport assembles + counts + lints in one pass — the island's live
// report. The zero contentType ("") means "prose", the TS default; pass
// tokenestimator.Auto for explicit auto-detection.
func BuildReport(blocks []PromptBlock, contentType tokenestimator.AutoType) PromptReport {
	// The TS default parameter is 'prose' (not auto). AutoType's Go zero
	// value is "", so map it to prose; an explicit "auto" is passed through
	// and detected per line by the estimator.
	ct := contentType
	if ct == "" {
		ct = tokenestimator.TypeProse
	}
	assembled := AssemblePrompt(blocks, AssembleOptions{})
	tokens := 0
	if assembled != "" {
		tokens = tokenestimator.EstimateTokens(assembled, &tokenestimator.EstimateOptions{ContentType: ct}).Tokens
	}
	warnings := []string{}
	if tokens > SystemPromptSoftLimitTokens {
		warnings = append(warnings,
			"Assembled prompt is ~"+groupDigits(tokens)+" tokens — beyond "+groupDigits(SystemPromptSoftLimitTokens)+
				" it starts crowding the context window on most models.")
	}
	if len(blocks) > 0 && !hasEnabledRole(blocks) {
		warnings = append(warnings,
			`No enabled "Role" block — stating who the model is tends to anchor every following instruction.`)
	}
	if len(blocks) > 0 && assembled == "" {
		warnings = append(warnings, "Every block is disabled or empty — the assembled prompt is empty.")
	}
	return PromptReport{Assembled: assembled, Tokens: tokens, Warnings: warnings}
}

func hasEnabledRole(blocks []PromptBlock) bool {
	for _, b := range blocks {
		if b.Enabled && strings.ToLower(strings.TrimSpace(b.Title)) == "role" {
			return true
		}
	}
	return false
}

// groupDigits renders n with en-US thousands separators, the format TS's
// toLocaleString('en-US') produces (2000 → "2,000").
func groupDigits(n int) string {
	s := strconv.Itoa(n)
	sign := ""
	if strings.HasPrefix(s, "-") {
		sign, s = "-", s[1:]
	}
	if len(s) <= 3 {
		return sign + s
	}
	var groups []string
	for len(s) > 3 {
		groups = append([]string{s[len(s)-3:]}, groups...)
		s = s[:len(s)-3]
	}
	return sign + s + "," + strings.Join(groups, ",")
}

// ---- shareable state codec (URL-safe, compact) ------------------------------
// Triples of [enabled(0/1), title, content] keep URLs far smaller than the
// full object shape; ids are regenerated on decode (they are UI-local).

// maxEncodedLength is the length beyond which the encoded form would make an
// uncomfortably long URL. Mirrors MAX_ENCODED_LENGTH.
const maxEncodedLength = 4000

// ErrMalformedState is returned by DecodeBlocks on any malformed input — the
// Go stand-in for the TS lib's null return (never throws).
var ErrMalformedState = errors.New("malformed encoded block state")

func toBase64URL(s string) string {
	// RawURLEncoding is unpadded and URL-safe: exactly the TS pipeline
	// (btoa → replace +→- , /→_ , strip trailing '=').
	return base64.RawURLEncoding.EncodeToString([]byte(s))
}

func fromBase64URL(s string) (string, error) {
	// Mirror of the TS pipeline: restore the standard alphabet, then pad.
	std := strings.NewReplacer("-", "+", "_", "/").Replace(strings.TrimRight(s, "="))
	b, err := base64.RawStdEncoding.DecodeString(std)
	if err != nil {
		return "", err
	}
	return string(b), nil
}

// EncodeBlocks encodes blocks to a compact base64url string; "" when blocks
// are empty.
func EncodeBlocks(blocks []PromptBlock) string {
	if len(blocks) == 0 {
		return ""
	}
	compact := make([][]any, 0, len(blocks))
	for _, b := range blocks {
		enabled := 0
		if b.Enabled {
			enabled = 1
		}
		compact = append(compact, []any{enabled, b.Title, b.Content})
	}
	// SetEscapeHTML(false) keeps the JSON byte-identical to JSON.stringify,
	// so the encoded strings match the TS lib character for character.
	var buf bytes.Buffer
	enc := json.NewEncoder(&buf)
	enc.SetEscapeHTML(false)
	if err := enc.Encode(compact); err != nil {
		return ""
	}
	return toBase64URL(strings.TrimSuffix(buf.String(), "\n"))
}

// EncodedTooLong reports whether the encoded form would make an
// uncomfortably long URL.
func EncodedTooLong(encoded string) bool {
	return len(encoded) > maxEncodedLength
}

// DecodeBlocks reverses EncodeBlocks; it regenerates ids (b1, b2, …) and
// returns ErrMalformedState on malformed input — never panics.
func DecodeBlocks(encoded string) ([]PromptBlock, error) {
	if encoded == "" {
		return []PromptBlock{}, nil
	}
	raw, err := fromBase64URL(encoded)
	if err != nil {
		return nil, ErrMalformedState
	}
	var top any
	if err := json.Unmarshal([]byte(raw), &top); err != nil {
		return nil, ErrMalformedState
	}
	arr, ok := top.([]any)
	if !ok {
		return nil, ErrMalformedState
	}
	blocks := make([]PromptBlock, 0, len(arr))
	for i, item := range arr {
		entry, ok := item.([]any)
		if !ok || len(entry) != 3 {
			return nil, ErrMalformedState
		}
		enabledNum, ok := entry[0].(float64)
		if !ok {
			return nil, ErrMalformedState
		}
		title, ok := entry[1].(string)
		if !ok {
			return nil, ErrMalformedState
		}
		content, ok := entry[2].(string)
		if !ok {
			return nil, ErrMalformedState
		}
		blocks = append(blocks, PromptBlock{
			ID:      "b" + strconv.Itoa(i+1),
			Title:   title,
			Content: content,
			Enabled: enabledNum == 1,
		})
	}
	return blocks, nil
}