// -*- coding: utf-8 -*-
// benchmark-server - configurable MCP server for benchmarking and load testing
//
// Copyright 2025
// SPDX-License-Identifier: Apache-2.0
//
// This file implements an MCP (Model Context Protocol) server written in Go
// that can expose an arbitrary number of tools, resources, and prompts for
// benchmarking purposes.
//
// Usage Examples:
//
//   # 1) STDIO transport with default 100 items
//   ./benchmark-server
//
//   # 2) Generate 1000 tools, 500 resources, 200 prompts
//   ./benchmark-server -tools=1000 -resources=500 -prompts=200
//
//   # 3) Large scale test with 10000 items each
//   ./benchmark-server -tools=10000 -resources=10000 -prompts=10000
//
//   # 4) SSE transport for web testing
//   ./benchmark-server -transport=sse -port=8080 -tools=500
//
//   # 5) HTTP transport with authentication
//   ./benchmark-server -transport=http -port=9090 -auth-token=secret123
//
//   # 6) Custom payload sizes per type
//   ./benchmark-server -tools=100 -tool-size=5000 -resource-size=10000 -prompt-size=2000
//
//   # 7) Large tools with small resources
//   ./benchmark-server -tools=1000 -tool-size=50000 -resources=100 -resource-size=500
//
//   # 8) Minimal setup for quick tests
//   ./benchmark-server -tools=10 -resources=5 -prompts=5 -log-level=debug
//
// Build:
//   go build -o benchmark-server .
//
// Claude Desktop Configuration (stdio):
//   {
//     "mcpServers": {
//       "benchmark": {
//         "command": "/path/to/benchmark-server",
//         "args": ["-tools=1000", "-resources=500", "-prompts=200"]
//       }
//     }
//   }

package main

import (
    "context"
    "encoding/json"
    "flag"
    "fmt"
    "io"
    "log"
    "net/http"
    "os"
    "os/signal"
    "strings"
    "sync"
    "syscall"
    "time"

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

/* ------------------------------------------------------------------ */
/*                             constants                              */
/* ------------------------------------------------------------------ */

const (
    appName    = "benchmark-server"
    appVersion = "1.0.0"

    // Default values
    defaultPort         = 8080
    defaultListen       = "0.0.0.0"
    defaultLogLevel     = "info"
    defaultServerCount  = 1
    defaultToolCount    = 100
    defaultResourceCnt  = 100
    defaultPromptCount  = 100
    defaultToolSize     = 1000
    defaultResourceSize = 1000
    defaultPromptSize   = 1000

    // Server timeouts
    shutdownTimeout = 10 * time.Second

    // Environment variables
    envAuthToken = "AUTH_TOKEN"
)

/* ------------------------------------------------------------------ */
/*                             logging                                */
/* ------------------------------------------------------------------ */

// logLvl represents logging verbosity levels
type logLvl int

const (
    logNone logLvl = iota
    logError
    logWarn
    logInfo
    logDebug
)

var (
    curLvl = logInfo
    logger = log.New(os.Stderr, "", log.LstdFlags)
)

// parseLvl converts a string log level to logLvl type
func parseLvl(s string) logLvl {
    switch strings.ToLower(s) {
    case "debug":
        return logDebug
    case "info":
        return logInfo
    case "warn", "warning":
        return logWarn
    case "error":
        return logError
    case "none", "off", "silent":
        return logNone
    default:
        return logInfo
    }
}

// logAt logs a message if the current log level permits
func logAt(l logLvl, f string, v ...any) {
    if curLvl >= l {
        logger.Printf(f, v...)
    }
}

/* ------------------------------------------------------------------ */
/*                    version / health helpers                        */
/* ------------------------------------------------------------------ */

var startTime = time.Now()

// versionJSON returns server version information as JSON
func versionJSON() string {
    return fmt.Sprintf(`{"name":%q,"version":%q,"mcp_version":"1.0"}`, appName, appVersion)
}

// healthJSON returns server health status as JSON
func healthJSON() string {
    return fmt.Sprintf(`{"status":"healthy","uptime_seconds":%d}`, int(time.Since(startTime).Seconds()))
}

/* ------------------------------------------------------------------ */
/*                      payload generation                            */
/* ------------------------------------------------------------------ */

// generatePayload creates a text payload of specified size
func generatePayload(name string, size int) string {
    base := fmt.Sprintf("Response from %s. ", name)
    if size <= len(base) {
        return base[:size]
    }

    // Fill the rest with repeating text
    filler := "This is benchmark data. "
    needed := size - len(base)
    count := (needed / len(filler)) + 1

    result := base
    for i := 0; i < count; i++ {
        result += filler
    }

    return result[:size]
}

/* ------------------------------------------------------------------ */
/*                     dynamic handler creators                       */
/* ------------------------------------------------------------------ */

// createToolHandler creates a handler function for a tool
func createToolHandler(toolName string, payloadSize int) func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) {
    return func(_ context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
        payload := generatePayload(toolName, payloadSize)

        response := map[string]interface{}{
            "tool":      toolName,
            "timestamp": time.Now().UTC().Format(time.RFC3339),
            "arguments": req.Params.Arguments,
            "data":      payload,
        }

        jsonData, err := json.Marshal(response)
        if err != nil {
            return mcp.NewToolResultError(fmt.Sprintf("failed to marshal response: %v", err)), nil
        }

        logAt(logDebug, "tool invoked: %s", toolName)
        return mcp.NewToolResultText(string(jsonData)), nil
    }
}

