package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"strconv"

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

func webSearchHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	query := request.Params.Arguments["query"].(string)
	client := &http.Client{}

	//https://www.baidu.com/s?wd=ssstn=json&rn=10
	u, err := url.Parse("https://www.baidu.com/s")
	v := url.Values{}
	v.Add("wd", query)
	v.Add("tn", "json")
	v.Add("rn", "10")
	u.RawQuery = v.Encode()

	fmt.Println(v.Encode())

	req, _ := http.NewRequest("GET", u.String(), nil)
	req.Header.Set("Accept", "text/html")
	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/134.0.0.0 Safari/537.36")

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	type Entry struct {
		Title string
		Url   string
		Abs   string
	}

	type Feed struct {
		Entry []Entry
	}
	type Result struct {
		Feed Feed
	}

	var r Result
	err = json.NewDecoder(resp.Body).Decode(&r)
	if err != nil {
		return nil, err
	}
	var messages []mcp.Content
	for _, entry := range r.Feed.Entry {
		if entry.Abs != "" {
			log.Println(entry.Title, entry.Url, entry.Abs)
			s := "标题:" + entry.Title + "\n链接:" + entry.Url + "\n简介:" + entry.Abs
			messages = append(messages, mcp.TextContent{Type: "text", Text: s})
		}

	}

	return &mcp.CallToolResult{
		Content: messages,
	}, nil
}

func NewMCPServer() *server.MCPServer {
	// Create a new MCP server
	s := server.NewMCPServer(
		"Calculator Demo",
		"1.0.0",
		server.WithResourceCapabilities(true, true),
		server.WithLogging(),
	)
	webSearchTool := mcp.NewTool("web_search",
		mcp.WithDescription("web search tool"),
		mcp.WithString("query", mcp.Description("The web search query"), mcp.Required()),
	)

	s.AddTool(webSearchTool, webSearchHandler)
	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" {
		/// Start the stdio server
		if err := server.ServeStdio(mcpServer); err != nil {
			fmt.Printf("Server error: %v\n", err)
		}
	} else {

		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)
		}
	}
}