package main

import (
	"context"
	"fmt"
	"io"
	"log"
	"net/url"
	"os"
	"strconv"
	"strings"
	"time"

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

type MCPProxy struct {
	remoteURL string
	apiKey    string
	debug     bool
	logger    *log.Logger
	client    *client.Client
	server    *server.MCPServer
}

func NewMCPProxy(remoteURL, apiKey string, debug bool) *MCPProxy {
	logger := log.New(os.Stderr, "[MCP-PROXY] ", log.LstdFlags)
	if !debug {
		logger.SetOutput(io.Discard)
	}

	// Convert /sse to /mcp if needed
	if strings.Contains(remoteURL, "/sse") {
		remoteURL = strings.Replace(remoteURL, "/sse", "/mcp", 1)
		logger.Printf("Converted SSE URL to MCP URL: %s", remoteURL)
	}

	return &MCPProxy{
		remoteURL: remoteURL,
		apiKey:    apiKey,
		debug:     debug,
		logger:    logger,
	}
}

func (p *MCPProxy) createRemoteClient() error {
	p.logger.Printf("Creating remote MCP client for: %s", p.remoteURL)

	// Parse URL and add API key as query parameter
	parsedURL, err := url.Parse(p.remoteURL)
	if err != nil {
		return fmt.Errorf("failed to parse remote URL: %w", err)
	}

	query := parsedURL.Query()
	query.Set("apiKey", p.apiKey)
	parsedURL.RawQuery = query.Encode()
	finalURL := parsedURL.String()

	p.logger.Printf("Final remote URL: %s", finalURL)

	// Create headers with API key
	headers := map[string]string{
		"Authorization": "Bearer " + p.apiKey,
		"User-Agent":    "flipside-mcp-proxy/1.0",
	}

	// Create streamable HTTP client options
	options := []transport.StreamableHTTPCOption{
		transport.WithHTTPHeaders(headers),
		transport.WithHTTPTimeout(30 * time.Second),
	}

	// Create the MCP client
	mcpClient, err := client.NewStreamableHttpClient(finalURL, options...)
	if err != nil {
		return fmt.Errorf("failed to create MCP client: %w", err)
	}

	p.client = mcpClient
	return nil
}

func (p *MCPProxy) setupProxyServer(ctx context.Context) error {
	p.logger.Printf("Setting up MCP proxy server")

	// Create MCP server
	mcpServer := server.NewMCPServer("flipside-remote-mcp-proxy", "1.0.0")

	// Initialize remote client first
	p.logger.Printf("Initializing remote client...")
	_, err := p.client.Initialize(ctx, mcp.InitializeRequest{
		Params: mcp.InitializeParams{
			ProtocolVersion: "2024-11-05",
			Capabilities:    mcp.ClientCapabilities{},
			ClientInfo: mcp.Implementation{
				Name:    "flipside-remote-mcp-proxy",
				Version: "1.0.0",
			},
		},
	})
	if err != nil {
		return fmt.Errorf("failed to initialize remote client: %w", err)
	}

	p.logger.Printf("Remote client initialized successfully")

	// Get tools from remote and add them to proxy server
	if err := p.addRemoteToolsToServer(ctx, mcpServer); err != nil {
		return fmt.Errorf("failed to add remote tools: %w", err)
	}

	// Try to add resources and prompts (may not be supported)
	p.addRemoteResourcesToServer(ctx, mcpServer)
	p.addRemotePromptsToServer(ctx, mcpServer)

	p.server = mcpServer
	return nil
}

func (p *MCPProxy) addRemoteToolsToServer(ctx context.Context, mcpServer *server.MCPServer) error {
	p.logger.Printf("Adding remote tools to proxy server")

	// List tools from remote client
	toolsResponse, err := p.client.ListTools(ctx, mcp.ListToolsRequest{})
	if err != nil {
		p.logger.Printf("Failed to list remote tools: %v", err)
		return err
	}

	p.logger.Printf("Found %d remote tools", len(toolsResponse.Tools))

	// Add each tool to the proxy server
	for _, tool := range toolsResponse.Tools {
		p.logger.Printf("Adding tool: %s", tool.Name)

		// Create a closure to capture the tool
		currentTool := tool

		mcpServer.AddTool(currentTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
			p.logger.Printf("Proxying tool call: %s", currentTool.Name)

			// Forward the tool call to the remote client
			response, err := p.client.CallTool(ctx, mcp.CallToolRequest{
				Params: mcp.CallToolParams{
					Name:      request.Params.Name,
					Arguments: request.Params.Arguments,
				},
			})

			if err != nil {
				p.logger.Printf("Error calling remote tool %s: %v", currentTool.Name, err)
				return nil, err
			}

			return response, nil
		})
	}

	return nil
}