// createResourceHandler creates a handler function for a resource
func createResourceHandler(resourceName string, payloadSize int) func(context.Context, mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
    return func(_ context.Context, _ mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
        payload := generatePayload(resourceName, payloadSize)

        data := map[string]interface{}{
            "resource":  resourceName,
            "timestamp": time.Now().UTC().Format(time.RFC3339),
            "data":      payload,
        }

        jsonData, err := json.Marshal(data)
        if err != nil {
            return nil, fmt.Errorf("failed to marshal resource data: %w", err)
        }

        logAt(logDebug, "resource requested: %s", resourceName)
        return []mcp.ResourceContents{
            mcp.TextResourceContents{
                URI:      fmt.Sprintf("benchmark://%s", resourceName),
                MIMEType: "application/json",
                Text:     string(jsonData),
            },
        }, nil
    }
}

// createPromptHandler creates a handler function for a prompt
func createPromptHandler(promptName string, payloadSize int) func(context.Context, mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
    return func(_ context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
        payload := generatePayload(promptName, payloadSize)

        var promptText strings.Builder
        promptText.WriteString(fmt.Sprintf("Prompt: %s\n\n", promptName))
        promptText.WriteString(fmt.Sprintf("Timestamp: %s\n\n", time.Now().UTC().Format(time.RFC3339)))

        // Include any arguments
        if len(req.Params.Arguments) > 0 {
            promptText.WriteString("Arguments:\n")
            for k, v := range req.Params.Arguments {
                promptText.WriteString(fmt.Sprintf("  - %s: %s\n", k, v))
            }
            promptText.WriteString("\n")
        }

        promptText.WriteString(payload)

        logAt(logDebug, "prompt requested: %s", promptName)
        return &mcp.GetPromptResult{
            Description: fmt.Sprintf("Benchmark prompt %s", promptName),
            Messages: []mcp.PromptMessage{
                {
                    Role:    mcp.RoleUser,
                    Content: mcp.TextContent{Type: "text", Text: promptText.String()},
                },
            },
        }, nil
    }
}

/* ------------------------------------------------------------------ */
/*                       authentication middleware                    */
/* ------------------------------------------------------------------ */

