// Create a basic server
s := server.NewMCPServer(
    "My Server",  // Server name
    "1.0.0",     // Version
)

// Start the server using stdio
if err := server.ServeStdio(s); err != nil {
    log.Fatalf("Server error: %v", err)
}

// Static resource example - exposing a README file
resource := mcp.NewResource(
    "docs://readme",
    "Project README",
    mcp.WithResourceDescription("The project's README file"), 
    mcp.WithMIMEType("text/markdown"),
)

// Add resource with its handler
s.AddResource(resource, func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
    content, err := os.ReadFile("README.md")
    if err != nil {
        return nil, err
    }
    
    return []mcp.ResourceContents{
        mcp.TextResourceContents{
            URI:      "docs://readme",
            MIMEType: "text/markdown",
            Text:     string(content),
        },
    }, nil
})

// Dynamic resource example - user profiles by ID
template := mcp.NewResourceTemplate(
    "users://{id}/profile",
    "User Profile",
    mcp.WithTemplateDescription("Returns user profile information"),
    mcp.WithTemplateMIMEType("application/json"),
)

// Add template with its handler
s.AddResourceTemplate(template, func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
    // Extract ID from the URI using regex matching
    // The server automatically matches URIs to templates
    userID := extractIDFromURI(request.Params.URI)
    
    profile, err := getUserProfile(userID)  // Your DB/API call here
    if err != nil {
        return nil, err
    }
    
    return []mcp.ResourceContents{
        mcp.TextResourceContents{
            URI:      request.Params.URI,
            MIMEType: "application/json",
            Text:     profile,
        },
    }, nil
})

// Example: A tool that requires task execution
processBatchTool := mcp.NewTool("process_batch",
    mcp.WithDescription("Process a batch of items asynchronously"),
    mcp.WithTaskSupport(mcp.TaskSupportRequired),
    mcp.WithArray("items",
        mcp.Description("Array of items to process"),
        mcp.WithStringItems(),
        mcp.Required(),
    ),
)

// Task tool handler returns CreateTaskResult instead of CallToolResult
s.AddTaskTool(processBatchTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CreateTaskResult, error) {
    items := request.GetStringSlice("items", []string{})
    
    // Long-running work here
    for i, item := range items {
        select {
        case <-ctx.Done():
            // Task was cancelled
            return nil, ctx.Err()
        default:
            // Process item...
            processItem(item)
        }
    }
    
    // Return result - task ID and metadata are managed by the server
    return &mcp.CreateTaskResult{
        Task: mcp.Task{
            // Task fields (ID, status, etc.) are populated by the server
        },
    }, nil
})

// Enable task capabilities when creating the server
s := server.NewMCPServer(
    "Task Server",
    "1.0.0",
    server.WithTaskCapabilities(
        true, // listTasks: allows clients to list all tasks
        true, // cancel: allows clients to cancel running tasks
        true, // toolCallTasks: enables task augmentation for tools
    ),
    server.WithMaxConcurrentTasks(10), // Optional: limit concurrent running tasks
)

// Tool with optional task support
analyzeTool := mcp.NewTool("analyze_data",
    mcp.WithDescription("Analyze data - can run sync or async"),
    mcp.WithTaskSupport(mcp.TaskSupportOptional),
    mcp.WithString("data", mcp.Required()),
)

// Use AddTaskTool for hybrid tools that support both modes
s.AddTaskTool(analyzeTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CreateTaskResult, error) {
    // This handler runs when called as a task
    data := request.GetString("data", "")
    result := analyzeData(data)
    
    return &mcp.CreateTaskResult{
        Task: mcp.Task{},
    }, nil
})

// The server automatically handles both sync and async invocations
// When called without task param: executes handler and returns immediately
// When called with task param: executes handler asynchronously

s := server.NewMCPServer(
    "Task Server",
    "1.0.0",
    server.WithTaskCapabilities(true, true, true),
    server.WithMaxConcurrentTasks(10), // Allow up to 10 concurrent running tasks
)

