Cyrene/backend/ai-core/internal/tools/random_tool.go

package tools

import (
	"context"
	"crypto/rand"
	"encoding/json"
	"fmt"
	"math/big"
	mathrand "math/rand"
	"strings"
)

// RandomTool provides random generation utilities for the LLM.
// Supports random numbers, UUIDs, passwords, and list operations.
type RandomTool struct{}

// NewRandomTool creates a random generation tool.
func NewRandomTool() *RandomTool {
	return &RandomTool{}
}

// Definition returns the tool definition for LLM function calling.
func (t *RandomTool) Definition() ToolDefinition {
	return ToolDefinition{
		Name:        "random",
		Description: "随机生成工具。生成随机数、UUID、安全密码，或从列表中随机选取/打乱元素。",
		Parameters: map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"action": map[string]interface{}{
					"type":        "string",
					"enum":        []string{"number", "uuid", "password", "pick", "shuffle"},
					"description": "操作类型。number: 生成随机整数；uuid: 生成UUID v4；password: 生成安全密码；pick: 从列表随机选取；shuffle: 随机打乱列表",
				},
				"min": map[string]interface{}{
					"type":        "number",
					"description": "随机数最小值（用于 number 操作），默认 0",
				},
				"max": map[string]interface{}{
					"type":        "number",
					"description": "随机数最大值（用于 number 操作），默认 100",
				},
				"length": map[string]interface{}{
					"type":        "integer",
					"description": "密码长度（用于 password 操作），默认 16",
				},
				"items": map[string]interface{}{
					"type":        "array",
					"description": "列表项（用于 pick/shuffle 操作），字符串数组",
					"items": map[string]interface{}{
						"type": "string",
					},
				},
				"count": map[string]interface{}{
					"type":        "integer",
					"description": "选取数量（用于 pick 操作），默认 1",
				},
			},
			"required": []string{"action"},
		},
	}
}

// Execute performs random generation operations.
func (t *RandomTool) Execute(ctx context.Context, arguments map[string]interface{}) (*ToolResult, error) {
	action, ok := arguments["action"].(string)
	if !ok || action == "" {
		return &ToolResult{
			ToolName: "random",
			Success:  false,
			Error:    "缺少 action 参数",
		}, nil
	}

	switch action {
	case "number":
		return t.handleNumber(arguments)
	case "uuid":
		return t.handleUUID()
	case "password":
		return t.handlePassword(arguments)
	case "pick":
		return t.handlePick(arguments)
	case "shuffle":
		return t.handleShuffle(arguments)
	default:
		return &ToolResult{
			ToolName: "random",
			Success:  false,
			Error:    fmt.Sprintf("未知操作: %s，支持: number, uuid, password, pick, shuffle", action),
		}, nil
	}
}

// handleNumber generates a random integer in [min, max].
func (t *RandomTool) handleNumber(arguments map[string]interface{}) (*ToolResult, error) {
	minVal := getFloatArg(arguments, "min", 0)
	maxVal := getFloatArg(arguments, "max", 100)

	if minVal > maxVal {
		minVal, maxVal = maxVal, minVal
	}

	minI := int64(minVal)
	maxI := int64(maxVal)

	// Use crypto/rand for secure random
	rangeVal := maxI - minI + 1
	if rangeVal <= 0 {
		return &ToolResult{
			ToolName: "random",
			Success:  false,
			Error:    "无效的数值范围",
		}, nil
	}

	n, err := rand.Int(rand.Reader, big.NewInt(rangeVal))
	if err != nil {
		// Fallback to math/rand
		result := minI + mathrand.Int63n(rangeVal)
		return &ToolResult{
			ToolName: "random",
			Success:  true,
			Data:     fmt.Sprintf("随机整数 [%d, %d]: %d", minI, maxI, result),
		}, nil
	}

	result := minI + n.Int64()
	return &ToolResult{
		ToolName: "random",
		Success:  true,
		Data:     fmt.Sprintf("随机整数 [%d, %d]: %d", minI, maxI, result),
	}, nil
}

// handleUUID generates a UUID v4 string.
func (t *RandomTool) handleUUID() (*ToolResult, error) {
	uuid := make([]byte, 16)
	_, err := rand.Read(uuid)
	if err != nil {
		return &ToolResult{
			ToolName: "random",
			Success:  false,
			Error:    fmt.Sprintf("生成UUID失败: %v", err),
		}, nil
	}

	// Set version 4 and variant bits
	uuid[6] = (uuid[6] & 0x0f) | 0x40 // Version 4
	uuid[8] = (uuid[8] & 0x3f) | 0x80 // Variant 10

	uuidStr := fmt.Sprintf("%08x-%04x-%04x-%04x-%012x",
		uuid[0:4], uuid[4:6], uuid[6:8], uuid[8:10], uuid[10:16])

	return &ToolResult{
		ToolName: "random",
		Success:  true,
		Data:     fmt.Sprintf("UUID v4: %s", uuidStr),
	}, nil
}

