// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"slices"
	"strconv"
	"strings"
	"time"

	"github.com/go-chi/chi/v5"
	"github.com/go-chi/chi/v5/middleware"
	"github.com/go-chi/cors"
	"github.com/go-chi/httplog/v3"
	"github.com/go-chi/render"
	"github.com/googleapis/mcp-toolbox/internal/auth"
	"github.com/googleapis/mcp-toolbox/internal/embeddingmodels"
	"github.com/googleapis/mcp-toolbox/internal/log"
	"github.com/googleapis/mcp-toolbox/internal/prompts"
	"github.com/googleapis/mcp-toolbox/internal/server/mcp/jsonrpc"
	"github.com/googleapis/mcp-toolbox/internal/server/resources"
	"github.com/googleapis/mcp-toolbox/internal/sources"
	"github.com/googleapis/mcp-toolbox/internal/telemetry"
	"github.com/googleapis/mcp-toolbox/internal/tools"
	"github.com/googleapis/mcp-toolbox/internal/util"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"
)

// Server contains info for running an instance of Toolbox. Should be instantiated with NewServer().
type Server struct {
	version             string
	sqlCommenterEnabled bool
	toolboxUrl          string
	srv                 *http.Server
	listener            net.Listener
	root                chi.Router
	logger              log.Logger
	instrumentation     *telemetry.Instrumentation
	sseManager          *sseManager
	ResourceMgr         *resources.ResourceManager
	mcpPrmFile          string
	httpMaxRequestBytes int64
	enableDraftSpecs    bool
}

