package sidekick

import (
  "context"
  "fmt"

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

// MCPClientInterface represents a scoped connection to an MCP server
type MCPClientInterface interface {
  CallTool(ctx context.Context, name string, args map[string]interface{}) (string, error)
  Close() error
}

type mcpGoClientWrapper struct {
  client *client.Client
}

func (w *mcpGoClientWrapper) CallTool(ctx context.Context, name string,
  args map[string]interface{}) (string, error) {
  req := mcp.CallToolRequest{}
  req.Params.Name = name
  req.Params.Arguments = args

  res, err := w.client.CallTool(ctx, req)
  if err != nil {
    return "", err
  }

  if res.IsError {
    return "", fmt.Errorf("tool returned error")
  }

  // Extract text from result content
  var output string
  for _, c := range res.Content {
    if textContent, ok := c.(mcp.TextContent); ok {
      output += textContent.Text + "\n"
    } else {
      output += fmt.Sprintf("%v\n", c)
    }
  }
  return output, nil
}

func (w *mcpGoClientWrapper) Close() error {
  // Client does not expose a close method directly in all transports, or we close transport
  return nil
}

// Global server configurations
var mcpServers map[string]MCPServerConfig

// GlobalLogger is an optional logger for MCP connections
var GlobalLogger func(format string, v ...interface{})

// InitMCPServers sets up the global registry for MCP factories
func InitMCPServers(servers map[string]MCPServerConfig) {
  mcpServers = servers
}

// defaultMCPFactory creates a real MCP client based on the global configuration
var defaultMCPFactory MCPClientFactory = func(toolset string) (MCPClientInterface, error) {
  cfg, ok := mcpServers[toolset]
  if !ok {
    return nil, fmt.Errorf("unknown MCP server: %s", toolset)
  }

  if GlobalLogger != nil {
    GlobalLogger("Initializing MCP connection to %s via %s", toolset, cfg.Transport)
  }

  var c *client.Client
  var err error

  switch cfg.Transport {
  case "stdio":
    c, err = client.NewStdioMCPClient(cfg.Command, cfg.Env, cfg.Args...)
  case "sse":
    var opts []transport.ClientOption
    headers := make(map[string]string)
    if cfg.AuthToken != "" {
      headers["Authorization"] = "Bearer " + cfg.AuthToken
    }
    for k, v := range cfg.Headers {
      headers[k] = v
    }
    if len(headers) > 0 {
      opts = append(opts, transport.WithHeaders(headers))
    }
    c, err = client.NewSSEMCPClient(cfg.URL, opts...)
  case "http":
    var opts []transport.StreamableHTTPCOption
    headers := make(map[string]string)
    if cfg.AuthToken != "" {
      headers["Authorization"] = "Bearer " + cfg.AuthToken
    }
    for k, v := range cfg.Headers {
      headers[k] = v
    }
    if len(headers) > 0 {
      opts = append(opts, transport.WithHTTPHeaders(headers))
    }
    c, err = client.NewStreamableHttpClient(cfg.URL, opts...)
  default:
    return nil, fmt.Errorf("unsupported MCP transport: %s", cfg.Transport)
  }

  if err != nil {
    return nil, err
  }

  // Initialize the client session
  initReq := mcp.InitializeRequest{}
  initReq.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
  initReq.Params.ClientInfo = mcp.Implementation{
    Name:    "Sidekick",
    Version: "1.0.0",
  }

  _, err = c.Initialize(context.Background(), initReq)
  if err != nil {
    return nil, fmt.Errorf("failed to initialize MCP client: %w", err)
  }

  return &mcpGoClientWrapper{client: c}, nil
}

// MCPClientFactory is a function that creates a new MCP client for a specific toolset
type MCPClientFactory func(toolset string) (MCPClientInterface, error)

// SetMCPClientFactory allows overriding the factory for testing
func SetMCPClientFactory(factory MCPClientFactory) {
  defaultMCPFactory = factory
}

// CallExternalTool creates a scoped MCP connection, calls the tool, and cleans up
func CallExternalTool(ctx context.Context, toolset string, toolName string,
  args map[string]interface{}) (string, error) {
  c, err := defaultMCPFactory(toolset)
  if err != nil {
    return "", fmt.Errorf("failed to create MCP client for %s: %w", toolset, err)
  }
  defer c.Close()

  return c.CallTool(ctx, toolName, args)
}