package mcp

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/url"
	"strings"
	"time"

	sdkmcp "github.com/modelcontextprotocol/go-sdk/mcp"

	"github.com/tidbcloud/AgentTeam/internal/extension/runner"
)

const (
	defaultCleanupTimeout = 5 * time.Second
	clientName            = "agentteam-extension-runtime"
	clientVersion         = "1.0.0"
)

var supportedProtocolVersions = map[string]struct{}{
	"2024-11-05": {},
	"2025-03-26": {},
	"2025-06-18": {},
	"2025-11-25": {},
	"2026-07-28": {},
}

// Options configures one reusable MCP endpoint definition. Every operation
// creates an invocation-local SDK client, session, authenticated transport,
// and byte budget; Client is therefore safe for concurrent use.
type Options struct {
	Endpoint string
	Auth     runner.AuthSpec

	// ProbeOnly creates a least-privilege client that can negotiate and clean
	// up a connection but has no tool/resource operation authority.
	ProbeOnly bool

	// AllowedTools and AllowedResourceURIs are constructor-time, code-defined
	// allowlists. At least one list must be non-empty unless ProbeOnly is true.
	// Runtime operation names must match one member exactly.
	AllowedTools        []string
	AllowedResourceURIs []string

	// Transport is a loopback test seam. Production leaves it nil and uses the
	// guarded provider dial path behind runtime.NewFixedProviderTransport.
	Transport http.RoundTripper

	// Timeout bounds the complete Connect + negotiation + operation lifecycle.
	// Zero selects runtime.DefaultRequestTimeout.
	Timeout time.Duration

	// CleanupTimeout independently bounds DELETE after the invocation context
	// has expired. Zero selects five seconds.
	CleanupTimeout time.Duration

	// MaxRequestBytes and MaxResponseBytes are cumulative across every
	// handshake and operation frame in one invocation. Zero selects the shared
	// 4 MiB ceiling.
	MaxRequestBytes  int64
	MaxResponseBytes int64
}

// Client is an immutable, concurrently reusable MCP endpoint definition.
type Client struct {
	endpoint            string
	auth                runner.AuthSpec
	probeOnly           bool
	allowedTools        map[string]struct{}
	allowedResourceURIs map[string]struct{}
	transport           http.RoundTripper
	timeout             time.Duration
	cleanupTimeout      time.Duration
	maxRequestBytes     int64
	maxResponseBytes    int64
}

// ToolResult is the SDK-independent result of one tools/call request.
type ToolResult struct {
	Content []json.RawMessage
	// StructuredContentPresent preserves the wire distinction between an
	// absent structuredContent field and an explicit JSON null. Providers may
	// use text fallback only for the absent case.
	StructuredContentPresent bool
	StructuredContent        json.RawMessage
}

// ResourceContent is one SDK-independent resources/read payload.
type ResourceContent struct {
	URI      string
	MIMEType string
	Text     string
	Blob     []byte
}

// ResourceResult is the SDK-independent result of one resources/read request.
type ResourceResult struct {
	Contents []ResourceContent
}

// NewClient validates and freezes one code-defined MCP endpoint.
func NewClient(options Options) (*Client, error) {
	if err := validateEndpoint(options.Endpoint); err != nil {
		return nil, validationError("MCP endpoint is invalid")
	}
	if err := options.Auth.ValidateSpec(); err != nil {
		return nil, validationError("MCP authentication contract is invalid")
	}
	allowedTools, err := buildAllowlist(options.AllowedTools)
	if err != nil {
		return nil, validationError("MCP tool allowlist is invalid")
	}
	allowedResources, err := buildAllowlist(options.AllowedResourceURIs)
	if err != nil {
		return nil, validationError("MCP resource allowlist is invalid")
	}
	if options.ProbeOnly &&
		(len(allowedTools) != 0 || len(allowedResources) != 0) {
		return nil, validationError(
			"MCP probe-only client cannot allow operations")
	}
	if !options.ProbeOnly &&
		len(allowedTools) == 0 &&
		len(allowedResources) == 0 {
		return nil, validationError("MCP client allows no operations")
	}
	timeout, err := boundedDuration(
		options.Timeout,
		runner.DefaultRequestTimeout,
	)
	if err != nil {
		return nil, validationError("MCP invocation timeout is invalid")
	}
	cleanupTimeout, err := boundedDuration(
		options.CleanupTimeout,
		defaultCleanupTimeout,
	)
	if err != nil {
		return nil, validationError("MCP cleanup timeout is invalid")
	}
	maxRequestBytes, err := boundedBytes(options.MaxRequestBytes)
	if err != nil {
		return nil, validationError("MCP request budget is invalid")
	}
	maxResponseBytes, err := boundedBytes(options.MaxResponseBytes)
	if err != nil {
		return nil, validationError("MCP response budget is invalid")
	}
	auth := options.Auth
	auth.Fields = append(
		[]runner.CredentialField(nil),
		options.Auth.Fields...,
	)
	return &Client{
		endpoint:            strings.TrimSpace(options.Endpoint),
		auth:                auth,
		probeOnly:           options.ProbeOnly,
		allowedTools:        allowedTools,
		allowedResourceURIs: allowedResources,
		transport:           options.Transport,
		timeout:             timeout,
		cleanupTimeout:      cleanupTimeout,
		maxRequestBytes:     maxRequestBytes,
		maxResponseBytes:    maxResponseBytes,
	}, nil
}