func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
	map[string]sources.Source,
	map[string]auth.AuthService,
	map[string]embeddingmodels.EmbeddingModel,
	map[string]tools.Tool,
	map[string]tools.Toolset,
	map[string]prompts.Prompt,
	map[string]prompts.Promptset,
	error,
) {
	if cfg.EnableAPI {
		for _, sc := range cfg.AuthServiceConfigs {
			if sc.IsMCPEnabled() {
				return nil, nil, nil, nil, nil, nil, nil, fmt.Errorf("MCP Auth cannot be enabled together with the legacy HTTP API (EnableAPI)")
			}
		}
	}

	metadataStr := cfg.Version
	if len(cfg.UserAgentMetadata) > 0 {
		metadataStr += "+" + strings.Join(cfg.UserAgentMetadata, "+")
	}
	ctx = util.WithUserAgent(ctx, metadataStr)
	instrumentation, err := util.InstrumentationFromContext(ctx)
	if err != nil {
		return nil, nil, nil, nil, nil, nil, nil, fmt.Errorf("failed to get instrumentation from context: %w", err)
	}

	l, err := util.LoggerFromContext(ctx)
	if err != nil {
		return nil, nil, nil, nil, nil, nil, nil, fmt.Errorf("failed to get logger from context: %w", err)
	}

	// initialize and validate the sources from configs
	sourcesMap := make(map[string]sources.Source)
	for name, sc := range cfg.SourceConfigs {
		s, err := func() (sources.Source, error) {
			childCtx, span := instrumentation.Tracer.Start(
				ctx,
				"toolbox/server/source/init",
				trace.WithAttributes(attribute.String("source_type", sc.SourceConfigType())),
				trace.WithAttributes(attribute.String("source_name", name)),
			)
			defer span.End()
			s, err := sc.Initialize(childCtx, instrumentation.Tracer)
			if err != nil {
				return nil, fmt.Errorf("unable to initialize source %q: %w", name, err)
			}
			return s, nil
		}()
		if err != nil {
			return nil, nil, nil, nil, nil, nil, nil, err
		}
		sourcesMap[name] = s
	}
	sourceNames := make([]string, 0, len(sourcesMap))
	for name := range sourcesMap {
		sourceNames = append(sourceNames, name)
	}
	l.InfoContext(ctx, fmt.Sprintf("Initialized %d sources: %s", len(sourcesMap), strings.Join(sourceNames, ", ")))

	// initialize and validate the auth services from configs
	authServicesMap := make(map[string]auth.AuthService)
	for name, sc := range cfg.AuthServiceConfigs {
		a, err := func() (auth.AuthService, error) {
			_, span := instrumentation.Tracer.Start(
				ctx,
				"toolbox/server/auth/init",
				trace.WithAttributes(attribute.String("auth_type", sc.AuthServiceConfigType())),
				trace.WithAttributes(attribute.String("auth_name", name)),
			)
			defer span.End()
			a, err := sc.Initialize()
			if err != nil {
				return nil, fmt.Errorf("unable to initialize auth service %q: %w", name, err)
			}
			return a, nil
		}()
		if err != nil {
			return nil, nil, nil, nil, nil, nil, nil, err
		}
		authServicesMap[name] = a
	}

	authServiceNames := make([]string, 0, len(authServicesMap))
	for name := range authServicesMap {
		authServiceNames = append(authServiceNames, name)
	}
	l.InfoContext(ctx, fmt.Sprintf("Initialized %d authServices: %s", len(authServicesMap), strings.Join(authServiceNames, ", ")))

	// Initialize and validate embedding models from configs.
	embeddingModelsMap := make(map[string]embeddingmodels.EmbeddingModel)
	for name, ec := range cfg.EmbeddingModelConfigs {
		em, err := func() (embeddingmodels.EmbeddingModel, error) {
			_, span := instrumentation.Tracer.Start(
				ctx,
				"toolbox/server/embeddingmodel/init",
				trace.WithAttributes(attribute.String("model_type", ec.EmbeddingModelConfigType())),
				trace.WithAttributes(attribute.String("model_name", name)),
			)
			defer span.End()
			em, err := ec.Initialize(ctx)
			if err != nil {
				return nil, fmt.Errorf("unable to initialize embedding model %q: %w", name, err)
			}
			return em, nil
		}()
		if err != nil {
			return nil, nil, nil, nil, nil, nil, nil, err
		}
		embeddingModelsMap[name] = em
	}
	embeddingModelNames := make([]string, 0, len(embeddingModelsMap))
	for name := range embeddingModelsMap {
		embeddingModelNames = append(embeddingModelNames, name)
	}
	l.InfoContext(ctx, fmt.Sprintf("Initialized %d embeddingModels: %s", len(embeddingModelsMap), strings.Join(embeddingModelNames, ", ")))

	toolsMap, err := initializeTools(ctx, cfg, instrumentation, l)
	if err != nil {
		return nil, nil, nil, nil, nil, nil, nil, err
	}

	toolsetsMap, err := initializeToolsets(ctx, cfg, toolsMap, instrumentation, l)
	if err != nil {
		return nil, nil, nil, nil, nil, nil, nil, err
	}

	// initialize and validate the prompts from configs
	promptsMap := make(map[string]prompts.Prompt)
	for name, pc := range cfg.PromptConfigs {
		p, err := func() (prompts.Prompt, error) {
			_, span := instrumentation.Tracer.Start(
				ctx,
				"toolbox/server/prompt/init",
				trace.WithAttributes(attribute.String("prompt_type", pc.PromptConfigType())),
				trace.WithAttributes(attribute.String("prompt_name", name)),
			)
			defer span.End()
			p, err := pc.Initialize()
			if err != nil {
				return nil, fmt.Errorf("unable to initialize prompt %q: %w", name, err)
			}
			return p, nil
		}()
		if err != nil {
			return nil, nil, nil, nil, nil, nil, nil, err
		}
		promptsMap[name] = p
	}
	promptNames := make([]string, 0, len(promptsMap))
	for name := range promptsMap {
		promptNames = append(promptNames, name)
	}
	l.InfoContext(ctx, fmt.Sprintf("Initialized %d prompts: %s", len(promptsMap), strings.Join(promptNames, ", ")))

	// create a default promptset that contains all prompts
	allPromptNames := make([]string, 0, len(promptsMap))
	for name := range promptsMap {
		allPromptNames = append(allPromptNames, name)
	}
	if cfg.PromptsetConfigs == nil {
		cfg.PromptsetConfigs = make(PromptsetConfigs)
	}
	cfg.PromptsetConfigs[""] = prompts.PromptsetConfig{Name: "", PromptNames: allPromptNames}

	// initialize and validate the promptsets from configs
	promptsetsMap := make(map[string]prompts.Promptset)
	for name, pc := range cfg.PromptsetConfigs {
		p, err := func() (prompts.Promptset, error) {
			_, span := instrumentation.Tracer.Start(
				ctx,
				"toolbox/server/prompset/init",
				trace.WithAttributes(attribute.String("prompset_name", name)),
			)
			defer span.End()
			p, err := pc.Initialize(cfg.Version, promptsMap)
			if err != nil {
				return prompts.Promptset{}, fmt.Errorf("unable to initialize promptset %q: %w", name, err)
			}
			return p, err
		}()
		if err != nil {
			return nil, nil, nil, nil, nil, nil, nil, err
		}
		promptsetsMap[name] = p
	}
	promptsetNames := make([]string, 0, len(promptsetsMap))
	for name := range promptsetsMap {
		if name == "" {
			promptsetNames = append(promptsetNames, "default")
		} else {
			promptsetNames = append(promptsetNames, name)
		}
	}
	l.InfoContext(ctx, fmt.Sprintf("Initialized %d promptsets: %s", len(promptsetsMap), strings.Join(promptsetNames, ", ")))

	return sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, nil
}

