package mcp

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

	"github.com/gin-gonic/gin"
	"github.com/tidbcloud/AgentTeam/desktop/internal/resultbudget"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

// rpcRequest represents a JSON-RPC 2.0 request.
type rpcRequest struct {
	JSONRPC string           `json:"jsonrpc"`
	ID      *json.RawMessage `json:"id,omitempty"`
	Method  string           `json:"method"`
	Params  json.RawMessage  `json:"params,omitempty"`
}

// rpcResponse represents a JSON-RPC 2.0 response.
type rpcResponse struct {
	JSONRPC string           `json:"jsonrpc"`
	ID      *json.RawMessage `json:"id"`
	Result  interface{}      `json:"result,omitempty"`
	Error   *rpcError        `json:"error,omitempty"`

	deliveryReceipts []string
	deliveryOwner    types.MCPToolResultDeliveryOwner
}

// rpcError represents a JSON-RPC 2.0 error object.
type rpcError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

// ---------------------------------------------------------------------------
// Tool result helpers
// ---------------------------------------------------------------------------

// toolText returns a standard MCP tool result containing text content.
func toolText(text string) *types.MCPToolResult {
	return &types.MCPToolResult{
		Content: []types.MCPToolContent{{Type: "text", Text: text}},
	}
}

// toolTextStructured returns a standard text result plus structured content
// for MCP clients that support non-prose result payloads.
func toolTextStructured(text string, structured any) *toolResult {
	return &toolResult{
		Content:           []toolContent{{Type: "text", Text: text}},
		StructuredContent: structured,
	}
}

// toolError returns an MCP tool result indicating an error.
func toolError(msg string) *types.MCPToolResult {
	return &types.MCPToolResult{
		IsError: true,
		Content: []types.MCPToolContent{{Type: "text", Text: "Error: " + msg}},
	}
}

// parseErrorResponse returns a JSON-RPC parse error response.
func parseErrorResponse(message string) *rpcResponse {
	null := json.RawMessage("null")
	return &rpcResponse{
		JSONRPC: "2.0",
		ID:      &null,
		Error:   &rpcError{Code: -32700, Message: message},
	}
}

// invalidRequestResponse returns a JSON-RPC invalid request response.
func invalidRequestResponse(message string) *rpcResponse {
	null := json.RawMessage("null")
	return &rpcResponse{
		JSONRPC: "2.0",
		ID:      &null,
		Error:   &rpcError{Code: -32600, Message: message},
	}
}

// ---------------------------------------------------------------------------
// HTTP helpers
// ---------------------------------------------------------------------------

const (
	oauthAuthorizationServerPath = "/.well-known/oauth-authorization-server"
	openIDConfigurationPath      = "/.well-known/openid-configuration"
	oauthProtectedResourcePath   = "/.well-known/oauth-protected-resource"
)

// writeMCPJSON writes JSON without HTML escaping. MCP protocol payloads
// contain raw text (markdown, code snippets) that must not be mangled by
// HTML entity encoding.
func writeMCPJSON(c *gin.Context, status int, payload interface{}) error {
	var encoded bytes.Buffer
	encoder := json.NewEncoder(&encoded)
	encoder.SetEscapeHTML(false)
	if err := encoder.Encode(payload); err != nil {
		return err
	}
	c.Header("Content-Type", "application/json")
	c.Status(status)
	n, err := c.Writer.Write(encoded.Bytes())
	if err != nil {
		return err
	}
	if n != encoded.Len() {
		return fmt.Errorf("%w: wrote %d of %d MCP response bytes", io.ErrShortWrite, n, encoded.Len())
	}
	return nil
}

func writeMCPStatus(c *gin.Context, status int) {
	c.Header("Content-Type", "application/json")
	c.Status(status)
	c.Writer.WriteHeaderNow()
}

func writeWellKnownUnsupported(c *gin.Context, resource string) {
	_ = writeMCPJSON(c, http.StatusNotFound, gin.H{
		"error":   "not_found",
		"message": fmt.Sprintf("%s discovery is not supported for agentteam MCP endpoints", resource),
	})
}

const maxMCPRequestBodyBytes = 4 << 20

// errRequestBodyTooLarge reports that an MCP HTTP body exceeded the protocol limit.
var errRequestBodyTooLarge = errors.New("mcp request body too large")

// readBody reads the request body and rejects bodies larger than 4 MiB.
func readBody(c *gin.Context) ([]byte, error) {
	body, err := io.ReadAll(io.LimitReader(c.Request.Body, maxMCPRequestBodyBytes+1))
	if err != nil {
		return nil, err
	}
	if len(body) > maxMCPRequestBodyBytes {
		return nil, errRequestBodyTooLarge
	}
	return body, nil
}

