package backend

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"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"
	"go.uber.org/zap"
)

var MCPClientInitTimeout = 1 * time.Minute

type ProxyBackend struct {
	logger *zap.Logger
	cmd    []string
	errCh  chan error
	client *client.Client
}

func NewProxyBackend(logger *zap.Logger, cmd []string) Backend {
	return &ProxyBackend{
		logger: logger,
		cmd:    cmd,
		errCh:  nil,
	}
}

func (p *ProxyBackend) Run(ctx context.Context) (http.Handler, error) {
	if p.errCh != nil {
		return nil, fmt.Errorf("proxy backend is already running")
	}
	execCmd := exec.CommandContext(ctx, p.cmd[0], p.cmd[1:]...)
	stdout, err := execCmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("failed to get stdout pipe: %w", err)
	}
	stdin, err := execCmd.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("failed to get stdin pipe: %w", err)
	}
	execCmd.Stderr = os.Stderr
	if err := execCmd.Start(); err != nil {
		return nil, fmt.Errorf("failed to start exec command: %w", err)
	}
	p.errCh = make(chan error)
	done := make(chan struct{})
	go func() {
		if err := execCmd.Wait(); err != nil {
			p.errCh <- fmt.Errorf("exec command failed: %s: %w", strings.Join(p.cmd, " "), err)
		}
	}()
	var handler http.Handler
	go func() {
		tr := transport.NewIO(stdout, stdin, os.Stderr)
		transport.WithCommandLogger(&zapLogger{p.logger})(tr)
		_client, _handler, err := setupProxy(ctx, p.logger, tr)
		if err != nil {
			p.errCh <- err
			return
		}
		p.client = _client
		handler = _handler
		done <- struct{}{}
	}()
	select {
	case err := <-p.errCh:
		return nil, err
	case <-done:
		return handler, nil
	}
}

func (p *ProxyBackend) Wait() error {
	if p.errCh == nil {
		return nil
	}
	return <-p.errCh
}

func (p *ProxyBackend) Close() error {
	if p.client != nil {
		return p.client.Close()
	}
	return nil
}

func setupProxy(ctx context.Context, logger *zap.Logger, tr transport.BidirectionalInterface) (*client.Client, http.Handler, error) {
	c := client.NewClient(tr)
	if err := c.Start(ctx); err != nil {
		return nil, nil, fmt.Errorf("failed to start MCP client: %w", err)
	}
	initCtx, cancel := context.WithTimeout(ctx, MCPClientInitTimeout)
	defer cancel()

	initRequest := mcp.InitializeRequest{
		Params: mcp.InitializeParams{
			ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION,
			ClientInfo: mcp.Implementation{
				Name:    "mcp-auth-proxy",
				Version: "dev",
			},
		},
	}
	init, err := c.Initialize(initCtx, initRequest)
	if err != nil {
		c.Close()
		return nil, nil, fmt.Errorf("failed to initialize MCP client: %w", err)
	}
	s := server.NewMCPServer(init.ServerInfo.Name, init.ServerInfo.Version)
	if init.Capabilities.Tools != nil {
		tools, err := c.ListTools(ctx, mcp.ListToolsRequest{})
		if err != nil {
			c.Close()
			return nil, nil, fmt.Errorf("failed to list tools: %w", err)
		}
		for _, tool := range tools.Tools {
			s.AddTool(tool, c.CallTool)
		}
	}
	if init.Capabilities.Prompts != nil {
		prompts, err := c.ListPrompts(ctx, mcp.ListPromptsRequest{})
		if err != nil {
			c.Close()
			return nil, nil, fmt.Errorf("failed to list prompts: %w", err)
		}
		for _, prompt := range prompts.Prompts {
			s.AddPrompt(prompt, c.GetPrompt)
		}
	}
	if init.Capabilities.Resources != nil {
		resources, err := c.ListResources(ctx, mcp.ListResourcesRequest{})
		if err != nil {
			c.Close()
			return nil, nil, fmt.Errorf("failed to list resources: %w", err)
		}
		for _, resource := range resources.Resources {
			s.AddResource(resource, func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
				res, err := c.ReadResource(ctx, request)
				if err != nil {
					return nil, err
				}
				return res.Contents, nil
			})
		}
	}
	ss := server.NewStreamableHTTPServer(
		s,
		server.WithEndpointPath("/mcp"),
		server.WithStateLess(true),
		server.WithLogger(&zapLogger{logger}),
	)
	return c, ss, nil
}

type zapLogger struct {
	logger *zap.Logger
}

func (l *zapLogger) Infof(format string, v ...any) {
	l.logger.Info(fmt.Sprintf(format, v...))
}

func (l *zapLogger) Errorf(format string, v ...any) {
	l.logger.Error(fmt.Sprintf(format, v...))
}