// CallTool opens one invocation-local session, calls one allowlisted tool, and
// closes the session before returning.
func (client *Client) CallTool(
	ctx context.Context,
	credentials runner.Credentials,
	tool string,
	arguments any,
) (ToolResult, error) {
	if client == nil {
		return ToolResult{}, upstreamError("MCP client is unavailable", nil)
	}
	if _, allowed := client.allowedTools[tool]; !allowed {
		return ToolResult{}, validationError("MCP tool is not allowed")
	}
	encodedArguments, err := encodeArguments(arguments)
	if err != nil {
		return ToolResult{}, validationError("MCP tool arguments are invalid")
	}
	return invoke(
		ctx,
		client,
		credentials,
		func(
			callCtx context.Context,
			session *sdkmcp.ClientSession,
		) (*sdkmcp.CallToolResult, error) {
			return session.CallTool(callCtx, &sdkmcp.CallToolParams{
				Name:      tool,
				Arguments: encodedArguments,
			})
		},
		convertToolResult,
	)
}

// ReadResource opens one invocation-local session, reads one allowlisted URI,
// and closes the session before returning.
func (client *Client) ReadResource(
	ctx context.Context,
	credentials runner.Credentials,
	uri string,
) (ResourceResult, error) {
	if client == nil {
		return ResourceResult{}, upstreamError(
			"MCP client is unavailable", nil)
	}
	if _, allowed := client.allowedResourceURIs[uri]; !allowed {
		return ResourceResult{}, validationError(
			"MCP resource URI is not allowed")
	}
	return invoke(
		ctx,
		client,
		credentials,
		func(
			callCtx context.Context,
			session *sdkmcp.ClientSession,
		) (*sdkmcp.ReadResourceResult, error) {
			return session.ReadResource(
				callCtx,
				&sdkmcp.ReadResourceParams{URI: uri},
			)
		},
		convertResourceResult,
	)
}

// Probe opens one invocation-local session, validates endpoint/auth/protocol
// negotiation, and closes it without listing or invoking any server
// capability. It is intended only for bounded connection validation.
func (client *Client) Probe(
	ctx context.Context,
	credentials runner.Credentials,
) error {
	if client == nil {
		return upstreamError("MCP client is unavailable", nil)
	}
	if !client.probeOnly {
		return validationError("MCP client is not probe-only")
	}
	_, err := invoke(
		ctx,
		client,
		credentials,
		func(
			context.Context,
			*sdkmcp.ClientSession,
		) (struct{}, error) {
			return struct{}{}, nil
		},
		func(struct{}, *invocationTransport) (struct{}, error) {
			return struct{}{}, nil
		},
	)
	return err
}

func invoke[SDKResult any, Result any](
	ctx context.Context,
	client *Client,
	credentials runner.Credentials,
	call func(context.Context, *sdkmcp.ClientSession) (SDKResult, error),
	convert func(SDKResult, *invocationTransport) (Result, error),
) (Result, error) {
	var zero Result
	if ctx == nil {
		return zero, validationError("MCP invocation context is unavailable")
	}
	callCtx, cancel := context.WithTimeout(ctx, client.timeout)
	defer cancel()

	invocation, err := newInvocationTransport(
		client,
		credentials,
		callCtx,
	)
	if err != nil {
		return zero, sanitizeError(callCtx, err, 0)
	}
	defer invocation.closeIdleConnections()
	sdkClient := sdkmcp.NewClient(
		&sdkmcp.Implementation{
			Name:    clientName,
			Version: clientVersion,
		},
		&sdkmcp.ClientOptions{
			Capabilities: &sdkmcp.ClientCapabilities{},
			MultiRoundTrip: &sdkmcp.MultiRoundTripOptions{
				Disabled: true,
			},
		},
	)
	transport := &sdkmcp.StreamableClientTransport{
		Endpoint:             client.endpoint,
		HTTPClient:           invocation.httpClient(),
		MaxRetries:           -1,
		DisableStandaloneSSE: true,
	}
	session, err := sdkClient.Connect(callCtx, transport, nil)
	if err != nil {
		status := invocation.operationStatus()
		_ = invocation.cleanup()
		return zero, sanitizeError(callCtx, err, status)
	}

	initializeResult := session.InitializeResult()
	if initializeResult == nil {
		_ = session.Close()
		_ = invocation.cleanup()
		return zero, upstreamError(
			"MCP server did not negotiate a protocol version", nil)
	}
	if _, supported := supportedProtocolVersions[initializeResult.ProtocolVersion]; !supported {
		_ = session.Close()
		_ = invocation.cleanup()
		return zero, upstreamError(
			"MCP server negotiated an unsupported protocol version", nil)
	}

	sdkResult, callErr := call(callCtx, session)
	callStatus := invocation.operationStatus()
	// Cleanup is mandatory but best-effort after the server has already
	// completed the requested operation. A failed DELETE must not turn a
	// successfully executed mutation into an apparent failure that callers may
	// retry. Connect and operation failures still retain their primary error.
	_ = session.Close()
	_ = invocation.cleanup()
	if callErr != nil {
		return zero, sanitizeError(callCtx, callErr, callStatus)
	}
	return convert(sdkResult, invocation)
}