// authMiddleware creates a middleware that checks for Bearer token authentication
func authMiddleware(token string, next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        // Skip auth for health and version endpoints
        if r.URL.Path == "/health" || r.URL.Path == "/version" {
            next.ServeHTTP(w, r)
            return
        }

        // Get Authorization header
        authHeader := r.Header.Get("Authorization")
        if authHeader == "" {
            logAt(logWarn, "missing authorization header from %s", r.RemoteAddr)
            w.Header().Set("WWW-Authenticate", `Bearer realm="MCP Server"`)
            http.Error(w, "Authorization required", http.StatusUnauthorized)
            return
        }

        // Check Bearer token format
        const bearerPrefix = "Bearer "
        if !strings.HasPrefix(authHeader, bearerPrefix) {
            logAt(logWarn, "invalid authorization format from %s", r.RemoteAddr)
            http.Error(w, "Invalid authorization format", http.StatusUnauthorized)
            return
        }

        // Verify token
        providedToken := strings.TrimPrefix(authHeader, bearerPrefix)
        if providedToken != token {
            logAt(logWarn, "invalid token from %s", r.RemoteAddr)
            http.Error(w, "Invalid token", http.StatusUnauthorized)
            return
        }

        // Token valid, proceed with request
        logAt(logDebug, "authenticated request from %s", r.RemoteAddr)
        next.ServeHTTP(w, r)
    })
}

/* ------------------------------------------------------------------ */
/*                      server creation helpers                       */
/* ------------------------------------------------------------------ */

// serverConfig holds configuration for creating an MCP server instance
type serverConfig struct {
    toolCount    int
    toolSize     int
    resourceCnt  int
    resourceSize int
    promptCount  int
    promptSize   int
}

// buildMCPServer creates and registers tools/resources/prompts for a server
func buildMCPServer(cfg serverConfig, serverID int, isMultiServer bool) *server.MCPServer {
    name := appName
    if isMultiServer {
        name = fmt.Sprintf("%s-%d", appName, serverID)
    }

    s := server.NewMCPServer(
        name,
        appVersion,
        server.WithToolCapabilities(false), // No progress reporting needed
        server.WithResourceCapabilities(false, true), // Enable resource capabilities
        server.WithPromptCapabilities(true),          // Enable prompt capabilities
        server.WithLogging(),                         // Enable MCP protocol logging
        server.WithRecovery(),                        // Recover from panics in handlers
    )

    // Register tools
    logAt(logInfo, "[%s] registering %d tools...", name, cfg.toolCount)
    for i := 0; i < cfg.toolCount; i++ {
        toolName := fmt.Sprintf("benchmark_tool_%d", i)
        tool := mcp.NewTool(toolName,
            mcp.WithDescription(fmt.Sprintf("Benchmark tool #%d - returns test data", i)),
            mcp.WithTitleAnnotation(fmt.Sprintf("Benchmark Tool %d", i)),
            mcp.WithReadOnlyHintAnnotation(true),
            mcp.WithString("param1",
                mcp.Description("Optional parameter 1"),
            ),
            mcp.WithString("param2",
                mcp.Description("Optional parameter 2"),
            ),
        )
        s.AddTool(tool, createToolHandler(toolName, cfg.toolSize))
    }
    logAt(logInfo, "[%s] registered %d tools", name, cfg.toolCount)

    // Register resources
    logAt(logInfo, "[%s] registering %d resources...", name, cfg.resourceCnt)
    for i := 0; i < cfg.resourceCnt; i++ {
        var resourceURI, resourceName string
        if isMultiServer {
            // Multi-server mode: include server ID in URI for uniqueness
            resourceName = fmt.Sprintf("resource_%d", i)
            resourceURI = fmt.Sprintf("benchmark://server%d/resource_%d", serverID, i)
        } else {
            // Single server mode: original naming
            resourceName = fmt.Sprintf("benchmark_resource_%d", i)
            resourceURI = fmt.Sprintf("benchmark://%s", resourceName)
        }

        resource := mcp.NewResource(
            resourceURI,
            fmt.Sprintf("Benchmark Resource %d", i),
            mcp.WithResourceDescription(fmt.Sprintf("Benchmark resource #%d - returns test data", i)),
            mcp.WithMIMEType("application/json"),
        )
        s.AddResource(resource, createResourceHandler(resourceName, cfg.resourceSize))
    }
    logAt(logInfo, "[%s] registered %d resources", name, cfg.resourceCnt)

    // Register prompts
    logAt(logInfo, "[%s] registering %d prompts...", name, cfg.promptCount)
    for i := 0; i < cfg.promptCount; i++ {
        promptName := fmt.Sprintf("benchmark_prompt_%d", i)
        prompt := mcp.NewPrompt(promptName,
            mcp.WithPromptDescription(fmt.Sprintf("Benchmark prompt #%d - returns test prompt", i)),
            mcp.WithArgument("arg1",
                mcp.ArgumentDescription("Optional argument 1"),
            ),
            mcp.WithArgument("arg2",
                mcp.ArgumentDescription("Optional argument 2"),
            ),
        )
        s.AddPrompt(prompt, createPromptHandler(promptName, cfg.promptSize))
    }
    logAt(logInfo, "[%s] registered %d prompts", name, cfg.promptCount)

    return s
}

