package mcp

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"

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

const (
	maxCleanupRequestBytes  = 4 << 10
	maxCleanupResponseBytes = 64 << 10
	maxOperationRequests    = 8
)

var (
	errResponseBudgetExceeded = errors.New("MCP response byte budget exceeded")
	errCleanupBudgetExceeded  = errors.New("MCP cleanup byte budget exceeded")
	errRedirectRejected       = errors.New("MCP redirect was rejected")
	errCleanupRejected        = errors.New("MCP cleanup was rejected")
)

type invocationTransport struct {
	endpoint       string
	base           http.RoundTripper
	invocationCtx  context.Context
	cleanupTimeout time.Duration
	requestBudget  byteBudget
	responseBudget byteBudget

	mu struct {
		sync.Mutex
		sessionID         string
		protocolVersion   string
		operationStatus   int
		operationRequests int
		cleanupAttempted  bool
		// Defaults to true for tools/call and changes to false only after the
		// complete wire response proves structuredContent was absent.
		toolStructuredContentPresent bool
	}
}

func newInvocationTransport(
	client *Client,
	credentials runner.Credentials,
	invocationCtx context.Context,
) (*invocationTransport, error) {
	base, err := runner.NewFixedProviderTransport(
		runner.FixedProviderTransportOptions{
			Endpoint:    client.endpoint,
			Auth:        client.auth,
			Credentials: credentials,
			Transport:   client.transport,
		},
	)
	if err != nil {
		return nil, err
	}
	invocation := &invocationTransport{
		endpoint:       client.endpoint,
		base:           base,
		invocationCtx:  invocationCtx,
		cleanupTimeout: client.cleanupTimeout,
		requestBudget:  byteBudget{max: client.maxRequestBytes},
		responseBudget: byteBudget{max: client.maxResponseBytes},
	}
	return invocation, nil
}