// agentID extracts the agentID set by AuthMiddleware.
func agentID(c *gin.Context) string {
	v, _ := c.Get("agentID")
	s, _ := v.(string)
	return s
}

// mcpSessionID returns the MCP session ID for a given agent.
func mcpSessionID(id string) string {
	return "agentteam-" + strings.TrimSpace(id)
}

func newStreamHandler() gin.HandlerFunc {
	return func(c *gin.Context) {
		accept := strings.ToLower(strings.TrimSpace(c.GetHeader("Accept")))
		if accept != "" && !strings.Contains(accept, "text/event-stream") && !strings.Contains(accept, "*/*") {
			_ = writeMCPJSON(c, http.StatusOK, gin.H{
				"name":    "agentteam",
				"session": mcpSessionID(agentID(c)),
			})
			return
		}

		flusher := c.Writer
		c.Header("Content-Type", "text/event-stream")
		c.Header("Cache-Control", "no-cache")
		c.Header("Connection", "keep-alive")
		c.Status(http.StatusOK)
		if _, err := io.WriteString(c.Writer, ": ready\n\n"); err != nil {
			return
		}
		flusher.Flush()

		ticker := time.NewTicker(25 * time.Second)
		defer ticker.Stop()

		for {
			select {
			case <-c.Request.Context().Done():
				return
			case <-ticker.C:
				if _, err := io.WriteString(c.Writer, ": keepalive\n\n"); err != nil {
					return
				}
				flusher.Flush()
			}
		}
	}
}

// ---------------------------------------------------------------------------
// DisconnectHandler — DELETE (204 No Content)
// ---------------------------------------------------------------------------

// newDisconnectHandler returns a gin handler that responds with 204 No Content.
func newDisconnectHandler() gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Status(http.StatusNoContent)
	}
}

// dispatchRPC processes a single JSON-RPC request and returns the response.
func dispatchRPC(ctx context.Context, id string, h types.MCPHandler, data []byte) (*rpcResponse, bool) {
	var req rpcRequest
	if err := json.Unmarshal(data, &req); err != nil {
		return parseErrorResponse("parse error"), true
	}
	if req.ID == nil {
		return nil, false
	}

	var result interface{}
	var rpcErr *rpcError
	var deliveryReceipts []string
	var deliveryOwner types.MCPToolResultDeliveryOwner

	switch req.Method {
	case "initialize":
		result = types.MCPInitializeResult{
			ProtocolVersion: "2024-11-05",
			Capabilities: types.MCPServerCapabilities{
				Tools: types.MCPToolCapabilities{ListChanged: true},
			},
			ServerInfo: types.MCPServerInfo{
				Name:    "agentteam",
				Version: "1.0.0",
			},
		}
	case "ping":
		result = struct{}{}
	case "tools/list":
		defs := h.Tools()
		for i := range defs {
			// Some tools are not created via NewToolDef, leaving Properties nil.
			// A nil Properties serializes to JSON null, which causes MCP clients
			// to fail to recognise the tool.
			if defs[i].InputSchema.Properties == nil {
				defs[i].InputSchema.Properties = map[string]types.MCPProperty{}
			}
		}
		result = types.MCPToolsListResult{Tools: defs}
	case "tools/call":
		var p struct {
			Name      string          `json:"name"`
			Arguments json.RawMessage `json:"arguments"`
		}
		if err := json.Unmarshal(req.Params, &p); err != nil {
			rpcErr = &rpcError{Code: -32602, Message: "invalid params"}
			break
		}
		toolResult := resultbudget.ApplyToolResult(h.CallTool(ctx, id, p.Name, p.Arguments))
		result = toolResult
		deliveryReceipts = types.RuntimeDeliveryReceipts(toolResult)
		if toolResult != nil {
			deliveryOwner = toolResult.DeliveryOwner
		}
	default:
		rpcErr = &rpcError{Code: -32601, Message: "method not found"}
	}

	return &rpcResponse{
		JSONRPC:          "2.0",
		ID:               req.ID,
		Result:           result,
		Error:            rpcErr,
		deliveryReceipts: deliveryReceipts,
		deliveryOwner:    deliveryOwner,
	}, true
}