calculatorTool := mcp.NewTool("calculate",
    mcp.WithDescription("Perform basic arithmetic calculations"),
    mcp.WithString("operation",
        mcp.Required(),
        mcp.Description("The arithmetic operation to perform"),
        mcp.Enum("add", "subtract", "multiply", "divide"),
    ),
    mcp.WithNumber("x",
        mcp.Required(),
        mcp.Description("First number"),
    ),
    mcp.WithNumber("y",
        mcp.Required(),
        mcp.Description("Second number"),
    ),
)

s.AddTool(calculatorTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
    args := request.GetArguments()
    op := args["operation"].(string)
    x := args["x"].(float64)
    y := args["y"].(float64)

    var result float64
    switch op {
    case "add":
        result = x + y
    case "subtract":
        result = x - y
    case "multiply":
        result = x * y
    case "divide":
        if y == 0 {
            return mcp.NewToolResultError("cannot divide by zero"), nil
        }
        result = x / y
    }
    
    return mcp.FormatNumberResult(result), nil
})

httpTool := mcp.NewTool("http_request",
    mcp.WithDescription("Make HTTP requests to external APIs"),
    mcp.WithString("method",
        mcp.Required(),
        mcp.Description("HTTP method to use"),
        mcp.Enum("GET", "POST", "PUT", "DELETE"),
    ),
    mcp.WithString("url",
        mcp.Required(),
        mcp.Description("URL to send the request to"),
        mcp.Pattern("^https?://.*"),
    ),
    mcp.WithString("body",
        mcp.Description("Request body (for POST/PUT)"),
    ),
)

s.AddTool(httpTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
    args := request.GetArguments()
    method := args["method"].(string)
    url := args["url"].(string)
    body := ""
    if b, ok := args["body"].(string); ok {
        body = b
    }

    // Create and send request
    var req *http.Request
    var err error
    if body != "" {
        req, err = http.NewRequest(method, url, strings.NewReader(body))
    } else {
        req, err = http.NewRequest(method, url, nil)
    }
    if err != nil {
        return mcp.NewToolResultErrorFromErr("unable to create request", err), nil
    }

    client := &http.Client{}
    resp, err := client.Do(req)
    if err != nil {
        return mcp.NewToolResultErrorFromErr("unable to execute request", err), nil
    }
    defer resp.Body.Close()

    // Return response
    respBody, err := io.ReadAll(resp.Body)
    if err != nil {
        return mcp.NewToolResultErrorFromErr("unable to read request response", err), nil
    }

    return mcp.NewToolResultText(fmt.Sprintf("Status: %d\nBody: %s", resp.StatusCode, string(respBody))), nil
})

// Simple greeting prompt
s.AddPrompt(mcp.NewPrompt("greeting",
    mcp.WithPromptDescription("A friendly greeting prompt"),
    mcp.WithArgument("name",
        mcp.ArgumentDescription("Name of the person to greet"),
    ),
), func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
    name := request.Params.Arguments["name"]
    if name == "" {
        name = "friend"
    }
    
    return mcp.NewGetPromptResult(
        "A friendly greeting",
        []mcp.PromptMessage{
            mcp.NewPromptMessage(
                mcp.RoleAssistant,
                mcp.NewTextContent(fmt.Sprintf("Hello, %s! How can I help you today?", name)),
            ),
        },
    ), nil
})

// Code review prompt with embedded resource
s.AddPrompt(mcp.NewPrompt("code_review",
    mcp.WithPromptDescription("Code review assistance"),
    mcp.WithArgument("pr_number",
        mcp.ArgumentDescription("Pull request number to review"),
        mcp.RequiredArgument(),
    ),
), func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
    prNumber := request.Params.Arguments["pr_number"]
    if prNumber == "" {
        return nil, fmt.Errorf("pr_number is required")
    }
    
    return mcp.NewGetPromptResult(
        "Code review assistance",
        []mcp.PromptMessage{
            mcp.NewPromptMessage(
                mcp.RoleUser,
                mcp.NewTextContent("Review the changes and provide constructive feedback."),
            ),
            mcp.NewPromptMessage(
                mcp.RoleAssistant,
                mcp.NewEmbeddedResource(mcp.ResourceContents{
                    URI: fmt.Sprintf("git://pulls/%s/diff", prNumber),
                    MIMEType: "text/x-diff",
                }),
            ),
        },
    ), nil
})