// handlePassword generates a secure random password.
func (t *RandomTool) handlePassword(arguments map[string]interface{}) (*ToolResult, error) {
	length := getIntArg(arguments, "length", 16)
	if length < 4 {
		length = 16
	}
	if length > 128 {
		length = 128
	}

	uppercase := "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
	lowercase := "abcdefghijklmnopqrstuvwxyz"
	digits := "0123456789"
	symbols := "!@#$%^&*()_+-=[]{}|;:,.<>?"

	allChars := uppercase + lowercase + digits + symbols

	password := make([]byte, length)

	// Ensure at least one of each character type
	password[0] = uppercase[secureIndex(len(uppercase))]
	password[1] = lowercase[secureIndex(len(lowercase))]
	password[2] = digits[secureIndex(len(digits))]
	password[3] = symbols[secureIndex(len(symbols))]

	// Fill remaining with random characters from all sets
	for i := 4; i < length; i++ {
		password[i] = allChars[secureIndex(len(allChars))]
	}

	// Shuffle the password
	shuffleBytes(password)

	passwordStr := string(password)

	return &ToolResult{
		ToolName: "random",
		Success:  true,
		Data: fmt.Sprintf("安全密码 (长度: %d):\n%s\n\n字符集: 大写字母 + 小写字母 + 数字 + 特殊符号",
			length, passwordStr),
	}, nil
}

// handlePick randomly picks items from a list.
func (t *RandomTool) handlePick(arguments map[string]interface{}) (*ToolResult, error) {
	items := getStringSliceArg(arguments, "items")
	if len(items) == 0 {
		return &ToolResult{
			ToolName: "random",
			Success:  false,
			Error:    "缺少 items 参数或列表为空",
		}, nil
	}

	count := getIntArg(arguments, "count", 1)
	if count < 1 {
		count = 1
	}
	if count > len(items) {
		count = len(items)
	}

	// Shuffle indices and pick first 'count'
	indices := make([]int, len(items))
	for i := range indices {
		indices[i] = i
	}
	shuffleInts(indices)

	picked := make([]string, 0, count)
	for i := 0; i < count; i++ {
		picked = append(picked, items[indices[i]])
	}

	var result strings.Builder
	result.WriteString(fmt.Sprintf("从 %d 个选项中随机选取 %d 个:\n", len(items), count))
	for i, p := range picked {
		result.WriteString(fmt.Sprintf("  %d. %s\n", i+1, p))
	}

	return &ToolResult{
		ToolName: "random",
		Success:  true,
		Data:     result.String(),
	}, nil
}

// handleShuffle randomly shuffles a list.
func (t *RandomTool) handleShuffle(arguments map[string]interface{}) (*ToolResult, error) {
	items := getStringSliceArg(arguments, "items")
	if len(items) == 0 {
		return &ToolResult{
			ToolName: "random",
			Success:  false,
			Error:    "缺少 items 参数或列表为空",
		}, nil
	}

	// Make a copy and shuffle
	shuffled := make([]string, len(items))
	copy(shuffled, items)
	shuffleStrings(shuffled)

	var result strings.Builder
	result.WriteString(fmt.Sprintf("随机打乱结果 (共 %d 项):\n", len(shuffled)))
	for i, s := range shuffled {
		result.WriteString(fmt.Sprintf("  %d. %s\n", i+1, s))
	}

	return &ToolResult{
		ToolName: "random",
		Success:  true,
		Data:     result.String(),
	}, nil
}

// --- Helper functions ---

// getFloatArg extracts a float64 argument with fallback.
func getFloatArg(arguments map[string]interface{}, key string, fallback float64) float64 {
	if v, ok := arguments[key]; ok {
		switch val := v.(type) {
		case float64:
			return val
		case int:
			return float64(val)
		case int64:
			return float64(val)
		case json.Number:
			f, err := val.Float64()
			if err == nil {
				return f
			}
		}
	}
	return fallback
}

// getIntArg extracts an int argument with fallback.
func getIntArg(arguments map[string]interface{}, key string, fallback int) int {
	if v, ok := arguments[key]; ok {
		switch val := v.(type) {
		case float64:
			return int(val)
		case int:
			return val
		case int64:
			return int(val)
		}
	}
	return fallback
}

// getStringSliceArg extracts a string slice argument.
func getStringSliceArg(arguments map[string]interface{}, key string) []string {
	if v, ok := arguments[key]; ok {
		switch val := v.(type) {
		case []interface{}:
			result := make([]string, 0, len(val))
			for _, item := range val {
				if s, ok := item.(string); ok {
					result = append(result, s)
				} else {
					result = append(result, fmt.Sprintf("%v", item))
				}
			}
			return result
		case []string:
			return val
		}
	}
	return nil
}

// secureIndex returns a cryptographically secure random index in [0, max).
func secureIndex(max int) int {
	if max <= 1 {
		return 0
	}
	n, err := rand.Int(rand.Reader, big.NewInt(int64(max)))
	if err != nil {
		return mathrand.Intn(max)
	}
	return int(n.Int64())
}

// shuffleBytes shuffles a byte slice using Fisher-Yates with crypto/rand.
func shuffleBytes(data []byte) {
	for i := len(data) - 1; i > 0; i-- {
		j := secureIndex(i + 1)
		data[i], data[j] = data[j], data[i]
	}
}

// shuffleInts shuffles an int slice using Fisher-Yates with crypto/rand.
func shuffleInts(data []int) {
	for i := len(data) - 1; i > 0; i-- {
		j := secureIndex(i + 1)
		data[i], data[j] = data[j], data[i]
	}
}

// shuffleStrings shuffles a string slice using Fisher-Yates with crypto/rand.
func shuffleStrings(data []string) {
	for i := len(data) - 1; i > 0; i-- {
		j := secureIndex(i + 1)
		data[i], data[j] = data[j], data[i]
	}
}