ttychart-mcp/tools/provider.go

// Package tools ttychart-mcp 工具实现包
//
// 该包提供所有 MCP 工具的具体实现，包括:
//   - plot_line: 折线图
//   - plot_bar: 柱状图
//   - plot_scatter: 散点图
//
// 作者: titor
// 创建日期: 2026-04-15
package tools

import (
	"context"
	"encoding/json"
	"fmt"
	"strconv"
	"strings"

	"hub.gaomia.site/titor/ttychart-mcp/tools/charts"

	"charm.land/lipgloss/v2"
	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
)

// NewServer 创建新的 MCP 服务器
func NewServer() *server.MCPServer {
	srv := server.NewMCPServer(
		"ttychart",
		"0.1.0",
		server.WithToolCapabilities(true),
	)

	// 折线图工具
	plotLineTool := mcp.NewTool(
		"plot_line",
		mcp.WithDescription("绘制折线图，用于展示数据随时间或类别的变化趋势"),
		mcp.WithNumber("data",
			mcp.Required(),
			mcp.Description("Y轴数据点数组"),
		),
		mcp.WithString("title",
			mcp.Description("图表标题"),
		),
		mcp.WithNumber("width",
			mcp.Description("图表宽度（字符数），默认自动计算"),
		),
		mcp.WithNumber("height",
			mcp.Description("图表高度（行数），默认自动计算"),
		),
		mcp.WithString("labels",
			mcp.Description("X轴标签，多个标签用逗号分隔"),
		),
		mcp.WithString("color",
			mcp.Description("线条颜色，可选值: red, green, blue, yellow, cyan, magenta"),
		),
	)

	// 柱状图工具
	plotBarTool := mcp.NewTool(
		"plot_bar",
		mcp.WithDescription("绘制柱状图，用于展示分类数据的比较"),
		mcp.WithString("data",
			mcp.Required(),
			mcp.Description("数据，格式: 标签1:数值1,标签2:数值2"),
		),
		mcp.WithString("title",
			mcp.Description("图表标题"),
		),
		mcp.WithNumber("width",
			mcp.Description("图表宽度"),
		),
		mcp.WithNumber("height",
			mcp.Description("图表高度"),
		),
		mcp.WithBoolean("horizontal",
			mcp.Description("是否为水平柱状图，默认为垂直"),
		),
		mcp.WithString("color",
			mcp.Description("柱体颜色"),
		),
	)

	// 散点图工具
	plotScatterTool := mcp.NewTool(
		"plot_scatter",
		mcp.WithDescription("绘制散点图，用于展示数据的相关性"),
		mcp.WithString("data",
			mcp.Required(),
			mcp.Description("数据点，格式: x1,y1 x2,y2 ..."),
		),
		mcp.WithString("title",
			mcp.Description("图表标题"),
		),
		mcp.WithNumber("width",
			mcp.Description("图表宽度"),
		),
		mcp.WithNumber("height",
			mcp.Description("图表高度"),
		),
		mcp.WithString("x_label",
			mcp.Description("X轴标签"),
		),
		mcp.WithString("y_label",
			mcp.Description("Y轴标签"),
		),
		mcp.WithString("color",
			mcp.Description("点颜色"),
		),
	)

	srv.AddTool(plotLineTool, handlePlotLineTool)
	srv.AddTool(plotBarTool, handlePlotBarTool)
	srv.AddTool(plotScatterTool, handlePlotScatterTool)

	return srv
}

// handlePlotLineTool 处理折线图工具调用
func handlePlotLineTool(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	// 将请求参数转换为 JSON 再解析
	jsonBytes, err := json.Marshal(request.Params)
	if err != nil {
		return mcp.NewToolResultError("参数解析失败"), err
	}

	var args struct {
		Data   interface{} `json:"data"`
		Title  string      `json:"title"`
		Width  int         `json:"width"`
		Height int         `json:"height"`
		Labels string      `json:"labels"`
		Color  string      `json:"color"`
	}

	if err := json.Unmarshal(jsonBytes, &args); err != nil {
		return mcp.NewToolResultError("参数解析失败"), err
	}

	// 处理数据
	var data []float64
	switch v := args.Data.(type) {
	case float64:
		data = []float64{v}
	case []interface{}:
		data = make([]float64, 0, len(v))
		for _, item := range v {
			if f, ok := toFloat64(item); ok {
				data = append(data, f)
			}
		}
	}

	if len(data) == 0 {
		return mcp.NewToolResultError("错误: data 参数无效"), nil
	}

	// 解析标签
	var labels []string
	if args.Labels != "" {
		labels = strings.Split(args.Labels, ",")
	}

	// 获取颜色
	color := getColorStyle(args.Color)

	// 生成图表
	graph := charts.NewLineChart(args.Width, args.Height)
	graph.SetData(data)
	graph.SetTitle(args.Title)
	graph.SetLabels(labels)
	graph.SetColor(color)

	return mcp.NewToolResultText(graph.Render()), nil
}