func (transport *invocationTransport) httpClient() *http.Client {
	return &http.Client{
		Transport: transport,
		CheckRedirect: func(*http.Request, []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}
}

func (transport *invocationTransport) RoundTrip(
	request *http.Request,
) (*http.Response, error) {
	if request == nil {
		return nil, validationError("MCP HTTP request is unavailable")
	}
	switch request.Method {
	case http.MethodPost:
		return transport.roundTripOperation(request)
	case http.MethodDelete:
		return transport.roundTripCleanup(request)
	default:
		return nil, validationError("MCP HTTP method is not allowed")
	}
}

func (transport *invocationTransport) roundTripOperation(
	request *http.Request,
) (*http.Response, error) {
	request, releaseContext, err := transport.bindInvocationContext(request)
	if err != nil {
		return nil, err
	}
	releaseOnReturn := true
	defer func() {
		if releaseOnReturn {
			releaseContext()
		}
	}()

	transport.mu.Lock()
	transport.mu.operationRequests++
	tooManyRequests := transport.mu.operationRequests > maxOperationRequests
	transport.mu.Unlock()
	if tooManyRequests {
		return nil, validationError(
			"MCP invocation exceeded its request count budget")
	}
	cloned, err := transport.cloneBudgetedRequest(
		request,
		&transport.requestBudget,
		validationError("MCP invocation exceeded its request byte budget"),
	)
	if err != nil {
		return nil, err
	}
	rpcMethod, rpcID := requestRPCIdentity(cloned)
	if rpcMethod == "tools/call" {
		transport.mu.Lock()
		transport.mu.toolStructuredContentPresent = true
		transport.mu.Unlock()
	}
	transport.mu.Lock()
	// The legacy initialize fallback intentionally omits the modern protocol
	// header. Replace rather than conditionally update this snapshot so a
	// failed initialize cannot send a stale discover version on DELETE.
	transport.mu.protocolVersion = cloned.Header.Get("Mcp-Protocol-Version")
	transport.mu.Unlock()
	response, err := transport.base.RoundTrip(cloned)
	if err != nil {
		return nil, err
	}
	transport.mu.Lock()
	transport.mu.operationStatus = response.StatusCode
	if version := response.Header.Get("Mcp-Protocol-Version"); version != "" {
		transport.mu.protocolVersion = version
	}
	if sessionID := response.Header.Get("Mcp-Session-Id"); sessionID != "" && transport.mu.sessionID == "" {
		transport.mu.sessionID = sessionID
	}
	transport.mu.Unlock()
	if response.StatusCode >= http.StatusMultipleChoices &&
		response.StatusCode < http.StatusBadRequest {
		_ = response.Body.Close()
		return nil, errRedirectRejected
	}
	if remaining := transport.responseBudget.remaining(); response.ContentLength > remaining {
		_ = response.Body.Close()
		return nil, errResponseBudgetExceeded
	}
	budgetedBody := &budgetReadCloser{
		ReadCloser: &releaseReadCloser{
			ReadCloser: response.Body,
			release:    releaseContext,
		},
		budget: &transport.responseBudget,
		err:    errResponseBudgetExceeded,
	}
	if rpcMethod == "tools/call" {
		response.Body = &captureReadCloser{
			ReadCloser: budgetedBody,
			onComplete: func(payload []byte) {
				transport.observeToolResponse(payload, rpcID)
			},
		}
	} else {
		response.Body = budgetedBody
	}
	releaseOnReturn = false
	return response, nil
}

func (transport *invocationTransport) bindInvocationContext(
	request *http.Request,
) (*http.Request, func(), error) {
	if err := transport.invocationCtx.Err(); err != nil {
		return nil, nil, err
	}
	requestCtx, cancel := context.WithCancel(request.Context())
	stop := context.AfterFunc(transport.invocationCtx, cancel)
	var once sync.Once
	release := func() {
		once.Do(func() {
			stop()
			cancel()
		})
	}
	return request.Clone(requestCtx), release, nil
}

func (transport *invocationTransport) roundTripCleanup(
	request *http.Request,
) (*http.Response, error) {
	transport.mu.Lock()
	transport.mu.cleanupAttempted = true
	if sessionID := request.Header.Get("Mcp-Session-Id"); sessionID != "" {
		transport.mu.sessionID = sessionID
	}
	transport.mu.Unlock()

	cleanupCtx, cancel := context.WithTimeout(
		context.Background(),
		transport.cleanupTimeout,
	)
	cloned := request.Clone(cleanupCtx)
	cloned.Header = request.Header.Clone()
	budget := byteBudget{max: maxCleanupRequestBytes}
	cloned, err := transport.cloneBudgetedRequest(
		cloned,
		&budget,
		errCleanupBudgetExceeded,
	)
	if err != nil {
		cancel()
		return nil, err
	}
	response, err := transport.base.RoundTrip(cloned)
	if err != nil {
		cancel()
		return nil, err
	}
	if response.StatusCode >= http.StatusMultipleChoices &&
		response.StatusCode < http.StatusBadRequest {
		_ = response.Body.Close()
		cancel()
		return nil, errRedirectRejected
	}
	if !cleanupStatusAccepted(response.StatusCode) {
		_ = response.Body.Close()
		cancel()
		return nil, errCleanupRejected
	}
	response.Body = &cancelReadCloser{
		ReadCloser: &budgetReadCloser{
			ReadCloser: response.Body,
			budget: &byteBudget{
				max: maxCleanupResponseBytes,
			},
			err: errCleanupBudgetExceeded,
		},
		cancel: cancel,
	}
	return response, nil
}

func (transport *invocationTransport) cleanup() error {
	transport.mu.Lock()
	if transport.mu.cleanupAttempted || transport.mu.sessionID == "" {
		transport.mu.Unlock()
		return nil
	}
	transport.mu.cleanupAttempted = true
	sessionID := transport.mu.sessionID
	protocolVersion := transport.mu.protocolVersion
	transport.mu.Unlock()

	cleanupCtx, cancel := context.WithTimeout(
		context.Background(),
		transport.cleanupTimeout,
	)
	defer cancel()
	request, err := http.NewRequestWithContext(
		cleanupCtx,
		http.MethodDelete,
		transport.endpoint,
		nil,
	)
	if err != nil {
		return errCleanupBudgetExceeded
	}
	request.Header.Set("Mcp-Session-Id", sessionID)
	if protocolVersion != "" {
		request.Header.Set("Mcp-Protocol-Version", protocolVersion)
	}
	response, err := transport.base.RoundTrip(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	if response.StatusCode >= http.StatusMultipleChoices &&
		response.StatusCode < http.StatusBadRequest {
		return errRedirectRejected
	}
	if !cleanupStatusAccepted(response.StatusCode) {
		return errCleanupRejected
	}
	read, err := io.Copy(
		io.Discard,
		io.LimitReader(response.Body, maxCleanupResponseBytes+1),
	)
	if err != nil {
		return err
	}
	if read > maxCleanupResponseBytes {
		return errCleanupBudgetExceeded
	}
	return nil
}

func (transport *invocationTransport) cloneBudgetedRequest(
	request *http.Request,
	budget *byteBudget,
	budgetErr error,
) (*http.Request, error) {
	cloned := request.Clone(request.Context())
	cloned.Header = request.Header.Clone()
	if request.Body == nil {
		return cloned, nil
	}
	remaining := budget.remaining()
	if request.ContentLength > remaining {
		_ = request.Body.Close()
		return nil, budgetErr
	}
	body, err := io.ReadAll(io.LimitReader(request.Body, remaining+1))
	_ = request.Body.Close()
	if err != nil {
		return nil, err
	}
	if int64(len(body)) > remaining || !budget.charge(int64(len(body))) {
		return nil, budgetErr
	}
	cloned.Body = io.NopCloser(bytes.NewReader(body))
	cloned.ContentLength = int64(len(body))
	cloned.GetBody = nil
	return cloned, nil
}

func (transport *invocationTransport) operationStatus() int {
	transport.mu.Lock()
	defer transport.mu.Unlock()
	return transport.mu.operationStatus
}

func (transport *invocationTransport) toolStructuredContentPresent() bool {
	transport.mu.Lock()
	defer transport.mu.Unlock()
	return transport.mu.toolStructuredContentPresent
}

func (transport *invocationTransport) observeToolResponse(
	payload []byte,
	rpcID json.RawMessage,
) {
	present, ok := toolStructuredContentPresence(payload, rpcID)
	if !ok {
		return
	}
	transport.mu.Lock()
	transport.mu.toolStructuredContentPresent = present
	transport.mu.Unlock()
}

func (transport *invocationTransport) closeIdleConnections() {
	if closer, ok := transport.base.(interface{ CloseIdleConnections() }); ok {
		closer.CloseIdleConnections()
	}
}

func requestRPCIdentity(
	request *http.Request,
) (string, json.RawMessage) {
	if request == nil || request.Body == nil {
		return "", nil
	}
	body, err := io.ReadAll(request.Body)
	_ = request.Body.Close()
	request.Body = io.NopCloser(bytes.NewReader(body))
	request.ContentLength = int64(len(body))
	if err != nil {
		return "", nil
	}
	var envelope struct {
		ID     json.RawMessage `json:"id"`
		Method string          `json:"method"`
	}
	if err := json.Unmarshal(body, &envelope); err != nil {
		return "", nil
	}
	rawID := bytes.TrimSpace(envelope.ID)
	if len(rawID) == 0 {
		return envelope.Method, nil
	}
	first := rawID[0]
	if first != '"' && first != '-' && (first < '0' || first > '9') {
		return envelope.Method, nil
	}
	if !json.Valid(rawID) {
		return envelope.Method, nil
	}
	return envelope.Method, append(json.RawMessage(nil), rawID...)
}

func toolStructuredContentPresence(
	payload []byte,
	expectedID json.RawMessage,
) (bool, bool) {
	if len(expectedID) == 0 {
		return false, false
	}
	if present, matched, valid := structuredContentPresenceInRPC(
		payload,
		expectedID,
	); valid && matched {
		return present, true
	}
	normalized := bytes.ReplaceAll(payload, []byte("\r\n"), []byte("\n"))
	matches := 0
	present := false
	for _, event := range bytes.Split(normalized, []byte("\n\n")) {
		var data []string
		for _, line := range strings.Split(string(event), "\n") {
			if strings.HasPrefix(line, "data:") {
				data = append(data, strings.TrimSpace(
					strings.TrimPrefix(line, "data:"),
				))
			}
		}
		if len(data) == 0 {
			continue
		}
		eventPresent, matched, valid := structuredContentPresenceInRPC(
			[]byte(strings.Join(data, "\n")),
			expectedID,
		)
		if !valid {
			return false, false
		}
		if !matched {
			continue
		}
		matches++
		if matches > 1 {
			return false, false
		}
		present = eventPresent
	}
	return present, matches == 1
}

func structuredContentPresenceInRPC(
	payload []byte,
	expectedID json.RawMessage,
) (bool, bool, bool) {
	if _, _, err := runner.DecodeUniqueJSONDocument(
		bytes.TrimSpace(payload),
	); err != nil {
		return false, false, false
	}
	var envelope map[string]json.RawMessage
	if err := json.Unmarshal(bytes.TrimSpace(payload), &envelope); err != nil {
		return false, false, false
	}
	var version string
	if err := json.Unmarshal(envelope["jsonrpc"], &version); err != nil ||
		version != "2.0" {
		return false, false, false
	}
	rawID, hasID := envelope["id"]
	if !hasID {
		if _, notification := envelope["method"]; notification {
			return false, false, true
		}
		return false, false, false
	}
	rawResult, exists := envelope["result"]
	if !exists {
		if _, isError := envelope["error"]; isError {
			return false, false, true
		}
		return false, false, false
	}
	var result map[string]json.RawMessage
	if err := json.Unmarshal(rawResult, &result); err != nil || result == nil {
		return false, false, false
	}
	if !bytes.Equal(bytes.TrimSpace(rawID), expectedID) {
		return false, false, true
	}
	_, present := result["structuredContent"]
	return present, true, true
}

func cleanupStatusAccepted(status int) bool {
	return (status >= http.StatusOK &&
		status < http.StatusMultipleChoices) ||
		status == http.StatusNotFound ||
		status == http.StatusGone
}

type byteBudget struct {
	mu   sync.Mutex
	max  int64
	used int64
}

func (budget *byteBudget) remaining() int64 {
	budget.mu.Lock()
	defer budget.mu.Unlock()
	return max(budget.max-budget.used, 0)
}

func (budget *byteBudget) charge(amount int64) bool {
	budget.mu.Lock()
	defer budget.mu.Unlock()
	if amount < 0 || amount > budget.max-budget.used {
		return false
	}
	budget.used += amount
	return true
}

type budgetReadCloser struct {
	io.ReadCloser
	budget *byteBudget
	err    error
}

func (reader *budgetReadCloser) Read(buffer []byte) (int, error) {
	if len(buffer) == 0 {
		return 0, nil
	}
	remaining := reader.budget.remaining()
	limit := int64(len(buffer))
	if limit > remaining+1 {
		limit = remaining + 1
	}
	n, err := reader.ReadCloser.Read(buffer[:int(limit)])
	if int64(n) > remaining {
		if remaining > 0 {
			_ = reader.budget.charge(remaining)
			return int(remaining), reader.err
		}
		return 0, reader.err
	}
	if n > 0 && !reader.budget.charge(int64(n)) {
		return 0, reader.err
	}
	return n, err
}

type cancelReadCloser struct {
	io.ReadCloser
	cancel context.CancelFunc
}

type captureReadCloser struct {
	io.ReadCloser
	mu         sync.Mutex
	payload    bytes.Buffer
	onComplete func([]byte)
	complete   bool
}

func (reader *captureReadCloser) Read(buffer []byte) (int, error) {
	n, err := reader.ReadCloser.Read(buffer)
	if n > 0 {
		reader.mu.Lock()
		_, _ = reader.payload.Write(buffer[:n])
		reader.mu.Unlock()
	}
	if err != nil {
		reader.finish()
	}
	return n, err
}

func (reader *captureReadCloser) Close() error {
	reader.finish()
	return reader.ReadCloser.Close()
}

func (reader *captureReadCloser) finish() {
	reader.mu.Lock()
	if reader.complete {
		reader.mu.Unlock()
		return
	}
	reader.complete = true
	payload := append([]byte(nil), reader.payload.Bytes()...)
	reader.payload.Reset()
	onComplete := reader.onComplete
	reader.mu.Unlock()
	if onComplete != nil {
		onComplete(payload)
	}
}

type releaseReadCloser struct {
	io.ReadCloser
	once    sync.Once
	release func()
}

func (reader *releaseReadCloser) Read(buffer []byte) (int, error) {
	n, err := reader.ReadCloser.Read(buffer)
	if err != nil {
		reader.once.Do(reader.release)
	}
	return n, err
}

func (reader *releaseReadCloser) Close() error {
	defer reader.once.Do(reader.release)
	return reader.ReadCloser.Close()
}

func (reader *cancelReadCloser) Close() error {
	err := reader.ReadCloser.Close()
	reader.cancel()
	return err
}