/* ------------------------------------------------------------------ */
/*                     HTTP server runner with shutdown               */
/* ------------------------------------------------------------------ */

// runHTTPServer starts an HTTP server and handles graceful shutdown
func runHTTPServer(ctx context.Context, wg *sync.WaitGroup, addr string, handler http.Handler, serverName string) {
    defer wg.Done()

    srv := &http.Server{
        Addr:              addr,
        Handler:           handler,
        ReadHeaderTimeout: shutdownTimeout,
    }

    // Start server in a goroutine
    go func() {
        logAt(logInfo, "[%s] listening on http://%s", serverName, addr)
        if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
            logAt(logError, "[%s] server error: %v", serverName, err)
        }
    }()

    // Wait for context cancellation (shutdown signal)
    <-ctx.Done()

    // Create shutdown context with timeout
    shutdownCtx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
    defer cancel()

    logAt(logInfo, "[%s] shutting down gracefully...", serverName)
    if err := srv.Shutdown(shutdownCtx); err != nil {
        logAt(logError, "[%s] shutdown error: %v", serverName, err)
    } else {
        logAt(logInfo, "[%s] shutdown complete", serverName)
    }
}

/* ------------------------------------------------------------------ */
/*                              main                                  */
/* ------------------------------------------------------------------ */