func acknowledgeMCPToolResultDelivery(handler types.MCPHandler, agentID string, runtimeSessionID string, responses []*rpcResponse, outcome types.MCPToolDeliveryOutcome) {
	seen := make(map[string]struct{})
	var fallbackReceipts []string
	for _, response := range responses {
		if response == nil {
			continue
		}
		responseReceipts := make([]string, 0, len(response.deliveryReceipts))
		for _, receipt := range response.deliveryReceipts {
			receipt = strings.TrimSpace(receipt)
			if !types.IsRuntimeDeliveryReceipt(receipt) {
				continue
			}
			if _, duplicate := seen[receipt]; duplicate {
				continue
			}
			seen[receipt] = struct{}{}
			responseReceipts = append(responseReceipts, receipt)
		}
		if len(responseReceipts) == 0 {
			continue
		}
		if response.deliveryOwner != nil {
			// The owner is bound to the exact Process that created these
			// receipts, so a concurrent Stop/replacement cannot redirect a
			// successful direct HTTP write to a different Process generation.
			response.deliveryOwner.AcknowledgeMCPToolResultDelivery(responseReceipts, outcome)
			continue
		}
		fallbackReceipts = append(fallbackReceipts, responseReceipts...)
	}
	if len(fallbackReceipts) > 0 {
		if acknowledger, ok := handler.(types.MCPToolResultDeliveryAcknowledger); ok {
			acknowledger.AcknowledgeMCPToolResultDelivery(agentID, strings.TrimSpace(runtimeSessionID), fallbackReceipts, outcome)
		}
	}
}

func mcpDeliveryOutcomeForWrite(c *gin.Context, writeErr error) (types.MCPToolDeliveryOutcome, bool) {
	deferred := strings.TrimSpace(c.GetHeader(types.MCPDeferredToolDeliveryHeader))
	if deferred == "1" || strings.EqualFold(deferred, "true") {
		// The Codex app-server owns a second hop even when this first HTTP write
		// reports short/failed. json.Encoder appends a newline, so a short write
		// can still leave a complete, parseable JSON response at the client. Keep
		// the exact lease transferred until the child write reports a terminal
		// outcome (or its bounded deferred lease/transport retirement settles it).
		return types.MCPToolDeliveryDeferred, true
	}
	if writeErr != nil {
		// A direct short/error write is ambiguous: the client may have received a
		// complete parseable response even though ResponseWriter returned short.
		// Do not release into replay reconciliation. Leaving the exact claim
		// unacknowledged retains its lifecycle lease until its bounded direct
		// orphan timer persists a fail-closed terminal Failed outcome.
		return "", false
	}
	return types.MCPToolDeliveryDelivered, true
}

// dispatchRPCBatch reads the request body, parses single or batch JSON-RPC
// requests, and dispatches them via DispatchRPC.
func dispatchRPCBatch(c *gin.Context, handler types.MCPHandler) {
	body, err := readBody(c)
	if err != nil {
		if errors.Is(err, errRequestBodyTooLarge) {
			writeMCPStatus(c, http.StatusRequestEntityTooLarge)
			return
		}
		writeMCPStatus(c, http.StatusBadRequest)
		return
	}

	trimmed := bytes.TrimSpace(body)
	if len(trimmed) == 0 {
		writeMCPStatus(c, http.StatusBadRequest)
		return
	}

	id := agentID(c)
	requestContext := context.WithValue(c.Request.Context(), requestHeadersContextKey{}, c.Request.Header.Clone())

	if trimmed[0] == '[' {
		var batch []json.RawMessage
		if err := json.Unmarshal(trimmed, &batch); err != nil {
			_ = writeMCPJSON(c, http.StatusBadRequest, parseErrorResponse("parse error"))
			return
		}
		if len(batch) == 0 {
			_ = writeMCPJSON(c, http.StatusBadRequest, invalidRequestResponse("invalid request"))
			return
		}

		responses := make([]*rpcResponse, 0, len(batch))
		for _, item := range batch {
			resp, ok := dispatchRPC(requestContext, id, handler, item)
			if ok {
				responses = append(responses, resp)
			}
		}
		if len(responses) == 0 {
			writeMCPStatus(c, http.StatusAccepted)
			return
		}
		writeErr := writeMCPJSON(c, http.StatusOK, responses)
		if outcome, acknowledge := mcpDeliveryOutcomeForWrite(c, writeErr); acknowledge {
			acknowledgeMCPToolResultDelivery(handler, id, c.GetHeader(types.MCPRuntimeSessionHeader), responses, outcome)
		}
		return
	}

	resp, ok := dispatchRPC(requestContext, id, handler, trimmed)
	if !ok {
		writeMCPStatus(c, http.StatusAccepted)
		return
	}
	writeErr := writeMCPJSON(c, http.StatusOK, resp)
	if outcome, acknowledge := mcpDeliveryOutcomeForWrite(c, writeErr); acknowledge {
		acknowledgeMCPToolResultDelivery(handler, id, c.GetHeader(types.MCPRuntimeSessionHeader), []*rpcResponse{resp}, outcome)
	}
}

// handlerResolver resolves a route key to its MCP handler.