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

package tools

import (
	"context"
	"fmt"
	"strconv"
	"strings"
	"time"
	"unicode"
)

// DateTimeTool provides date/time operations for the LLM.
// Supports current time, formatting, date arithmetic, and timezone listing.
type DateTimeTool struct{}

// NewDateTimeTool creates a date/time tool.
func NewDateTimeTool() *DateTimeTool {
	return &DateTimeTool{}
}

// Definition returns the tool definition for LLM function calling.
func (t *DateTimeTool) Definition() ToolDefinition {
	return ToolDefinition{
		Name:        "datetime",
		Description: "日期时间工具。获取当前时间、格式化日期、日期加减、计算日期差、查看可用时区。",
		Parameters: map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"action": map[string]interface{}{
					"type":        "string",
					"enum":        []string{"now", "format", "add", "diff", "timezone_list"},
					"description": "操作类型。now: 获取当前时间；format: 格式化日期；add: 日期加减；diff: 计算两个日期的差值；timezone_list: 列出常用时区",
				},
				"format": map[string]interface{}{
					"type":        "string",
					"description": "日期格式串（Go风格）。默认 \"2006-01-02 15:04:05\"。常用: \"2006-01-02\"（仅日期）、\"15:04:05\"（仅时间）",
				},
				"timezone": map[string]interface{}{
					"type":        "string",
					"description": "时区标识，如 \"Asia/Shanghai\"、\"America/New_York\"、\"UTC\"。默认使用服务器本地时区",
				},
				"date": map[string]interface{}{
					"type":        "string",
					"description": "基准日期，格式为 \"2006-01-02 15:04:05\" 或 \"2006-01-02\"",
				},
				"duration": map[string]interface{}{
					"type":        "string",
					"description": "时长字符串，如 \"24h\"、\"7d\"、\"30m\"、\"1h30m\"。支持单位: s(秒), m(分钟), h(小时), d(天), w(周), M(月), y(年)",
				},
				"date2": map[string]interface{}{
					"type":        "string",
					"description": "第二个日期（用于 diff 操作），格式同 date",
				},
			},
			"required": []string{"action"},
		},
	}
}

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

	switch action {
	case "now":
		return t.handleNow(arguments)
	case "format":
		return t.handleFormat(arguments)
	case "add":
		return t.handleAdd(arguments)
	case "diff":
		return t.handleDiff(arguments)
	case "timezone_list":
		return t.handleTimezoneList()
	default:
		return &ToolResult{
			ToolName: "datetime",
			Success:  false,
			Error:    fmt.Sprintf("未知操作: %s，支持: now, format, add, diff, timezone_list", action),
		}, nil
	}
}

// handleNow returns the current date/time in the specified timezone.
func (t *DateTimeTool) handleNow(arguments map[string]interface{}) (*ToolResult, error) {
	tz, err := t.getTimezone(arguments)
	if err != nil {
		return &ToolResult{
			ToolName: "datetime",
			Success:  false,
			Error:    err.Error(),
		}, nil
	}

	format := t.getFormat(arguments)
	now := time.Now().In(tz)

	return &ToolResult{
		ToolName: "datetime",
		Success:  true,
		Data: fmt.Sprintf("当前时间: %s\n时区: %s\nUnix时间戳: %d",
			now.Format(format), tz.String(), now.Unix()),
	}, nil
}

// handleFormat formats a given date string.
func (t *DateTimeTool) handleFormat(arguments map[string]interface{}) (*ToolResult, error) {
	dateStr, _ := arguments["date"].(string)
	if dateStr == "" {
		return &ToolResult{
			ToolName: "datetime",
			Success:  false,
			Error:    "format 操作需要 date 参数",
		}, nil
	}

	parsed, err := t.parseDate(dateStr)
	if err != nil {
		return &ToolResult{
			ToolName: "datetime",
			Success:  false,
			Error:    fmt.Sprintf("日期解析失败: %v", err),
		}, nil
	}

	tz, err := t.getTimezone(arguments)
	if err != nil {
		return &ToolResult{
			ToolName: "datetime",
			Success:  false,
			Error:    err.Error(),
		}, nil
	}

	format := t.getFormat(arguments)
	formatted := parsed.In(tz).Format(format)

	return &ToolResult{
		ToolName: "datetime",
		Success:  true,
		Data:     fmt.Sprintf("原始: %s\n格式化: %s\n时区: %s", dateStr, formatted, tz.String()),
	}, nil
}