// InitializeOfflineConfigs initializes only tools and toolsets from the config,
// skipping sources, auth services, and embedding models. It backs flows like
// skills-generate that need tool metadata without live source connections.
func InitializeOfflineConfigs(ctx context.Context, cfg ServerConfig) (
	map[string]tools.Tool,
	map[string]tools.Toolset,
	error,
) {
	instrumentation, err := util.InstrumentationFromContext(ctx)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to get instrumentation from context: %w", err)
	}

	l, err := util.LoggerFromContext(ctx)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to get logger from context: %w", err)
	}

	toolsMap, err := initializeTools(ctx, cfg, instrumentation, l)
	if err != nil {
		return nil, nil, err
	}

	toolsetsMap, err := initializeToolsets(ctx, cfg, toolsMap, instrumentation, l)
	if err != nil {
		return nil, nil, err
	}

	return toolsMap, toolsetsMap, nil
}

// initializeTools initializes and validates the tools from the config.
func initializeTools(ctx context.Context, cfg ServerConfig, instrumentation *telemetry.Instrumentation, l log.Logger) (map[string]tools.Tool, error) {
	toolsMap := make(map[string]tools.Tool)
	for name, tc := range cfg.ToolConfigs {
		t, err := func() (tools.Tool, error) {
			_, span := instrumentation.Tracer.Start(
				ctx,
				"toolbox/server/tool/init",
				trace.WithAttributes(attribute.String("tool_type", tc.ToolConfigType())),
				trace.WithAttributes(attribute.String("tool_name", name)),
			)
			defer span.End()
			t, err := tc.Initialize(ctx)
			if err != nil {
				return nil, fmt.Errorf("unable to initialize tool %q: %w", name, err)
			}
			return t, nil
		}()
		if err != nil {
			return nil, err
		}
		toolsMap[name] = t
	}
	toolNames := make([]string, 0, len(toolsMap))
	for name := range toolsMap {
		toolNames = append(toolNames, name)
	}
	l.InfoContext(ctx, fmt.Sprintf("Initialized %d tools: %s", len(toolsMap), strings.Join(toolNames, ", ")))
	return toolsMap, nil
}

// initializeToolsets seeds a default toolset containing all tools, then
// initializes and validates the toolsets from the config.
func initializeToolsets(ctx context.Context, cfg ServerConfig, toolsMap map[string]tools.Tool, instrumentation *telemetry.Instrumentation, l log.Logger) (map[string]tools.Toolset, error) {
	// create a default toolset that contains all tools
	allToolNames := make([]string, 0, len(toolsMap))
	for name := range toolsMap {
		allToolNames = append(allToolNames, name)
	}
	if cfg.ToolsetConfigs == nil {
		cfg.ToolsetConfigs = make(ToolsetConfigs)
	}
	cfg.ToolsetConfigs[""] = tools.ToolsetConfig{Name: "", ToolNames: allToolNames}

	toolsetsMap := make(map[string]tools.Toolset)
	for name, tc := range cfg.ToolsetConfigs {
		if cfg.IgnoreUnknownTools {
			filteredToolNames := make([]string, 0, len(tc.ToolNames))
			for _, tn := range tc.ToolNames {
				if _, ok := toolsMap[tn]; ok {
					filteredToolNames = append(filteredToolNames, tn)
				} else {
					l.WarnContext(ctx, fmt.Sprintf("Skipping missing tool %q in toolset %q", tn, name))
				}
			}
			tc.ToolNames = filteredToolNames
			cfg.ToolsetConfigs[name] = tc
		}

		t, err := func() (tools.Toolset, error) {
			_, span := instrumentation.Tracer.Start(
				ctx,
				"toolbox/server/toolset/init",
				trace.WithAttributes(attribute.String("toolset.name", name)),
			)
			defer span.End()
			t, err := tc.Initialize(cfg.Version, toolsMap)
			if err != nil {
				return tools.Toolset{}, fmt.Errorf("unable to initialize toolset %q: %w", name, err)
			}
			return t, err
		}()
		if err != nil {
			return nil, err
		}
		toolsetsMap[name] = t
	}
	toolsetNames := make([]string, 0, len(toolsetsMap))
	for name := range toolsetsMap {
		if name == "" {
			toolsetNames = append(toolsetNames, "default")
		} else {
			toolsetNames = append(toolsetNames, name)
		}
	}
	l.InfoContext(ctx, fmt.Sprintf("Initialized %d toolsets: %s", len(toolsetsMap), strings.Join(toolsetNames, ", ")))

	return toolsetsMap, nil
}

