package mcp

import (
	"context"
	"errors"
	"net/http"

	"github.com/modelcontextprotocol/go-sdk/jsonrpc"

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

func validationError(message string) error {
	return runner.Errorf(
		runner.ErrorValidation,
		message,
		nil,
	)
}

func upstreamError(message string, cause error) error {
	return runner.Errorf(
		runner.ErrorUpstreamFailure,
		message,
		cause,
	)
}

func sanitizeError(ctx context.Context, err error, status int) error {
	if err == nil {
		return nil
	}
	if ctx != nil && ctx.Err() != nil {
		return upstreamError("MCP invocation did not complete", ctx.Err())
	}
	switch {
	case errors.Is(err, context.DeadlineExceeded):
		return upstreamError(
			"MCP invocation did not complete", context.DeadlineExceeded)
	case errors.Is(err, context.Canceled):
		return upstreamError(
			"MCP invocation did not complete", context.Canceled)
	}

	switch runner.CodeOf(err) {
	case runner.ErrorValidation:
		return validationError("MCP request failed local validation")
	case runner.ErrorPermissionDenied:
		return runner.Errorf(
			runner.ErrorPermissionDenied,
			"MCP request was not permitted",
			nil,
		)
	case runner.ErrorConnectionUnavailable,
		runner.ErrorProviderAuthUnavailable:
		return runner.Errorf(
			runner.ErrorProviderAuthUnavailable,
			"MCP connection credentials are unavailable",
			nil,
		)
	case runner.ErrorNotFound:
		return runner.Errorf(
			runner.ErrorNotFound,
			"MCP resource was not found",
			nil,
		)
	case runner.ErrorRateLimited:
		return runner.Errorf(
			runner.ErrorRateLimited,
			"MCP request was rate limited",
			nil,
		)
	case runner.ErrorUpstreamFailure:
		return upstreamError("MCP request failed", nil)
	}

	switch {
	case status == http.StatusUnauthorized ||
		status == http.StatusForbidden:
		return runner.Errorf(
			runner.ErrorProviderAuthUnavailable,
			"MCP server rejected the connection credentials",
			nil,
		)
	case status == http.StatusNotFound:
		return runner.Errorf(
			runner.ErrorNotFound,
			"MCP endpoint was not found",
			nil,
		)
	case status == http.StatusTooManyRequests:
		return runner.Errorf(
			runner.ErrorRateLimited,
			"MCP server rate limited the request",
			nil,
		)
	case status >= http.StatusBadRequest &&
		status < http.StatusInternalServerError:
		return validationError("MCP server rejected the request")
	}

	var rpcError *jsonrpc.Error
	if errors.As(err, &rpcError) && rpcError != nil {
		switch rpcError.Code {
		case jsonrpc.CodeMethodNotFound:
			return runner.Errorf(
				runner.ErrorNotFound,
				"MCP method was not found",
				nil,
			)
		case jsonrpc.CodeInvalidParams:
			return validationError("MCP method parameters were invalid")
		default:
			return upstreamError("MCP protocol request failed", nil)
		}
	}
	return upstreamError("MCP request failed", nil)
}