package tools

import (
	"context"
	"fmt"
	"maps"
	"strings"
	"sync"
	"time"

	"charm.land/fantasy"
	"github.com/mark3labs/kit/internal/config"
	"github.com/mark3labs/mcp-go/client"
	"github.com/mark3labs/mcp-go/client/transport"
	"github.com/mark3labs/mcp-go/mcp"
)

// ConnectionPoolConfig defines configuration parameters for the MCP connection pool.
// It controls connection lifecycle, health checking, and error handling behaviors.
type ConnectionPoolConfig struct {
	MaxIdleTime         time.Duration // Maximum time a connection can remain idle before being marked unhealthy
	MaxRetries          int           // Maximum number of retry attempts for failed operations
	HealthCheckInterval time.Duration // Interval between background health checks of all connections
	MaxErrorCount       int           // Maximum consecutive errors before marking a connection unhealthy
	ReconnectDelay      time.Duration // Delay before attempting to reconnect after connection failure
}

// DefaultConnectionPoolConfig returns a connection pool configuration with sensible defaults.
// Default values: 5 minute max idle time, 3 retries, 30 second health check interval,
// 3 max errors before marking unhealthy, and 2 second reconnect delay.
// These defaults are suitable for most MCP server deployments.
func DefaultConnectionPoolConfig() *ConnectionPoolConfig {
	return &ConnectionPoolConfig{
		MaxIdleTime:         5 * time.Minute,
		MaxRetries:          3,
		HealthCheckInterval: 30 * time.Second,
		MaxErrorCount:       3,
		ReconnectDelay:      2 * time.Second,
	}
}

// MCPConnection represents a single MCP client connection with health tracking and metadata.
// It wraps an MCP client and maintains state about connection health, usage patterns,
// and error history. Access to connection state is protected by a read-write mutex
// for thread-safe concurrent access.
type MCPConnection struct {
	client       client.MCPClient
	serverName   string
	serverConfig config.MCPServerConfig
	lastUsed     time.Time
	isHealthy    bool
	errorCount   int
	lastError    error
	mu           sync.RWMutex
}

// MCPConnectionPool manages a pool of MCP client connections with automatic health checking,
// connection reuse, and failure recovery. It provides thread-safe connection management
// across multiple MCP servers, automatically handling connection lifecycle including
// creation, health monitoring, and cleanup. The pool runs background health checks
// to proactively identify and remove unhealthy connections.
type MCPConnectionPool struct {
	connections map[string]*MCPConnection
	config      *ConnectionPoolConfig
	mu          sync.RWMutex
	model       fantasy.LanguageModel
	ctx         context.Context
	cancel      context.CancelFunc
	debug       bool
	debugLogger DebugLogger
	oauthFlow   *OAuthFlowRunner
}

// NewMCPConnectionPool creates a new MCP connection pool with the specified configuration.
// If config is nil, default configuration values will be used. The pool starts a background
// goroutine for periodic health checks that runs until Close is called.
// The model parameter is used for MCP servers that require sampling support.
// Thread-safe for concurrent use immediately after creation.
func NewMCPConnectionPool(config *ConnectionPoolConfig, model fantasy.LanguageModel, debug bool, authHandler MCPAuthHandler) *MCPConnectionPool {
	if config == nil {
		config = DefaultConnectionPoolConfig()
	}

	ctx, cancel := context.WithCancel(context.Background())
	pool := &MCPConnectionPool{
		connections: make(map[string]*MCPConnection),
		config:      config,
		model:       model,
		ctx:         ctx,
		cancel:      cancel,
		debug:       debug,
	}

	if authHandler != nil {
		pool.oauthFlow = NewOAuthFlowRunner(authHandler)
	}

	go pool.startHealthCheck()
	return pool
}

// SetDebugLogger sets the debug logger for the connection pool.
// The logger will be used to output detailed information about connection lifecycle,
// health checks, and error conditions. Thread-safe and can be called at any time.
func (p *MCPConnectionPool) SetDebugLogger(logger DebugLogger) {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.debugLogger = logger
}

// SetOAuthFlow sets the OAuth flow runner for the connection pool.
// When set, the pool can trigger OAuth re-authorization when a tool call fails
// with an OAuth error (e.g. expired token). Thread-safe and can be called at any time.
func (p *MCPConnectionPool) SetOAuthFlow(flow *OAuthFlowRunner) {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.oauthFlow = flow
}

