package mcpserver

import (
	"context"
	"fmt"

	"github.com/giantswarm/muster/pkg/logging"
	"github.com/giantswarm/muster/pkg/observability"

	"github.com/mark3labs/mcp-go/client"
	"github.com/mark3labs/mcp-go/client/transport"
	"github.com/mark3labs/mcp-go/mcp"
	mcpotel "github.com/mark3labs/mcp-go/otel"
	"go.opentelemetry.io/otel"
)

// StreamableHTTPClient implements the MCPClient interface using StreamableHTTP transport.
// It connects to remote MCP servers using HTTP with streaming support.
type StreamableHTTPClient struct {
	baseMCPClient
	url        string
	headers    map[string]string
	headerFunc transport.HTTPHeaderFunc // Dynamic header function called on each request
}

// NewStreamableHTTPClientWithHeaders creates a new StreamableHTTP-based MCP client with custom headers
func NewStreamableHTTPClientWithHeaders(url string, headers map[string]string) *StreamableHTTPClient {
	if headers == nil {
		headers = make(map[string]string)
	}
	return &StreamableHTTPClient{
		url:     url,
		headers: headers,
	}
}

// NewStreamableHTTPClientWithHeaderFunc creates a new StreamableHTTP-based MCP client
// with a dynamic header function that is called on every request. This enables token
// refresh by resolving the latest token at call time instead of baking in a static header.
func NewStreamableHTTPClientWithHeaderFunc(url string, headerFunc transport.HTTPHeaderFunc) *StreamableHTTPClient {
	return &StreamableHTTPClient{
		url:        url,
		headers:    make(map[string]string),
		headerFunc: headerFunc,
	}
}

// Initialize establishes the connection and performs protocol handshake
func (c *StreamableHTTPClient) Initialize(ctx context.Context) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.connected {
		return nil
	}

	logging.Debug("StreamableHTTPClient", "Creating StreamableHTTP client for URL: %s", c.url)

	// Build client options including headers if provided
	var opts []transport.StreamableHTTPCOption
	if c.headerFunc != nil {
		opts = append(opts, transport.WithHTTPHeaderFunc(c.headerFunc))
		logging.Debug("StreamableHTTPClient", "Configured dynamic header function")
	} else if len(c.headers) > 0 {
		opts = append(opts, transport.WithHTTPHeaders(c.headers))
		logging.Debug("StreamableHTTPClient", "Configured %d custom headers", len(c.headers))
	}

	// Enable receiving server-pushed notifications outside active requests.
	// This opens a long-lived GET connection to the server per the MCP spec.
	opts = append(opts, transport.WithContinuousListening())

	mcpClient, err := client.NewStreamableHttpClient(c.url, opts...)
	if err != nil {
		return fmt.Errorf("failed to create StreamableHTTP client: %w", err)
	}
	mcpotel.WithClientTracing(otel.Tracer(observability.TracerName))(mcpClient)

	// Start with a background context so the continuous GET listener goroutine
	// survives after the caller's initialization context (which may be short-lived) completes.
	// The listener goroutine is separately cancelled when the client is closed.
	if err := mcpClient.Start(context.Background()); err != nil {
		_ = mcpClient.Close()
		if authErr := CheckForAuthRequiredError(ctx, err, c.url); authErr != nil {
			logging.Debug("StreamableHTTPClient", "Authentication required for URL: %s", c.url)
			return authErr
		}
		return fmt.Errorf("failed to start StreamableHTTP transport: %w", err)
	}

	initResult, err := mcpClient.Initialize(ctx, mcp.InitializeRequest{
		Params: struct {
			ProtocolVersion string                 `json:"protocolVersion"`
			Capabilities    mcp.ClientCapabilities `json:"capabilities"`
			ClientInfo      mcp.Implementation     `json:"clientInfo"`
		}{
			ProtocolVersion: "2024-11-05",
			ClientInfo: mcp.Implementation{
				Name:    "muster-aggregator",
				Version: "1.0.0",
			},
			Capabilities: mcp.ClientCapabilities{},
		},
	})
	if err != nil {
		_ = mcpClient.Close()

		// Check if this is a 401 authentication error
		if authErr := CheckForAuthRequiredError(ctx, err, c.url); authErr != nil {
			logging.Debug("StreamableHTTPClient", "Authentication required for URL: %s", c.url)
			return authErr
		}

		return fmt.Errorf("failed to initialize MCP protocol: %w", err)
	}

	c.client = mcpClient
	c.connected = true
	c.wireNotificationHandler()

	logging.Debug("StreamableHTTPClient", "StreamableHTTP client initialized. Server: %s, Version: %s",
		initResult.ServerInfo.Name, initResult.ServerInfo.Version)

	return nil
}

// Close cleanly shuts down the client connection
func (c *StreamableHTTPClient) Close() error {
	return c.closeClient()
}

// ListTools returns all available tools from the server
func (c *StreamableHTTPClient) ListTools(ctx context.Context) ([]mcp.Tool, error) {
	return c.listTools(ctx)
}

// CallTool executes a specific tool and returns the result
func (c *StreamableHTTPClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*mcp.CallToolResult, error) {
	return c.callTool(ctx, name, args)
}

// ListResources returns all available resources from the server
func (c *StreamableHTTPClient) ListResources(ctx context.Context) ([]mcp.Resource, error) {
	return c.listResources(ctx)
}

// ReadResource retrieves a specific resource
func (c *StreamableHTTPClient) ReadResource(ctx context.Context, uri string) (*mcp.ReadResourceResult, error) {
	return c.readResource(ctx, uri)
}

// ListPrompts returns all available prompts from the server
func (c *StreamableHTTPClient) ListPrompts(ctx context.Context) ([]mcp.Prompt, error) {
	return c.listPrompts(ctx)
}

// GetPrompt retrieves a specific prompt
func (c *StreamableHTTPClient) GetPrompt(ctx context.Context, name string, args map[string]interface{}) (*mcp.GetPromptResult, error) {
	return c.getPrompt(ctx, name, args)
}

// Ping checks if the server is responsive
func (c *StreamableHTTPClient) Ping(ctx context.Context) error {
	return c.ping(ctx)
}

// OnNotification registers a handler for server-pushed notifications.
func (c *StreamableHTTPClient) OnNotification(handler func(mcp.JSONRPCNotification)) {
	c.onNotification(handler)
}