func hostCheck(allowedHosts map[string]struct{}) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			_, hasWildcard := allowedHosts["*"]
			hostname := r.Host
			if host, _, err := net.SplitHostPort(r.Host); err == nil {
				hostname = host
			}
			_, hostIsAllowed := allowedHosts[hostname]
			if !hasWildcard && !hostIsAllowed {
				// Return 403 Forbidden to block the attack
				http.Error(w, "Invalid Host header", http.StatusForbidden)
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}

// NewServer returns a Server object based on provided Config.
func NewServer(ctx context.Context, cfg ServerConfig) (*Server, error) {
	instrumentation, err := util.InstrumentationFromContext(ctx)
	if err != nil {
		return nil, err
	}

	ctx, span := instrumentation.Tracer.Start(ctx, "toolbox/server/init")
	defer span.End()

	l, err := util.LoggerFromContext(ctx)
	if err != nil {
		return nil, err
	}

	// set up http serving
	r := chi.NewRouter()
	r.Use(middleware.Recoverer)

	// logging
	logLevel, err := log.SeverityToLevel(cfg.LogLevel.String())
	if err != nil {
		return nil, fmt.Errorf("unable to initialize http log: %w", err)
	}

	schema := *httplog.SchemaGCP
	schema.Level = cfg.LogLevel.String()
	schema.Concise(true)
	httpOpts := &httplog.Options{
		Level:  logLevel,
		Schema: &schema,
	}
	logger := l.SlogLogger()
	r.Use(httplog.RequestLogger(logger, httpOpts))

	sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := InitializeConfigs(ctx, cfg)
	if err != nil {
		return nil, fmt.Errorf("unable to initialize configs: %w", err)
	}

	addr := net.JoinHostPort(cfg.Address, strconv.Itoa(cfg.Port))
	srv := &http.Server{Addr: addr, Handler: r}

	sseManager := newSseManager(ctx)

	resourceManager := resources.NewResourceManager(sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap)

	limit := cfg.HttpMaxRequestBytes
	if limit <= 0 {
		limit = DefaultHTTPMaxRequestBytes
	}

	s := &Server{
		version:             cfg.Version,
		sqlCommenterEnabled: cfg.SQLCommenter,
		srv:                 srv,
		root:                r,
		logger:              l,
		instrumentation:     instrumentation,
		sseManager:          sseManager,
		ResourceMgr:         resourceManager,
		toolboxUrl:          cfg.ToolboxUrl,
		mcpPrmFile:          cfg.McpPrmFile,
		httpMaxRequestBytes: limit,
		enableDraftSpecs:    cfg.EnableDraftSpecs,
	}

	if s.enableDraftSpecs {
		s.logger.WarnContext(ctx, "Flag --enable-draft-specs is active. Please note that draft specs are subject to breaking changes and will be completely removed (not redirected) once stable MCP specifications are released. Do not use this configuration in production.")
	}

	// cors
	if slices.Contains(cfg.AllowedOrigins, "*") {
		s.logger.WarnContext(ctx, "wildcard (*) allows any website to access the resources. This creates a security risk regardless of whether you are in a production or local development environment. Recommended to use --allowed-origins with specific local addresses.")
	}
	corsOpts := cors.Options{
		AllowedOrigins:   cfg.AllowedOrigins,
		AllowedMethods:   []string{"GET", "POST", "DELETE", "OPTIONS"},
		AllowCredentials: true, // required since Toolbox uses auth headers
		AllowedHeaders:   []string{"Accept", "Authorization", "Content-Type", "X-CSRF-Token", "Mcp-Session-Id", "MCP-Protocol-Version"},
		ExposedHeaders:   []string{"Mcp-Session-Id"}, // headers that are sent to clients
		MaxAge:           300,                        // cache preflight results for 5 minutes
	}
	r.Use(cors.Handler(corsOpts))
	// validate hosts for DNS rebinding attacks
	if slices.Contains(cfg.AllowedHosts, "*") {
		s.logger.WarnContext(ctx, "wildcard (*) hosts allow any domain to access this resource, making it vulnerable to DNS rebinding attacks regardless of whether you are in a production or local development environment. For improved security, use the --allowed-hosts flag to specify trusted domains.")
	}
	allowedHostsMap := make(map[string]struct{}, len(cfg.AllowedHosts))
	for _, h := range cfg.AllowedHosts {
		hostname := h
		if host, _, err := net.SplitHostPort(h); err == nil {
			hostname = host
		}
		allowedHostsMap[hostname] = struct{}{}
	}
	r.Use(hostCheck(allowedHostsMap))

	// Host OAuth Protected Resource Metadata endpoint
	mcpAuthEnabled := false
	for _, authSvc := range s.ResourceMgr.GetAuthServiceMap() {
		if mSvc, ok := authSvc.(auth.MCPAuthService); ok && mSvc.IsMCPEnabled() {
			mcpAuthEnabled = true
			break
		}
	}

	// Manual PRM override
	var cachedPrmBytes []byte
	var prmConfig ProtectedResourceMetadata
	if s.mcpPrmFile != "" {
		var err error
		cachedPrmBytes, err = os.ReadFile(s.mcpPrmFile)
		if err != nil {
			return nil, fmt.Errorf("failed to read manual PRM file at startup: %w", err)
		}
		// Unmarshal into the struct to strictly validate the schema
		if err := json.Unmarshal(cachedPrmBytes, &prmConfig); err != nil {
			return nil, fmt.Errorf("manual PRM file does not match expected schema: %w", err)
		}
	}

	// Register route if auth is enabled or a manual file is provided
	if mcpAuthEnabled || s.mcpPrmFile != "" {
		r.Get("/.well-known/oauth-protected-resource", func(w http.ResponseWriter, req *http.Request) {
			// Serve from memory if file was loaded
			if s.mcpPrmFile != "" {
				w.Header().Set("Content-Type", "application/json")
				w.WriteHeader(http.StatusOK)
				if _, err := w.Write(cachedPrmBytes); err != nil {
					s.logger.ErrorContext(req.Context(), "failed to write manual PRM file response", "error", err)
				}
				return
			}

			prmHandler(s, w, req)
		})
	}

	// control plane
	mcpR, err := mcpRouter(s)
	if err != nil {
		return nil, err
	}

	r.Mount("/mcp", mcpR)
	if cfg.EnableAPI {
		apiR, err := apiRouter(s)
		if err != nil {
			return nil, err
		}
		r.Mount("/api", apiR)
	} else {
		r.Handle("/api/*", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			err := errors.New("/api native endpoints are disabled by default. Please use the standard /mcp JSON-RPC endpoint")
			_ = render.Render(w, r, newErrResponse(err, http.StatusGone))
		}))
	}
	if cfg.UI {
		webR, err := webRouter()
		if err != nil {
			return nil, err
		}
		r.Mount("/ui", webR)
	}
	// default endpoint for validating server is running
	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte("🧰 Hello, World! 🧰"))
	})

	return s, nil
}