func main() {
    /* ---------------------------- flags --------------------------- */
    var (
        transport    = flag.String("transport", "stdio", "Transport: stdio | sse | http")
        addrFlag     = flag.String("addr", "", "Full listen address (host:port) - overrides -listen/-port")
        listenHost   = flag.String("listen", defaultListen, "Listen interface for sse/http")
        port         = flag.Int("port", defaultPort, "TCP port for sse/http")
        serverCount  = flag.Int("server-count", defaultServerCount, "Number of servers to spawn (for multi-server testing)")
        startPort    = flag.Int("start-port", defaultPort, "Starting port for multi-server mode")
        publicURL    = flag.String("public-url", "", "External base URL advertised to SSE clients")
        authToken    = flag.String("auth-token", "", "Bearer token for authentication (SSE/HTTP only)")
        logLevel     = flag.String("log-level", defaultLogLevel, "Logging level: debug|info|warn|error|none")
        toolCount    = flag.Int("tools", defaultToolCount, "Number of tools to generate")
        resourceCnt  = flag.Int("resources", defaultResourceCnt, "Number of resources to generate")
        promptCount  = flag.Int("prompts", defaultPromptCount, "Number of prompts to generate")
        toolSize     = flag.Int("tool-size", defaultToolSize, "Size of tool response payload in bytes")
        resourceSize = flag.Int("resource-size", defaultResourceSize, "Size of resource response payload in bytes")
        promptSize   = flag.Int("prompt-size", defaultPromptSize, "Size of prompt response payload in bytes")
        showHelp     = flag.Bool("help", false, "Show help message")
    )

    // Custom usage function
    flag.Usage = func() {
        const ind = "  "
        fmt.Fprintf(flag.CommandLine.Output(),
            "%s %s - configurable MCP server for benchmarking\n\n",
            appName, appVersion)
        fmt.Fprintln(flag.CommandLine.Output(), "Options:")
        flag.VisitAll(func(fl *flag.Flag) {
            fmt.Fprintf(flag.CommandLine.Output(), ind+"-%s\n", fl.Name)
            fmt.Fprintf(flag.CommandLine.Output(), ind+ind+"%s (default %q)\n\n",
                fl.Usage, fl.DefValue)
        })
        fmt.Fprint(flag.CommandLine.Output(),
            "Examples:\n",
            ind, appName, " -tools=1000 -resources=500 -prompts=200\n",
            ind, appName, " -transport=sse -port=8080 -tools=500\n",
            ind, appName, " -tools=100 -tool-size=5000 -resource-size=10000 -prompt-size=2000\n",
            ind, appName, " -tools=10000 -resources=0 -prompts=0 -tool-size=500\n",
            ind, appName, " -transport=sse -server-count=10 -start-port=8080 -tools=100\n",
            ind, appName, " -transport=http -server-count=5 -start-port=9000\n\n",
            "Environment Variables:\n",
            ind, "AUTH_TOKEN - Bearer token for authentication (overrides -auth-token flag)\n",
        )
    }

    flag.Parse()

    if *showHelp {
        flag.Usage()
        os.Exit(0)
    }

    /* ----------------------- configuration setup ------------------ */
    // Check for auth token in environment variable (overrides flag)
    if envToken := os.Getenv(envAuthToken); envToken != "" {
        *authToken = envToken
        logAt(logDebug, "using auth token from environment variable")
    }

    /* ------------------------- logging setup ---------------------- */
    curLvl = parseLvl(*logLevel)
    if curLvl == logNone {
        logger.SetOutput(io.Discard)
    }

    logAt(logInfo, "starting %s %s", appName, appVersion)
    logAt(logInfo, "configuration: tools=%d (size=%d) resources=%d (size=%d) prompts=%d (size=%d)",
        *toolCount, *toolSize, *resourceCnt, *resourceSize, *promptCount, *promptSize)

    if *authToken != "" && *transport != "stdio" {
        logAt(logInfo, "authentication enabled with Bearer token")
    }

    // Validate server count for multi-server mode
    if *serverCount < 1 {
        logger.Fatalf("server-count must be at least 1, got %d", *serverCount)
    }

    // stdio transport doesn't support multi-server
    if strings.ToLower(*transport) == "stdio" && *serverCount > 1 {
        logAt(logWarn, "stdio transport does not support multi-server mode, ignoring server-count")
        *serverCount = 1
    }

    if *serverCount > 1 {
        logAt(logInfo, "multi-server mode: spawning %d servers starting at port %d", *serverCount, *startPort)
    }

    /* ----------------------- build server config ------------------ */
    cfg := serverConfig{
        toolCount:    *toolCount,
        toolSize:     *toolSize,
        resourceCnt:  *resourceCnt,
        resourceSize: *resourceSize,
        promptCount:  *promptCount,
        promptSize:   *promptSize,
    }

    /* -------------------- choose transport & serve ---------------- */
    switch strings.ToLower(*transport) {

    /* ---------------------------- stdio -------------------------- */
    case "stdio":
        if *authToken != "" {
            logAt(logWarn, "auth-token is ignored for stdio transport")
        }
        logAt(logInfo, "serving via stdio transport")
        s := buildMCPServer(cfg, 0, false)
        if err := server.ServeStdio(s); err != nil {
            logger.Fatalf("stdio server error: %v", err)
        }

    /* ----------------------------- sse --------------------------- */
    case "sse":
        // Setup signal handling for graceful shutdown
        ctx, cancel := context.WithCancel(context.Background())
        sigChan := make(chan os.Signal, 1)
        signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)

        var wg sync.WaitGroup

        // Spawn servers
        for i := 0; i < *serverCount; i++ {
            serverID := i
            currentPort := *startPort + i

            // Override addr for multi-server mode
            var addr string
            if *serverCount > 1 {
                addr = fmt.Sprintf("%s:%d", *listenHost, currentPort)
            } else {
                addr = effectiveAddr(*addrFlag, *listenHost, *port)
            }

            // Build MCP server instance
            s := buildMCPServer(cfg, serverID, *serverCount > 1)

            // Create mux and register handlers
            mux := http.NewServeMux()

            // Configure SSE options
            opts := []server.SSEOption{}
            if *publicURL != "" {
                // Append port to public URL for multi-server
                if *serverCount > 1 {
                    opts = append(opts, server.WithBaseURL(fmt.Sprintf("%s:%d", strings.TrimRight(*publicURL, "/"), currentPort)))
                } else {
                    opts = append(opts, server.WithBaseURL(strings.TrimRight(*publicURL, "/")))
                }
            }

            // Register SSE handler
            sseHandler := server.NewSSEServer(s, opts...)
            mux.Handle("/", sseHandler)

            // Register health and version endpoints
            registerHealthAndVersion(mux)

            // Create handler chain
            var handler http.Handler = mux
            handler = loggingHTTPMiddleware(handler)
            if *authToken != "" {
                handler = authMiddleware(*authToken, handler)
            }

            // Log server info
            serverName := fmt.Sprintf("%s-%d", appName, serverID)
            if *serverCount == 1 {
                serverName = appName
                logAt(logInfo, "SSE server ready on http://%s", addr)
                logAt(logInfo, "  MCP SSE events:   /sse")
                logAt(logInfo, "  MCP SSE messages: /messages")
                logAt(logInfo, "  Health check:     /health")
                logAt(logInfo, "  Version info:     /version")
                if *publicURL != "" {
                    logAt(logInfo, "  Public URL:       %s", *publicURL)
                }
                if *authToken != "" {
                    logAt(logInfo, "  Authentication:   Bearer token required")
                }
            }

            // Start server
            wg.Add(1)
            go runHTTPServer(ctx, &wg, addr, handler, serverName)
        }

        // Wait for interrupt signal
        <-sigChan
        logAt(logInfo, "received shutdown signal, shutting down %d server(s)...", *serverCount)

        // Cancel context to trigger graceful shutdown
        cancel()

        // Wait for all servers to finish shutting down
        wg.Wait()
        logAt(logInfo, "all servers stopped")

    /* ----------------------- streamable http --------------------- */
    case "http":
        // Setup signal handling for graceful shutdown
        ctx, cancel := context.WithCancel(context.Background())
        sigChan := make(chan os.Signal, 1)
        signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)

        var wg sync.WaitGroup

        // Spawn servers
        for i := 0; i < *serverCount; i++ {
            serverID := i
            currentPort := *startPort + i

            // Override addr for multi-server mode
            var addr string
            if *serverCount > 1 {
                addr = fmt.Sprintf("%s:%d", *listenHost, currentPort)
            } else {
                addr = effectiveAddr(*addrFlag, *listenHost, *port)
            }

            // Build MCP server instance
            s := buildMCPServer(cfg, serverID, *serverCount > 1)

            // Create mux and register handlers
            mux := http.NewServeMux()

            // Register HTTP handler
            httpHandler := server.NewStreamableHTTPServer(s)
            mux.Handle("/", httpHandler)

            // Register health and version endpoints
            registerHealthAndVersion(mux)

            // Create handler chain
            var handler http.Handler = mux
            handler = loggingHTTPMiddleware(handler)
            if *authToken != "" {
                handler = authMiddleware(*authToken, handler)
            }

            // Log server info
            serverName := fmt.Sprintf("%s-%d", appName, serverID)
            if *serverCount == 1 {
                serverName = appName
                logAt(logInfo, "HTTP server ready on http://%s", addr)
                logAt(logInfo, "  MCP endpoint:     / (POST with JSON-RPC)")
                logAt(logInfo, "  Health check:     /health")
                logAt(logInfo, "  Version info:     /version")
                if *authToken != "" {
                    logAt(logInfo, "  Authentication:   Bearer token required")
                }
            }

            // Start server
            wg.Add(1)
            go runHTTPServer(ctx, &wg, addr, handler, serverName)
        }

        // Wait for interrupt signal
        <-sigChan
        logAt(logInfo, "received shutdown signal, shutting down %d server(s)...", *serverCount)

        // Cancel context to trigger graceful shutdown
        cancel()

        // Wait for all servers to finish shutting down
        wg.Wait()
        logAt(logInfo, "all servers stopped")

    default:
        fmt.Fprintf(os.Stderr, "Error: unknown transport %q\n\n", *transport)
        flag.Usage()
        os.Exit(2)
    }
}