// Database query builder prompt
s.AddPrompt(mcp.NewPrompt("query_builder",
    mcp.WithPromptDescription("SQL query builder assistance"),
    mcp.WithArgument("table",
        mcp.ArgumentDescription("Name of the table to query"),
        mcp.RequiredArgument(),
    ),
), func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
    tableName := request.Params.Arguments["table"]
    if tableName == "" {
        return nil, fmt.Errorf("table name is required")
    }
    
    return mcp.NewGetPromptResult(
        "SQL query builder assistance",
        []mcp.PromptMessage{
            mcp.NewPromptMessage(
                mcp.RoleUser,
                mcp.NewTextContent("Help construct efficient and safe queries for the provided schema."),
            ),
            mcp.NewPromptMessage(
                mcp.RoleUser,
                mcp.NewEmbeddedResource(mcp.ResourceContents{
                    URI: fmt.Sprintf("db://schema/%s", tableName),
                    MIMEType: "application/json",
                }),
            ),
        },
    ), nil
})

// Create a server with session capabilities
s := server.NewMCPServer(
    "Session Demo",
    "1.0.0",
    server.WithToolCapabilities(true),
)

// Implement your own ClientSession
type MySession struct {
    id           string
    notifChannel chan mcp.JSONRPCNotification
    isInitialized bool
    // Add custom fields for your application
}

// Implement the ClientSession interface
func (s *MySession) SessionID() string {
    return s.id
}

func (s *MySession) NotificationChannel() chan<- mcp.JSONRPCNotification {
    return s.notifChannel
}

func (s *MySession) Initialize() {
    s.isInitialized = true
}

func (s *MySession) Initialized() bool {
    return s.isInitialized
}

// Register a session
session := &MySession{
    id:           "user-123",
    notifChannel: make(chan mcp.JSONRPCNotification, 10),
}
if err := s.RegisterSession(context.Background(), session); err != nil {
    log.Printf("Failed to register session: %v", err)
}

// Send notification to a specific client
err := s.SendNotificationToSpecificClient(
    session.SessionID(),
    "notification/update",
    map[string]any{"message": "New data available!"},
)
if err != nil {
    log.Printf("Failed to send notification: %v", err)
}

// Unregister session when done
s.UnregisterSession(context.Background(), session.SessionID())

// Implement SessionWithTools interface for per-session tools
type MyAdvancedSession struct {
    MySession  // Embed the basic session
    sessionTools map[string]server.ServerTool
}

// Implement additional methods for SessionWithTools
func (s *MyAdvancedSession) GetSessionTools() map[string]server.ServerTool {
    return s.sessionTools
}

func (s *MyAdvancedSession) SetSessionTools(tools map[string]server.ServerTool) {
    s.sessionTools = tools
}

// Create and register a session with tools support
advSession := &MyAdvancedSession{
    MySession: MySession{
        id:           "user-456",
        notifChannel: make(chan mcp.JSONRPCNotification, 10),
    },
    sessionTools: make(map[string]server.ServerTool),
}
if err := s.RegisterSession(context.Background(), advSession); err != nil {
    log.Printf("Failed to register session: %v", err)
}

// Add session-specific tools
userSpecificTool := mcp.NewTool(
    "user_data",
    mcp.WithDescription("Access user-specific data"),
)
// You can use AddSessionTool (similar to AddTool)
err := s.AddSessionTool(
    advSession.SessionID(),
    userSpecificTool,
    func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
        // This handler is only available to this specific session
        return mcp.NewToolResultText("User-specific data for " + advSession.SessionID()), nil
    },
)
if err != nil {
    log.Printf("Failed to add session tool: %v", err)
}

// Or use AddSessionTools directly with ServerTool
/*
err := s.AddSessionTools(
    advSession.SessionID(),
    server.ServerTool{
        Tool: userSpecificTool,
        Handler: func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
            // This handler is only available to this specific session
            return mcp.NewToolResultText("User-specific data for " + advSession.SessionID()), nil
        },
    },
)
if err != nil {
    log.Printf("Failed to add session tool: %v", err)
}
*/