func (p *MCPProxy) addRemoteResourcesToServer(ctx context.Context, mcpServer *server.MCPServer) {
	p.logger.Printf("Adding remote resources to proxy server")

	// List resources from remote client
	resourcesResponse, err := p.client.ListResources(ctx, mcp.ListResourcesRequest{})
	if err != nil {
		p.logger.Printf("Remote server doesn't support resources: %v", err)
		return
	}

	p.logger.Printf("Found %d remote resources", len(resourcesResponse.Resources))

	// Add each resource to the proxy server
	for _, resource := range resourcesResponse.Resources {
		p.logger.Printf("Adding resource: %s", resource.URI)

		mcpServer.AddResource(resource, func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
			p.logger.Printf("Proxying read resource: %s", request.Params.URI)

			response, err := p.client.ReadResource(ctx, request)
			if err != nil {
				return nil, err
			}

			return response.Contents, nil
		})
	}
}

func (p *MCPProxy) addRemotePromptsToServer(ctx context.Context, mcpServer *server.MCPServer) {
	p.logger.Printf("Adding remote prompts to proxy server")

	// List prompts from remote client
	promptsResponse, err := p.client.ListPrompts(ctx, mcp.ListPromptsRequest{})
	if err != nil {
		p.logger.Printf("Remote server doesn't support prompts: %v", err)
		return
	}

	p.logger.Printf("Found %d remote prompts", len(promptsResponse.Prompts))

	// Add each prompt to the proxy server
	for _, prompt := range promptsResponse.Prompts {
		p.logger.Printf("Adding prompt: %s", prompt.Name)

		currentPrompt := prompt

		mcpServer.AddPrompt(currentPrompt, func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
			p.logger.Printf("Proxying get prompt: %s", request.Params.Name)

			response, err := p.client.GetPrompt(ctx, request)
			if err != nil {
				return nil, err
			}

			return response, nil
		})
	}
}

func (p *MCPProxy) run() error {
	ctx := context.Background()

	// Create remote client
	if err := p.createRemoteClient(); err != nil {
		return fmt.Errorf("failed to create remote client: %w", err)
	}

	// Setup proxy server (this initializes the remote client and discovers tools)
	if err := p.setupProxyServer(ctx); err != nil {
		return fmt.Errorf("failed to setup proxy server: %w", err)
	}

	p.logger.Printf("Starting MCP proxy server with stdio transport")
	p.logger.Printf("Remote URL: %s", p.remoteURL)
	p.logger.Printf("Debug mode: %v", p.debug)

	// Start the proxy server with stdio transport
	if err := server.ServeStdio(p.server); err != nil {
		return fmt.Errorf("server error: %w", err)
	}

	return nil
}

func maskAPIKey(apiKey string) string {
	if len(apiKey) <= 8 {
		return "***"
	}
	return apiKey[:4] + "..." + apiKey[len(apiKey)-4:]
}

func main() {
	// Always enable initial logging to stderr for startup diagnostics
	startupLogger := log.New(os.Stderr, "[MCP-PROXY-STARTUP] ", log.LstdFlags)
	startupLogger.Println("Starting Flipside MCP Remote Proxy...")

	remoteURL := os.Getenv("MCP_REMOTE_URL")
	if remoteURL == "" {
		startupLogger.Fatal("ERROR: MCP_REMOTE_URL environment variable is required")
	}
	startupLogger.Printf("Remote URL configured: %s", remoteURL)

	apiKey := os.Getenv("FLIPSIDE_API_KEY")
	if apiKey == "" {
		startupLogger.Fatal("ERROR: FLIPSIDE_API_KEY environment variable is required")
	}
	startupLogger.Printf("API key configured: %s", maskAPIKey(apiKey))

	debugStr := os.Getenv("MCP_DEBUG")
	debug, err := strconv.ParseBool(debugStr)
	if err != nil {
		debug = false // Default to false if parsing fails
	}
	startupLogger.Printf("Debug mode: %v", debug)

	proxy := NewMCPProxy(remoteURL, apiKey, debug)

	startupLogger.Println("Starting MCP proxy with mcp-go client/server architecture...")

	if err := proxy.run(); err != nil {
		startupLogger.Fatalf("Proxy error: %v", err)
	}
}