/* ------------------------------------------------------------------ */
/*                        helper functions                            */
/* ------------------------------------------------------------------ */

// effectiveAddr determines the actual address to listen on
func effectiveAddr(addrFlag, listen string, port int) string {
    if addrFlag != "" {
        return addrFlag
    }
    return fmt.Sprintf("%s:%d", listen, port)
}

// registerHealthAndVersion adds health and version endpoints to the mux
func registerHealthAndVersion(mux *http.ServeMux) {
    // Health endpoint - JSON response
    mux.HandleFunc("/health", func(w http.ResponseWriter, _ *http.Request) {
        w.Header().Set("Content-Type", "application/json")
        w.WriteHeader(http.StatusOK)
        _, _ = w.Write([]byte(healthJSON()))
    })

    // Version endpoint - JSON response
    mux.HandleFunc("/version", func(w http.ResponseWriter, _ *http.Request) {
        w.Header().Set("Content-Type", "application/json")
        w.WriteHeader(http.StatusOK)
        _, _ = w.Write([]byte(versionJSON()))
    })
}

/* -------------------- HTTP middleware ----------------------------- */

// loggingHTTPMiddleware provides request logging when log level permits
func loggingHTTPMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        if curLvl < logInfo {
            next.ServeHTTP(w, r)
            return
        }

        start := time.Now()

        // Wrap response writer to capture status code
        rw := &statusWriter{ResponseWriter: w, status: http.StatusOK, written: false}

        // Call the next handler
        next.ServeHTTP(rw, r)

        // Log the request
        duration := time.Since(start)
        if r.Method == "POST" && curLvl >= logDebug {
            logAt(logDebug, "%s %s %s %d (Content-Length: %s) %v",
                r.RemoteAddr, r.Method, r.URL.Path, rw.status, r.Header.Get("Content-Length"), duration)
        } else {
            logAt(logInfo, "%s %s %s %d %v",
                r.RemoteAddr, r.Method, r.URL.Path, rw.status, duration)
        }
    })
}

// statusWriter wraps http.ResponseWriter to capture status code
type statusWriter struct {
    http.ResponseWriter
    status  int
    written bool
}

func (sw *statusWriter) WriteHeader(code int) {
    if !sw.written {
        sw.status = code
        sw.written = true
        sw.ResponseWriter.WriteHeader(code)
    }
}

func (sw *statusWriter) Write(b []byte) (int, error) {
    if !sw.written {
        sw.WriteHeader(http.StatusOK)
    }
    return sw.ResponseWriter.Write(b)
}

// Flush lets the underlying handler stream (needed for SSE)
func (sw *statusWriter) Flush() {
    if f, ok := sw.ResponseWriter.(http.Flusher); ok {
        if !sw.written {
            sw.WriteHeader(http.StatusOK)
        }
        f.Flush()
    }
}