func mcpAuthMiddleware(s *Server) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// Find McpEnabled auth service
			var mcpSvc auth.MCPAuthService
			for _, authSvc := range s.ResourceMgr.GetAuthServiceMap() {
				if mSvc, ok := authSvc.(auth.MCPAuthService); ok && mSvc.IsMCPEnabled() {
					mcpSvc = mSvc
					break
				}
			}

			// MCP Auth not enabled
			if mcpSvc == nil {
				next.ServeHTTP(w, r)
				return
			}

			claims, err := mcpSvc.ValidateMCPAuth(r.Context(), r.Header)
			if err != nil {
				var mcpErr *auth.MCPAuthError
				if errors.As(err, &mcpErr) {
					switch mcpErr.Code {
					case http.StatusUnauthorized:
						scopesArg := ""
						if len(mcpErr.ScopesRequired) > 0 {
							scopesArg = fmt.Sprintf(`, scope="%s"`, strings.Join(mcpErr.ScopesRequired, " "))
						}
						w.Header().Set("WWW-Authenticate", fmt.Sprintf(`Bearer resource_metadata="%s"%s`, s.toolboxUrl+"/.well-known/oauth-protected-resource", scopesArg))
						render.Status(r, http.StatusUnauthorized)
						render.JSON(w, r, jsonrpc.NewError(nil, jsonrpc.UNAUTHORIZED, mcpErr.Message, nil))
						return
					case http.StatusForbidden:
						w.Header().Set("WWW-Authenticate", fmt.Sprintf(`Bearer error="insufficient_scope", scope="%s", resource_metadata="%s", error_description="%s"`, strings.Join(mcpErr.ScopesRequired, " "), s.toolboxUrl+"/.well-known/oauth-protected-resource", mcpErr.Message))
						render.Status(r, http.StatusForbidden)
						render.JSON(w, r, jsonrpc.NewError(nil, jsonrpc.FORBIDDEN, mcpErr.Message, nil))
						return
					}
				}
				// Fail closed on unexpected errors
				s.logger.ErrorContext(r.Context(), "unexpected error during MCP auth validation", "error", err)
				render.Status(r, http.StatusInternalServerError)
				render.JSON(w, r, jsonrpc.NewError(nil, jsonrpc.INTERNAL_ERROR, "Internal Server Error", nil))
				return
			}

			ctx := util.WithAuthTokenClaims(r.Context(), claims)
			r = r.WithContext(ctx)

			next.ServeHTTP(w, r)
		})
	}
}