// handleAdd adds/subtracts a duration from a date.
func (t *DateTimeTool) handleAdd(arguments map[string]interface{}) (*ToolResult, error) {
	durationStr, _ := arguments["duration"].(string)
	if durationStr == "" {
		return &ToolResult{
			ToolName: "datetime",
			Success:  false,
			Error:    "add 操作需要 duration 参数",
		}, nil
	}

	dateStr, _ := arguments["date"].(string)
	var base time.Time
	if dateStr != "" {
		var err error
		base, err = t.parseDate(dateStr)
		if err != nil {
			return &ToolResult{
				ToolName: "datetime",
				Success:  false,
				Error:    fmt.Sprintf("日期解析失败: %v", err),
			}, nil
		}
	} else {
		tz, _ := t.getTimezone(arguments)
		base = time.Now().In(tz)
	}

	dur, err := t.parseDuration(durationStr)
	if err != nil {
		return &ToolResult{
			ToolName: "datetime",
			Success:  false,
			Error:    fmt.Sprintf("时长解析失败: %v", err),
		}, nil
	}

	tz, _ := t.getTimezone(arguments)

	result := base.In(tz)

	// Extract months and years from the duration string (not handled by time.Duration)
	months := extractDurationUnit(durationStr, 'M')
	years := extractDurationUnit(durationStr, 'y')

	if months != 0 || years != 0 {
		result = result.AddDate(years, months, 0)
	}

	// Add the standard duration part
	if dur != 0 {
		result = result.Add(dur)
	}

	format := t.getFormat(arguments)
	return &ToolResult{
		ToolName: "datetime",
		Success:  true,
		Data: fmt.Sprintf("基准日期: %s\n操作: %s\n结果: %s",
			base.In(tz).Format(format), durationStr, result.Format(format)),
	}, nil
}

// handleDiff calculates the difference between two dates.
func (t *DateTimeTool) handleDiff(arguments map[string]interface{}) (*ToolResult, error) {
	dateStr, _ := arguments["date"].(string)
	date2Str, _ := arguments["date2"].(string)

	if dateStr == "" || date2Str == "" {
		return &ToolResult{
			ToolName: "datetime",
			Success:  false,
			Error:    "diff 操作需要 date 和 date2 参数",
		}, nil
	}

	d1, err := t.parseDate(dateStr)
	if err != nil {
		return &ToolResult{
			ToolName: "datetime",
			Success:  false,
			Error:    fmt.Sprintf("date 解析失败: %v", err),
		}, nil
	}

	d2, err := t.parseDate(date2Str)
	if err != nil {
		return &ToolResult{
			ToolName: "datetime",
			Success:  false,
			Error:    fmt.Sprintf("date2 解析失败: %v", err),
		}, nil
	}

	diff := d2.Sub(d1)
	absDiff := diff
	if absDiff < 0 {
		absDiff = -absDiff
	}

	days := int(absDiff.Hours() / 24)
	hours := int(absDiff.Hours()) % 24
	minutes := int(absDiff.Minutes()) % 60
	seconds := int(absDiff.Seconds()) % 60

	sign := ""
	if diff < 0 {
		sign = "-"
	}

	return &ToolResult{
		ToolName: "datetime",
		Success:  true,
		Data: fmt.Sprintf("日期1: %s\n日期2: %s\n差值: %s%d天 %d小时 %d分钟 %d秒 (总计 %s%.0f秒)",
			dateStr, date2Str, sign, days, hours, minutes, seconds, sign, absDiff.Seconds()),
	}, nil
}