// GetConnection retrieves or creates a connection for the specified MCP server.
// If a healthy, non-idle connection exists in the pool, it will be reused.
// Otherwise, a new connection is created and added to the pool.
// Returns an error if connection creation or initialization fails.
// Thread-safe for concurrent calls.
func (p *MCPConnectionPool) GetConnection(ctx context.Context, serverName string, serverConfig config.MCPServerConfig) (*MCPConnection, error) {
	p.mu.Lock()
	defer p.mu.Unlock()

	if conn, exists := p.connections[serverName]; exists {
		conn.mu.RLock()
		isHealthy := conn.isHealthy && time.Since(conn.lastUsed) < p.config.MaxIdleTime
		conn.mu.RUnlock()

		if isHealthy {
			conn.mu.Lock()
			conn.lastUsed = time.Now()
			conn.mu.Unlock()
			if p.debugLogger != nil && p.debugLogger.IsDebugEnabled() {
				p.debugLogger.LogDebug(fmt.Sprintf("[POOL] Reusing connection for %s", serverName))
			}
			return conn, nil
		} else {
			if p.debugLogger != nil && p.debugLogger.IsDebugEnabled() {
				p.debugLogger.LogDebug(fmt.Sprintf("[POOL] Connection %s unhealthy, removing", serverName))
			}
			_ = conn.client.Close()
			delete(p.connections, serverName)
		}
	}

	if p.debugLogger != nil && p.debugLogger.IsDebugEnabled() {
		p.debugLogger.LogDebug(fmt.Sprintf("[POOL] Creating new connection for %s", serverName))
	}
	conn, err := p.createConnection(ctx, serverName, serverConfig)
	if err != nil {
		return nil, fmt.Errorf("failed to create connection for %s: %w", serverName, err)
	}

	p.connections[serverName] = conn
	return conn, nil
}

// GetConnectionWithHealthCheck retrieves a connection with an additional proactive health check.
// Unlike GetConnection, this method performs a health check on existing connections before
// returning them, ensuring the connection is truly healthy. This is useful for critical
// operations where connection reliability is paramount. Creates a new connection if the
// existing one fails the health check or doesn't exist.
// Thread-safe for concurrent calls.
func (p *MCPConnectionPool) GetConnectionWithHealthCheck(ctx context.Context, serverName string, serverConfig config.MCPServerConfig) (*MCPConnection, error) {
	p.mu.Lock()
	defer p.mu.Unlock()

	if conn, exists := p.connections[serverName]; exists {
		conn.mu.RLock()
		isHealthy := conn.isHealthy && time.Since(conn.lastUsed) < p.config.MaxIdleTime
		conn.mu.RUnlock()

		if isHealthy {
			// Perform proactive health check before reusing connection
			if p.performHealthCheck(ctx, conn) {
				conn.mu.Lock()
				conn.lastUsed = time.Now()
				conn.mu.Unlock()
				if p.debugLogger != nil && p.debugLogger.IsDebugEnabled() {
					p.debugLogger.LogDebug(fmt.Sprintf("[POOL] Reusing healthy connection for %s", serverName))
				}
				return conn, nil
			} else {
				if p.debugLogger != nil && p.debugLogger.IsDebugEnabled() {
					p.debugLogger.LogDebug(fmt.Sprintf("[POOL] Connection %s failed health check, removing", serverName))
				}
				_ = conn.client.Close()
				delete(p.connections, serverName)
			}
		} else {
			if p.debugLogger != nil && p.debugLogger.IsDebugEnabled() {
				p.debugLogger.LogDebug(fmt.Sprintf("[POOL] Connection %s unhealthy, removing", serverName))
			}
			_ = conn.client.Close()
			delete(p.connections, serverName)
		}
	}

	if p.debugLogger != nil && p.debugLogger.IsDebugEnabled() {
		p.debugLogger.LogDebug(fmt.Sprintf("[POOL] Creating new connection for %s", serverName))
	}
	conn, err := p.createConnection(ctx, serverName, serverConfig)
	if err != nil {
		return nil, fmt.Errorf("failed to create connection for %s: %w", serverName, err)
	}

	p.connections[serverName] = conn
	return conn, nil
}

// performHealthCheck performs a quick health check on the connection
func (p *MCPConnectionPool) performHealthCheck(ctx context.Context, conn *MCPConnection) bool {
	// Create a short timeout context for health check
	healthCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	// Try to list tools as a health check - this is a lightweight operation
	_, err := conn.client.ListTools(healthCtx, mcp.ListToolsRequest{})
	if err != nil {
		if p.debugLogger != nil && p.debugLogger.IsDebugEnabled() {
			p.debugLogger.LogDebug(fmt.Sprintf("[HEALTH_CHECK] Connection %s failed health check: %v", conn.serverName, err))
		}
		conn.mu.Lock()
		conn.isHealthy = false
		conn.errorCount++
		conn.lastError = err
		conn.mu.Unlock()
		return false
	}

	// Reset error count on successful health check
	conn.mu.Lock()
	conn.errorCount = 0
	conn.lastError = nil
	conn.mu.Unlock()

	return true
}

