go-svc-mcp/internal/tools/get_date_range.go

package tools

import (
	"encoding/json"
	"log"
	"time"

	"git.x2erp.com/qdy/go-svc-mcp/internal/mcp"
)

func init() {
	mcp.Register("get_date_range", "根据时间单位和偏移量计算日期范围",
		map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"unit": map[string]interface{}{
					"type":        "string",
					"description": "时间单位",
					"enum":        []string{"day", "week", "month", "quarter", "year"},
				},
				"offset": map[string]interface{}{
					"type":        "integer",
					"description": "偏移量（负数表示过去，正数表示未来）",
					"default":     0,
				},
			},
			"required": []string{"unit", "offset"},
		},
		func(input json.RawMessage, deps *mcp.ToolDependencies) (interface{}, error) {
			var params struct {
				Unit   string `json:"unit"`
				Offset int    `json:"offset"`
			}
			log.Printf("AI 调用工具-get_date_range： %s", time.Now())
			if len(input) > 0 {
				if err := json.Unmarshal(input, &params); err != nil {
					return nil, err
				}
			}

			now := time.Now()
			var startDate, endDate time.Time

			switch params.Unit {
			case "day":
				targetDate := now.AddDate(0, 0, params.Offset)
				startDate = targetDate
				endDate = targetDate

			case "week":
				weekStart := getWeekStart(now)
				targetWeekStart := weekStart.AddDate(0, 0, params.Offset*7)
				targetWeekEnd := targetWeekStart.AddDate(0, 0, 6)
				startDate = targetWeekStart
				endDate = targetWeekEnd

			case "month":
				targetMonth := now.AddDate(0, params.Offset, 0)
				monthStart := time.Date(targetMonth.Year(), targetMonth.Month(), 1, 0, 0, 0, 0, targetMonth.Location())
				monthEnd := monthStart.AddDate(0, 1, -1)
				startDate = monthStart
				endDate = monthEnd

			case "quarter":
				// 基于当前季度偏移
				currentMonth := int(now.Month())
				currentQuarter := (currentMonth-1)/3 + 1
				targetQuarter := currentQuarter + params.Offset
				targetYear := now.Year()

				// 处理跨年
				if targetQuarter < 1 {
					yearsBack := (1 - targetQuarter + 3) / 4
					targetYear -= yearsBack
					targetQuarter += yearsBack * 4
				} else if targetQuarter > 4 {
					yearsForward := (targetQuarter - 1) / 4
					targetYear += yearsForward
					targetQuarter -= yearsForward * 4
				}

				// 计算季度范围
				quarterStartMonth := time.Month((targetQuarter-1)*3 + 1)
				startDate = time.Date(targetYear, quarterStartMonth, 1, 0, 0, 0, 0, now.Location())
				endDate = time.Date(targetYear, quarterStartMonth+3, 0, 0, 0, 0, 0, now.Location())

			case "year":
				targetYear := now.AddDate(params.Offset, 0, 0)
				yearStart := time.Date(targetYear.Year(), 1, 1, 0, 0, 0, 0, targetYear.Location())
				yearEnd := time.Date(targetYear.Year(), 12, 31, 0, 0, 0, 0, targetYear.Location())
				startDate = yearStart
				endDate = yearEnd

			default:
				startDate = now
				endDate = now
			}

			return map[string]interface{}{
				"unit":       params.Unit,
				"offset":     params.Offset,
				"start_date": startDate.Format("2006-01-02"),
				"end_date":   endDate.Format("2006-01-02"),
			}, nil
		},
	)
}

func getWeekStart(t time.Time) time.Time {
	weekday := int(t.Weekday())
	if weekday == 0 {
		return t.AddDate(0, 0, -6)
	}
	return t.AddDate(0, 0, -(weekday - 1))
}