package main  

import (  
	"context"  
	"flag"  
	"fmt"  
	"log"  
	"strconv"  

	"github.com/mark3labs/mcp-go/mcp"  
	"github.com/mark3labs/mcp-go/server"  
)  

// 两数求和 handler  
func addHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {  
	// 参数断言 float64 支持整数和小数  
	a, ok1 := request.Params.Arguments["a"].(float64)  
	b, ok2 := request.Params.Arguments["b"].(float64)  
	if !ok1 || !ok2 {  
		return nil, fmt.Errorf("参数类型错误，期望输入两个数字")  
	}  

	sum := a + b  
	resultStr := fmt.Sprintf("两数之和为: %.2f", sum)  

	return &mcp.CallToolResult{  
		Content: []mcp.Content{  
			mcp.TextContent{Type: "text", Text: resultStr},  
		},  
	}, nil  
}  

func NewMCPServer() *server.MCPServer {  
	// 创建 MCP server  
	s := server.NewMCPServer(  
		"Calculator Demo",  
		"1.0.0",  
		server.WithResourceCapabilities(true, true),  
		server.WithLogging(),  
	)  

	// 注册“求和”工具  
	addTool := mcp.NewTool("add",  
		mcp.WithDescription("两数求和工具"),  
		mcp.WithNumber("a", mcp.Description("第一个数字"), mcp.Required()),  
		mcp.WithNumber("b", mcp.Description("第二个数字"), mcp.Required()),  
	)  

	s.AddTool(addTool, addHandler)  
	return s  
}  

func main() {  
	var transport string  
	var port int  
	flag.StringVar(&transport, "t", "sse", "Transport type (stdio or sse)")  
	flag.StringVar(&transport, "transport", "sse", "Transport type (stdio or sse)")  
	flag.IntVar(&port, "p", 8080, "Port number (default 8080)")  
	flag.IntVar(&port, "port", 8080, "Port number (default 8080)")  
	flag.Parse()  

	mcpServer := NewMCPServer()  

	if transport == "stdio" {  
		// 启动 stdio server  
		if err := server.ServeStdio(mcpServer); err != nil {  
			fmt.Printf("Server error: %v\n", err)  
		}  
	} else {  
		// 启动 SSE server  
		sseServer := server.NewSSEServer(mcpServer)  
		log.Printf("SSE server listening on :" + strconv.Itoa(port))  
		if err := sseServer.Start(":" + strconv.Itoa(port)); err != nil {  
			log.Fatalf("Server error: %v", err)  
		}  
	}  
}