// createConnection creates a new connection
func (p *MCPConnectionPool) createConnection(ctx context.Context, serverName string, serverConfig config.MCPServerConfig) (*MCPConnection, error) {
	mcpClient, err := p.createMCPClient(ctx, serverName, serverConfig)
	if err != nil {
		// SSE transport can return OAuth error during Start()
		if p.oauthFlow != nil && IsOAuthError(err) {
			if flowErr := p.oauthFlow.RunAuthFlow(ctx, serverName, err); flowErr != nil {
				return nil, fmt.Errorf("OAuth authorization failed: %w", flowErr)
			}
			// Retry after successful auth
			mcpClient, err = p.createMCPClient(ctx, serverName, serverConfig)
			if err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	}

	if err := p.initializeClient(ctx, mcpClient); err != nil {
		// Streamable HTTP transport returns OAuth error during Initialize()
		if p.oauthFlow != nil && IsOAuthError(err) {
			if flowErr := p.oauthFlow.RunAuthFlow(ctx, serverName, err); flowErr != nil {
				_ = mcpClient.Close()
				return nil, fmt.Errorf("OAuth authorization failed: %w", flowErr)
			}
			// Retry initialization after successful auth
			if err := p.initializeClient(ctx, mcpClient); err != nil {
				_ = mcpClient.Close()
				return nil, err
			}
		} else {
			_ = mcpClient.Close()
			return nil, err
		}
	}

	conn := &MCPConnection{
		client:       mcpClient,
		serverName:   serverName,
		serverConfig: serverConfig,
		lastUsed:     time.Now(),
		isHealthy:    true,
		errorCount:   0,
		lastError:    nil,
	}

	if p.debugLogger != nil && p.debugLogger.IsDebugEnabled() {
		p.debugLogger.LogDebug(fmt.Sprintf("[POOL] Created connection for %s", serverName))
	}
	return conn, nil
}

// createMCPClient creates an MCP client
func (p *MCPConnectionPool) createMCPClient(ctx context.Context, serverName string, serverConfig config.MCPServerConfig) (client.MCPClient, error) {
	transportType := serverConfig.GetTransportType()

	switch transportType {
	case "stdio":
		return p.createStdioClient(ctx, serverConfig)
	case "sse":
		return p.createSSEClient(ctx, serverConfig)
	case "streamable":
		return p.createStreamableClient(ctx, serverConfig)
	default:
		return nil, fmt.Errorf("unsupported transport type '%s' for server %s", transportType, serverName)
	}
}

// createStdioClient creates a STDIO client
func (p *MCPConnectionPool) createStdioClient(ctx context.Context, serverConfig config.MCPServerConfig) (client.MCPClient, error) {
	var env []string
	var command string
	var args []string

	if len(serverConfig.Command) > 0 {
		command = serverConfig.Command[0]
		if len(serverConfig.Command) > 1 {
			args = serverConfig.Command[1:]
		} else if len(serverConfig.Args) > 0 {
			args = serverConfig.Args
		}
	}

	if serverConfig.Environment != nil {
		for k, v := range serverConfig.Environment {
			env = append(env, fmt.Sprintf("%s=%s", k, v))
		}
	}

	stdioTransport := transport.NewStdio(command, env, args...)
	stdioClient := client.NewClient(stdioTransport)

	if err := stdioTransport.Start(ctx); err != nil {
		return nil, fmt.Errorf("failed to start stdio transport: %v", err)
	}

	time.Sleep(100 * time.Millisecond)
	return stdioClient, nil
}

// createSSEClient creates an SSE client
func (p *MCPConnectionPool) createSSEClient(ctx context.Context, serverConfig config.MCPServerConfig) (client.MCPClient, error) {
	var options []transport.ClientOption

	if len(serverConfig.Headers) > 0 {
		headers := make(map[string]string)
		for _, header := range serverConfig.Headers {
			parts := strings.SplitN(header, ":", 2)
			if len(parts) == 2 {
				key := strings.TrimSpace(parts[0])
				value := strings.TrimSpace(parts[1])
				headers[key] = value
			}
		}
		if len(headers) > 0 {
			options = append(options, transport.WithHeaders(headers))
		}
	}

	// Enable OAuth for remote transports when an auth handler is configured.
	// The OAuthConfig uses PKCE and the handler's redirect URI. Client ID and
	// scopes are discovered automatically via dynamic client registration and
	// server metadata (RFC 9728).
	if p.oauthFlow != nil {
		tokenStore, tsErr := NewFileTokenStore(serverConfig.URL)
		if tsErr != nil {
			return nil, fmt.Errorf("failed to create token store: %w", tsErr)
		}
		options = append(options, transport.WithOAuth(transport.OAuthConfig{
			RedirectURI: p.oauthFlow.handler.RedirectURI(),
			PKCEEnabled: true,
			TokenStore:  tokenStore,
		}))
	}

	sseClient, err := client.NewSSEMCPClient(serverConfig.URL, options...)
	if err != nil {
		return nil, err
	}

	if err := sseClient.Start(ctx); err != nil {
		return nil, fmt.Errorf("failed to start SSE client: %w", err)
	}

	return sseClient, nil
}

// createStreamableClient creates a Streamable client
func (p *MCPConnectionPool) createStreamableClient(ctx context.Context, serverConfig config.MCPServerConfig) (client.MCPClient, error) {
	var options []transport.StreamableHTTPCOption

	if len(serverConfig.Headers) > 0 {
		headers := make(map[string]string)
		for _, header := range serverConfig.Headers {
			parts := strings.SplitN(header, ":", 2)
			if len(parts) == 2 {
				key := strings.TrimSpace(parts[0])
				value := strings.TrimSpace(parts[1])
				headers[key] = value
			}
		}
		if len(headers) > 0 {
			options = append(options, transport.WithHTTPHeaders(headers))
		}
	}

	// Enable OAuth for remote transports when an auth handler is configured.
	// The OAuthConfig uses PKCE and the handler's redirect URI. Client ID and
	// scopes are discovered automatically via dynamic client registration and
	// server metadata (RFC 9728).
	if p.oauthFlow != nil {
		tokenStore, tsErr := NewFileTokenStore(serverConfig.URL)
		if tsErr != nil {
			return nil, fmt.Errorf("failed to create token store: %w", tsErr)
		}
		options = append(options, transport.WithHTTPOAuth(transport.OAuthConfig{
			RedirectURI: p.oauthFlow.handler.RedirectURI(),
			PKCEEnabled: true,
			TokenStore:  tokenStore,
		}))
	}

	streamableClient, err := client.NewStreamableHttpClient(serverConfig.URL, options...)
	if err != nil {
		return nil, err
	}

	if err := streamableClient.Start(ctx); err != nil {
		return nil, fmt.Errorf("failed to start streamable HTTP client: %w", err)
	}

	return streamableClient, nil
}

// initializeClient initializes the client
func (p *MCPConnectionPool) initializeClient(ctx context.Context, client client.MCPClient) error {
	initCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
	defer cancel()

	initRequest := mcp.InitializeRequest{}
	initRequest.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
	initRequest.Params.ClientInfo = mcp.Implementation{
		Name:    "kit",
		Version: "1.0.0",
	}
	initRequest.Params.Capabilities = mcp.ClientCapabilities{}

	_, err := client.Initialize(initCtx, initRequest)
	if err != nil {
		return fmt.Errorf("initialization timeout or failed: %w", err)
	}

	if p.debugLogger != nil && p.debugLogger.IsDebugEnabled() {
		p.debugLogger.LogDebug("[POOL] Initialized MCP client")
	}
	return nil
}

// startHealthCheck starts the health check routine
func (p *MCPConnectionPool) startHealthCheck() {
	ticker := time.NewTicker(p.config.HealthCheckInterval)
	defer ticker.Stop()

	for {
		select {
		case <-p.ctx.Done():
			return
		case <-ticker.C:
			p.checkConnectionsHealth()
		}
	}
}

// checkConnectionsHealth checks the health of all connections
func (p *MCPConnectionPool) checkConnectionsHealth() {
	p.mu.RLock()
	connections := make(map[string]*MCPConnection)
	maps.Copy(connections, p.connections)
	p.mu.RUnlock()

	for serverName, conn := range connections {
		conn.mu.Lock()

		if time.Since(conn.lastUsed) > p.config.MaxIdleTime {
			conn.isHealthy = false
			if p.debugLogger != nil && p.debugLogger.IsDebugEnabled() {
				p.debugLogger.LogDebug(fmt.Sprintf("[HEALTH_CHECK] Connection %s marked as unhealthy due to inactivity", serverName))
			}
		}

		if conn.errorCount > p.config.MaxErrorCount {
			conn.isHealthy = false
			if p.debugLogger != nil && p.debugLogger.IsDebugEnabled() {
				p.debugLogger.LogDebug(fmt.Sprintf("[HEALTH_CHECK] Connection %s marked as unhealthy due to errors", serverName))
			}
		}

		conn.mu.Unlock()
	}
}

// HandleConnectionError records and handles errors for a specific connection.
// It increments the error count and may mark the connection as unhealthy based on
// the error type and configured thresholds. Connection errors (network, transport, 404)
// immediately mark the connection as unhealthy for removal on next access.
// Thread-safe for concurrent error reporting.
func (p *MCPConnectionPool) HandleConnectionError(serverName string, err error) {
	p.mu.RLock()
	conn, exists := p.connections[serverName]
	p.mu.RUnlock()

	if !exists {
		return
	}

	conn.mu.Lock()
	defer conn.mu.Unlock()

	conn.errorCount++
	conn.lastError = err

	if isConnectionError(err) {
		conn.isHealthy = false
		if p.debugLogger != nil && p.debugLogger.IsDebugEnabled() {
			p.debugLogger.LogDebug(fmt.Sprintf("[POOL] Connection %s unhealthy, removing", serverName))
		}

		if strings.Contains(err.Error(), "404") {
			if p.debugLogger != nil && p.debugLogger.IsDebugEnabled() {
				p.debugLogger.LogDebug(fmt.Sprintf("[POOL] 404 error for %s, will recreate on next request", serverName))
			}
		}
	}
}

// GetConnectionStats returns detailed statistics for all connections in the pool.
// The returned map includes health status, last usage time, error counts, and
// last error for each connection. Useful for monitoring and debugging connection
// pool behavior. The returned data is a snapshot and safe for concurrent access.
func (p *MCPConnectionPool) GetConnectionStats() map[string]any {
	p.mu.RLock()
	defer p.mu.RUnlock()

	stats := make(map[string]any)
	for serverName, conn := range p.connections {
		conn.mu.RLock()
		stats[serverName] = map[string]any{
			"is_healthy":  conn.isHealthy,
			"last_used":   conn.lastUsed,
			"last_error":  conn.lastError,
			"error_count": conn.errorCount,
			"server_name": conn.serverName,
		}
		conn.mu.RUnlock()
	}
	return stats
}

// ServerName returns the server name associated with this MCP connection.
// This is the configured name from the KIT configuration, not necessarily
// the actual server implementation name.
func (c *MCPConnection) ServerName() string {
	return c.serverName
}

// GetClients returns a map of all MCP clients currently in the pool.
// The map keys are server names and values are the corresponding MCP client instances.
// The returned map is a copy and modifications won't affect the pool.
// Note that clients may be unhealthy; use GetConnectionStats to check health status.
func (p *MCPConnectionPool) GetClients() map[string]client.MCPClient {
	p.mu.RLock()
	defer p.mu.RUnlock()

	clients := make(map[string]client.MCPClient)
	for name, conn := range p.connections {
		clients[name] = conn.client
	}
	return clients
}

// Close gracefully shuts down the connection pool, closing all client connections
// and stopping the background health check goroutine. It attempts to close all
// connections even if some fail, logging any errors encountered.
// Safe to call multiple times; subsequent calls are no-ops.
// Always call Close when done with the pool to prevent resource leaks.
func (p *MCPConnectionPool) Close() error {
	p.cancel()

	p.mu.Lock()
	defer p.mu.Unlock()

	for name, conn := range p.connections {
		if err := conn.client.Close(); err != nil {
			if p.debugLogger != nil && p.debugLogger.IsDebugEnabled() {
				p.debugLogger.LogDebug(fmt.Sprintf("[POOL] Failed to close connection %s: %v", name, err))
			}
		}
	}

	if p.debugLogger != nil && p.debugLogger.IsDebugEnabled() {
		p.debugLogger.LogDebug("[POOL] Connection pool closed")
	}
	return nil
}

// isConnectionError checks if the error is connection-related
func isConnectionError(err error) bool {
	if IsOAuthError(err) {
		return false // OAuth errors are recoverable, not connection failures
	}
	errStr := err.Error()
	return strings.Contains(errStr, "Connection not found") ||
		strings.Contains(errStr, "transport error") ||
		strings.Contains(errStr, "request failed with status 404") ||
		strings.Contains(errStr, "connection refused") ||
		strings.Contains(errStr, "no such host") ||
		strings.Contains(errStr, "Client.Timeout exceeded")
}