func convertToolResult(
	result *sdkmcp.CallToolResult,
	invocation *invocationTransport,
) (ToolResult, error) {
	if result == nil {
		return ToolResult{}, upstreamError(
			"MCP tool returned no result", nil)
	}
	if result.NeedsInput() {
		return ToolResult{}, upstreamError(
			"MCP tool requires an unsupported additional round trip", nil)
	}
	if result.IsError {
		return ToolResult{}, upstreamError(
			"MCP tool reported an upstream failure", nil)
	}
	converted := ToolResult{
		Content: make([]json.RawMessage, 0, len(result.Content)),
	}
	if invocation != nil {
		converted.StructuredContentPresent =
			invocation.toolStructuredContentPresent()
	}
	for _, content := range result.Content {
		encoded, err := json.Marshal(content)
		if err != nil {
			return ToolResult{}, upstreamError(
				"MCP tool content was not encodable", nil)
		}
		converted.Content = append(
			converted.Content,
			json.RawMessage(encoded),
		)
	}
	if result.StructuredContent != nil {
		encoded, err := json.Marshal(result.StructuredContent)
		if err != nil {
			return ToolResult{}, upstreamError(
				"MCP structured content was not encodable", nil)
		}
		converted.StructuredContent = json.RawMessage(encoded)
	}
	return converted, nil
}

func convertResourceResult(
	result *sdkmcp.ReadResourceResult,
	_ *invocationTransport,
) (ResourceResult, error) {
	if result == nil {
		return ResourceResult{}, upstreamError(
			"MCP resource returned no result", nil)
	}
	if result.NeedsInput() {
		return ResourceResult{}, upstreamError(
			"MCP resource requires an unsupported additional round trip", nil)
	}
	converted := ResourceResult{
		Contents: make([]ResourceContent, 0, len(result.Contents)),
	}
	for _, content := range result.Contents {
		if content == nil {
			return ResourceResult{}, upstreamError(
				"MCP resource returned invalid content", nil)
		}
		converted.Contents = append(converted.Contents, ResourceContent{
			URI:      content.URI,
			MIMEType: content.MIMEType,
			Text:     content.Text,
			Blob:     append([]byte(nil), content.Blob...),
		})
	}
	return converted, nil
}

func validateEndpoint(raw string) error {
	endpoint, err := url.Parse(strings.TrimSpace(raw))
	if err != nil ||
		endpoint.Scheme != "https" ||
		endpoint.Host == "" ||
		endpoint.User != nil ||
		endpoint.Opaque != "" ||
		endpoint.RawQuery != "" ||
		endpoint.ForceQuery ||
		endpoint.Fragment != "" ||
		endpoint.RawFragment != "" {
		return validationError("MCP endpoint is invalid")
	}
	return nil
}

func buildAllowlist(values []string) (map[string]struct{}, error) {
	result := make(map[string]struct{}, len(values))
	for _, value := range values {
		if value == "" || strings.TrimSpace(value) != value {
			return nil, validationError("MCP allowlist member is invalid")
		}
		if _, duplicate := result[value]; duplicate {
			return nil, validationError("MCP allowlist member is duplicated")
		}
		result[value] = struct{}{}
	}
	return result, nil
}

func boundedDuration(value, fallback time.Duration) (time.Duration, error) {
	if value < 0 {
		return 0, validationError("MCP duration is invalid")
	}
	if value == 0 {
		return fallback, nil
	}
	return value, nil
}

func boundedBytes(value int64) (int64, error) {
	if value < 0 || value > runner.MaxResponseBytes {
		return 0, validationError("MCP byte budget is invalid")
	}
	if value == 0 {
		return runner.MaxResponseBytes, nil
	}
	return value, nil
}

func encodeArguments(arguments any) (json.RawMessage, error) {
	if arguments == nil {
		return json.RawMessage(`{}`), nil
	}
	encoded, err := json.Marshal(arguments)
	if err != nil {
		return nil, err
	}
	trimmed := bytes.TrimSpace(encoded)
	if len(trimmed) < 2 || trimmed[0] != '{' {
		return nil, validationError(
			"MCP tool arguments must encode as an object")
	}
	return json.RawMessage(trimmed), nil
}