// handlePlotBarTool 处理柱状图工具调用
func handlePlotBarTool(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	jsonBytes, err := json.Marshal(request.Params)
	if err != nil {
		return mcp.NewToolResultError("参数解析失败"), err
	}

	var args struct {
		Data       interface{} `json:"data"`
		Title      string      `json:"title"`
		Width      int         `json:"width"`
		Height     int         `json:"height"`
		Horizontal bool        `json:"horizontal"`
		Color      string      `json:"color"`
	}

	if err := json.Unmarshal(jsonBytes, &args); err != nil {
		return mcp.NewToolResultError("参数解析失败"), err
	}

	dataStr, ok := args.Data.(string)
	if !ok {
		return mcp.NewToolResultError("错误: data 参数格式无效"), nil
	}

	// 解析数据: "标签1:数值1,标签2:数值2"
	data := parseBarData(dataStr)
	if len(data) == 0 {
		return mcp.NewToolResultError("错误: data 参数格式无效"), nil
	}

	color := getColorStyle(args.Color)

	graph := charts.NewBarChart(args.Width, args.Height, args.Horizontal)
	graph.SetData(data)
	graph.SetTitle(args.Title)
	graph.SetColor(color)

	return mcp.NewToolResultText(graph.Render()), nil
}

// handlePlotScatterTool 处理散点图工具调用
func handlePlotScatterTool(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	jsonBytes, err := json.Marshal(request.Params)
	if err != nil {
		return mcp.NewToolResultError("参数解析失败"), err
	}

	var args struct {
		Data   interface{} `json:"data"`
		Title  string      `json:"title"`
		Width  int         `json:"width"`
		Height int         `json:"height"`
		XLabel string      `json:"x_label"`
		YLabel string      `json:"y_label"`
		Color  string      `json:"color"`
	}

	if err := json.Unmarshal(jsonBytes, &args); err != nil {
		return mcp.NewToolResultError("参数解析失败"), err
	}

	dataStr, ok := args.Data.(string)
	if !ok {
		return mcp.NewToolResultError("错误: data 参数格式无效"), nil
	}

	// 解析数据: "x1,y1 x2,y2 ..."
	data := parseScatterData(dataStr)
	if len(data) == 0 {
		return mcp.NewToolResultError("错误: data 参数格式无效"), nil
	}

	color := getColorStyle(args.Color)

	graph := charts.NewScatterChart(args.Width, args.Height)
	graph.SetData(data)
	graph.SetTitle(args.Title)
	graph.SetXLabel(args.XLabel)
	graph.SetYLabel(args.YLabel)
	graph.SetColor(color)

	return mcp.NewToolResultText(graph.Render()), nil
}

// parseBarData 解析柱状图数据
// 格式: "标签1:数值1,标签2:数值2"
func parseBarData(str string) []charts.BarDataItem {
	var result []charts.BarDataItem
	items := strings.Split(str, ",")
	for _, item := range items {
		parts := strings.SplitN(item, ":", 2)
		if len(parts) == 2 {
			if value, err := strconv.ParseFloat(strings.TrimSpace(parts[1]), 64); err == nil {
				result = append(result, charts.BarDataItem{
					Label: strings.TrimSpace(parts[0]),
					Value: value,
				})
			}
		}
	}
	return result
}

// parseScatterData 解析散点图数据
// 格式: "x1,y1 x2,y2 ..."
func parseScatterData(str string) []charts.ScatterDataItem {
	var result []charts.ScatterDataItem
	items := strings.Fields(str)
	for _, item := range items {
		parts := strings.SplitN(item, ",", 2)
		if len(parts) == 2 {
			x, err1 := strconv.ParseFloat(strings.TrimSpace(parts[0]), 64)
			y, err2 := strconv.ParseFloat(strings.TrimSpace(parts[1]), 64)
			if err1 == nil && err2 == nil {
				result = append(result, charts.ScatterDataItem{
					X: x,
					Y: y,
				})
			}
		}
	}
	return result
}

// getColorStyle 获取颜色样式
func getColorStyle(colorName string) lipgloss.Style {
	switch colorName {
	case "red":
		return lipgloss.NewStyle().Foreground(lipgloss.Color("9"))
	case "green":
		return lipgloss.NewStyle().Foreground(lipgloss.Color("10"))
	case "blue":
		return lipgloss.NewStyle().Foreground(lipgloss.Color("12"))
	case "yellow":
		return lipgloss.NewStyle().Foreground(lipgloss.Color("11"))
	case "cyan":
		return lipgloss.NewStyle().Foreground(lipgloss.Color("14"))
	case "magenta":
		return lipgloss.NewStyle().Foreground(lipgloss.Color("13"))
	default:
		return lipgloss.NewStyle().Foreground(lipgloss.Color("15"))
	}
}

// toFloat64 将 interface{} 转换为 float64
func toFloat64(v interface{}) (float64, bool) {
	switch n := v.(type) {
	case float64:
		return n, true
	case int:
		return float64(n), true
	case int64:
		return float64(n), true
	default:
		return 0, false
	}
}

// 这里添加 fmt 的使用以避免导入错误
var _ = fmt.Sprintf