// Delete session-specific tools when no longer needed
err = s.DeleteSessionTools(advSession.SessionID(), "user_data")
if err != nil {
    log.Printf("Failed to delete session tool: %v", err)
}

// Add a tool filter that only shows tools with certain prefixes
s := server.NewMCPServer(
    "Tool Filtering Demo",
    "1.0.0",
    server.WithToolCapabilities(true),
    server.WithToolFilter(func(ctx context.Context, tools []mcp.Tool) []mcp.Tool {
        // Get session from context
        session := server.ClientSessionFromContext(ctx)
        if session == nil {
            return tools // Return all tools if no session
        }
        
        // Example: filter tools based on session ID prefix
        if strings.HasPrefix(session.SessionID(), "admin-") {
            // Admin users get all tools
            return tools
        } else {
            // Regular users only get tools with "public-" prefix
            var filteredTools []mcp.Tool
            for _, tool := range tools {
                if strings.HasPrefix(tool.Name, "public-") {
                    filteredTools = append(filteredTools, tool)
                }
            }
            return filteredTools
        }
    }),
)

s.AddTool(mcp.NewTool("session_aware"), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
    // Get the current session from context
    session := server.ClientSessionFromContext(ctx)
    if session == nil {
        return mcp.NewToolResultError("No active session"), nil
    }
    
    return mcp.NewToolResultText("Hello, session " + session.SessionID()), nil
})

// When using handlers in HTTP/SSE servers, you need to pass the context with the session
httpHandler := func(w http.ResponseWriter, r *http.Request) {
    // Get session from somewhere (like a cookie or header)
    session := getSessionFromRequest(r)
    
    // Add session to context
    ctx := s.WithContext(r.Context(), session)
    
    // Use this context when handling requests
    // ...
}

type MyPromptCompletionProvider struct{}

func (p *MyPromptCompletionProvider) CompletePromptArgument(
    ctx context.Context,
    promptName string,
    argument mcp.CompleteArgument,
    context mcp.CompleteContext,
) (*mcp.Completion, error) {
    // Example: provide style suggestions for a "code_review" prompt
    if promptName == "code_review" && argument.Name == "style" {
        styles := []string{"formal", "casual", "technical", "creative"}
        var suggestions []string
        
        // Filter based on current input
        for _, style := range styles {
            if strings.HasPrefix(style, argument.Value) {
                suggestions = append(suggestions, style)
            }
        }
        
        return &mcp.Completion{
            Values: suggestions,
        }, nil
    }
    
    // Return empty suggestions for unhandled cases
    return &mcp.Completion{Values: []string{}}, nil
}

type MyResourceCompletionProvider struct{}

func (p *MyResourceCompletionProvider) CompleteResourceArgument(
    ctx context.Context,
    uri string,
    argument mcp.CompleteArgument,
    context mcp.CompleteContext,
) (*mcp.Completion, error) {
    // Example: provide file path completions
    if uri == "file:///{path}" && argument.Name == "path" {
        // You can access previously completed arguments from context.Arguments
        // context.Arguments is a map[string]string of already-resolved arguments
        
        paths := getMatchingPaths(argument.Value) // Your custom logic
        
        return &mcp.Completion{
            Values:  paths[:min(len(paths), 100)], // Max 100 items
            Total:   len(paths),                    // Total available matches
            HasMore: len(paths) > 100,              // More results available
        }, nil
    }
    
    return &mcp.Completion{Values: []string{}}, nil
}

// Register the provider
mcpServer := server.NewMCPServer(
    "my-server",
    "1.0.0",
    server.WithCompletions(),
    server.WithPromptCompletionProvider(&MyPromptCompletionProvider{}),
    server.WithResourceCompletionProvider(&MyResourceCompletionProvider{}),
)

func (p *MyProvider) CompleteResourceArgument(
    ctx context.Context,
    uri string,
    argument mcp.CompleteArgument,
    context mcp.CompleteContext,
) (*mcp.Completion, error) {
    // Access previously completed arguments
    if previousValue, ok := context.Arguments["previous_arg"]; ok {
        // Provide suggestions based on previous_arg value
        return getSuggestionsFor(argument.Value, previousValue), nil
    }
    
    return &mcp.Completion{Values: []string{}}, nil
}