// handleTimezoneList returns a list of common timezones.
func (t *DateTimeTool) handleTimezoneList() (*ToolResult, error) {
	zones := []string{
		"UTC",
		"Asia/Shanghai (北京时间)",
		"Asia/Tokyo (东京时间)",
		"Asia/Seoul (首尔时间)",
		"Asia/Singapore (新加坡时间)",
		"Asia/Kolkata (印度时间)",
		"Asia/Dubai (迪拜时间)",
		"Europe/London (伦敦时间)",
		"Europe/Paris (巴黎时间)",
		"Europe/Berlin (柏林时间)",
		"Europe/Moscow (莫斯科时间)",
		"America/New_York (纽约时间)",
		"America/Chicago (芝加哥时间)",
		"America/Denver (丹佛时间)",
		"America/Los_Angeles (洛杉矶时间)",
		"America/Sao_Paulo (圣保罗时间)",
		"Australia/Sydney (悉尼时间)",
		"Pacific/Auckland (奥克兰时间)",
	}

	var result strings.Builder
	result.WriteString("常用时区列表:\n\n")
	for i, z := range zones {
		result.WriteString(fmt.Sprintf("  %2d. %s\n", i+1, z))
	}

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

// getTimezone extracts the timezone from arguments, defaulting to Asia/Shanghai.
func (t *DateTimeTool) getTimezone(arguments map[string]interface{}) (*time.Location, error) {
	tzName, _ := arguments["timezone"].(string)
	if tzName == "" {
		loc, err := time.LoadLocation("Asia/Shanghai")
		if err != nil {
			return time.Local, nil
		}
		return loc, nil
	}
	loc, err := time.LoadLocation(tzName)
	if err != nil {
		return nil, fmt.Errorf("无效时区: %s", tzName)
	}
	return loc, nil
}

// getFormat extracts the format string from arguments, defaulting to standard format.
func (t *DateTimeTool) getFormat(arguments map[string]interface{}) string {
	format, _ := arguments["format"].(string)
	if format == "" {
		return "2006-01-02 15:04:05"
	}
	return format
}

// parseDate parses a date string with multiple format attempts.
func (t *DateTimeTool) parseDate(s string) (time.Time, error) {
	formats := []string{
		"2006-01-02 15:04:05",
		"2006-01-02T15:04:05Z",
		"2006-01-02T15:04:05",
		"2006-01-02",
		"2006/01/02 15:04:05",
		"2006/01/02",
		time.RFC3339,
		time.RFC3339Nano,
	}
	for _, f := range formats {
		if t, err := time.Parse(f, s); err == nil {
			return t, nil
		}
	}
	return time.Time{}, fmt.Errorf("无法解析日期: %s", s)
}

// parseDuration parses a human-friendly duration string like "24h", "7d", "1h30m".
func (t *DateTimeTool) parseDuration(s string) (time.Duration, error) {
	// First try standard Go duration parsing
	if d, err := time.ParseDuration(s); err == nil {
		return d, nil
	}

	// Custom parsing for days and weeks
	var total time.Duration
	remaining := s

	for len(remaining) > 0 {
		// find the number
		numStart := 0
		for numStart < len(remaining) && !unicode.IsDigit(rune(remaining[numStart])) && remaining[numStart] != '-' {
			numStart++
		}
		if numStart >= len(remaining) {
			break
		}

		numEnd := numStart
		for numEnd < len(remaining) && (unicode.IsDigit(rune(remaining[numEnd])) || remaining[numEnd] == '.') {
			numEnd++
		}

		val, err := strconv.ParseFloat(remaining[numStart:numEnd], 64)
		if err != nil {
			return 0, fmt.Errorf("无效时长数字: %s", remaining[numStart:numEnd])
		}

		unitEnd := numEnd
		for unitEnd < len(remaining) && unicode.IsLetter(rune(remaining[unitEnd])) {
			unitEnd++
		}
		unit := remaining[numEnd:unitEnd]

		switch unit {
		case "s":
			total += time.Duration(val * float64(time.Second))
		case "m":
			total += time.Duration(val * float64(time.Minute))
		case "h":
			total += time.Duration(val * float64(time.Hour))
		case "d":
			total += time.Duration(val * 24 * float64(time.Hour))
		case "w":
			total += time.Duration(val * 7 * 24 * float64(time.Hour))
		default:
			// skip unknown units (M and y handled elsewhere)
		}

		remaining = remaining[unitEnd:]
	}

	return total, nil
}

// extractDurationUnit extracts numeric value for a given unit character from a duration string.
// e.g., extractDurationUnit("3M", 'M') returns 3, extractDurationUnit("1y2M", 'y') returns 1.
func extractDurationUnit(s string, unit byte) int {
	for i := 0; i < len(s); i++ {
		if s[i] == unit {
			// Scan backwards to find the start of the number
			j := i - 1
			for j >= 0 && (unicode.IsDigit(rune(s[j])) || s[j] == '.') {
				j--
			}
			numStr := s[j+1 : i]
			val, err := strconv.Atoi(numStr)
			if err != nil {
				return 0
			}
			return val
		}
	}
	return 0
}