// Listen starts a listener for the given Server instance.
func (s *Server) Listen(ctx context.Context, certFile, keyFile string) error {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	if s.listener != nil {
		return fmt.Errorf("server is already listening: %s", s.listener.Addr().String())
	}
	lc := net.ListenConfig{KeepAlive: 30 * time.Second}
	ln, err := lc.Listen(ctx, "tcp", s.srv.Addr)
	if err != nil {
		return fmt.Errorf("failed to open listener for %q: %w", s.srv.Addr, err)
	}

	if certFile != "" || keyFile != "" {
		// Load the certificates
		cert, err := tls.LoadX509KeyPair(certFile, keyFile)
		if err != nil {
			ln.Close()
			return fmt.Errorf("failed to load TLS key pair (cert: %q, key: %q): %w", certFile, keyFile, err)
		}
		// Wrap the listener with TLS
		config := &tls.Config{Certificates: []tls.Certificate{cert}, MinVersion: tls.VersionTLS12}
		s.listener = tls.NewListener(ln, config)
		s.logger.DebugContext(ctx, fmt.Sprintf("secure server listening on %s", s.srv.Addr))
	} else {
		s.listener = ln
		s.logger.DebugContext(ctx, fmt.Sprintf("server listening on %s", s.srv.Addr))
	}
	return nil
}

// Serve starts an HTTP server for the given Server instance.
func (s *Server) Serve(ctx context.Context) error {
	s.logger.DebugContext(ctx, "Starting a HTTP server.")
	return s.srv.Serve(s.listener)
}

// ServeStdio starts a new stdio session for mcp.
func (s *Server) ServeStdio(ctx context.Context, stdin io.Reader, stdout io.Writer) error {
	stdioServer := NewStdioSession(s, stdin, stdout)
	return stdioServer.Start(ctx)
}

// Shutdown gracefully shuts down the server without interrupting any active
// connections. It uses http.Server.Shutdown() and has the same functionality.
func (s *Server) Shutdown(ctx context.Context) error {
	s.logger.DebugContext(ctx, "shutting down the server.")
	return s.srv.Shutdown(ctx)
}

func (s *Server) Addr() string {
	return s.listener.Addr().String()
}