package aifin

import (
	"bytes"
	"encoding/json"
	"math"
	"strconv"
	"strings"

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

func normalizeToolResult(
	result mcp.ToolResult,
) (json.RawMessage, error) {
	var candidate json.RawMessage
	structured := bytes.TrimSpace(result.StructuredContent)
	if result.StructuredContentPresent ||
		len(result.StructuredContent) > 0 {
		if len(structured) == 0 ||
			bytes.Equal(structured, []byte("null")) {
			return nil, invalidUpstreamResult()
		}
		candidate = append(json.RawMessage(nil), structured...)
	} else {
		fallback, err := uniqueTextFallback(result.Content)
		if err != nil {
			return nil, err
		}
		candidate = fallback
	}
	compact, decoded, err := runner.DecodeUniqueJSONDocument(candidate)
	if err != nil || decoded == nil || hasRejectedResultShape(decoded) {
		return nil, invalidUpstreamResult()
	}
	encoded, err := json.Marshal(struct {
		Result json.RawMessage `json:"result"`
	}{Result: compact})
	if err != nil || len(encoded) > runner.MaxResponseBytes {
		return nil, invalidUpstreamResult()
	}
	return json.RawMessage(encoded), nil
}

func uniqueTextFallback(
	content []json.RawMessage,
) (json.RawMessage, error) {
	if len(content) != 1 {
		return nil, invalidUpstreamResult()
	}
	_, decoded, err := runner.DecodeUniqueJSONDocument(content[0])
	if err != nil {
		return nil, invalidUpstreamResult()
	}
	item, ok := decoded.(map[string]any)
	if !ok {
		return nil, invalidUpstreamResult()
	}
	for key := range item {
		switch key {
		case "type", "text", "annotations", "_meta":
		default:
			return nil, invalidUpstreamResult()
		}
	}
	for _, key := range []string{"annotations", "_meta"} {
		if metadata, exists := item[key]; exists {
			if _, ok := metadata.(map[string]any); !ok {
				return nil, invalidUpstreamResult()
			}
		}
	}
	contentType, typeOK := item["type"].(string)
	text, textOK := item["text"].(string)
	if !typeOK ||
		!textOK ||
		contentType != "text" ||
		strings.TrimSpace(text) == "" {
		return nil, invalidUpstreamResult()
	}
	return json.RawMessage(text), nil
}

func hasRejectedResultShape(value any) bool {
	switch typed := value.(type) {
	case []any:
		for _, item := range typed {
			if hasRejectedResultShapeAt(item, false) {
				return true
			}
		}
		return false
	case map[string]any:
		return hasRejectedResultRecord(typed, true)
	default:
		return false
	}
}

func hasRejectedResultShapeAt(value any, root bool) bool {
	switch typed := value.(type) {
	case []any:
		for _, item := range typed {
			if hasRejectedResultShapeAt(item, false) {
				return true
			}
		}
	case map[string]any:
		return hasRejectedResultRecord(typed, root)
	}
	return false
}

func hasRejectedResultRecord(record map[string]any, root bool) bool {
	if looksLikeMCPEnvelope(record, root) {
		return true
	}
	if resultType, ok := record["resultType"].(string); ok &&
		resultType == "input_required" {
		return true
	}
	if requests, exists := record["inputRequests"]; exists {
		requestMap, ok := requests.(map[string]any)
		if !ok || len(requestMap) > 0 {
			return true
		}
	}
	if code, exists := record["mcp_tool_error_code"]; exists &&
		isNonZeroWindCode(code) {
		return true
	}
	if embedded, exists := record["error"]; exists &&
		hasMeaningfulEmbeddedError(embedded) {
		return true
	}
	if root && windDirectEmbeddedDataFailed(record) {
		return true
	}
	for _, item := range record {
		if hasRejectedResultShapeAt(item, false) {
			return true
		}
	}
	return false
}

// windDirectEmbeddedDataFailed mirrors the pinned Wind wrapper: only the
// top-level result's direct data child can carry a finite, number-coercible
// status code. Non-zero 2xx codes succeed only with an OK/SUCCESS message.
func windDirectEmbeddedDataFailed(record map[string]any) bool {
	embedded, ok := record["data"].(map[string]any)
	if !ok {
		return false
	}
	rawCode, exists := embedded["code"]
	if !exists {
		return false
	}
	code, ok := numericBusinessCode(rawCode)
	if !ok ||
		math.IsNaN(code) ||
		math.IsInf(code, 0) ||
		code == 0 {
		return false
	}
	message, _ := embedded["message"].(string)
	normalizedMessage := strings.ToUpper(strings.TrimSpace(message))
	return code < 200 ||
		code >= 300 ||
		(normalizedMessage != "OK" && normalizedMessage != "SUCCESS")
}

func looksLikeMCPEnvelope(record map[string]any, root bool) bool {
	if version, ok := record["jsonrpc"].(string); ok &&
		version == "2.0" {
		return true
	}
	if _, exists := record["isError"]; exists {
		return true
	}
	if _, exists := record["structuredContent"]; exists {
		return true
	}
	content, exists := record["content"]
	if !exists {
		return false
	}
	blocks, isArray := content.([]any)
	if !isArray {
		return false
	}
	return root || isHighConfidenceMCPContent(blocks)
}

func isHighConfidenceMCPContent(blocks []any) bool {
	if len(blocks) == 0 {
		return false
	}
	for _, block := range blocks {
		record, ok := block.(map[string]any)
		if !ok || !isMCPContentBlock(record) {
			return false
		}
	}
	return true
}

func isMCPContentBlock(record map[string]any) bool {
	contentType, ok := record["type"].(string)
	if !ok {
		return false
	}
	switch contentType {
	case "text":
		_, ok := record["text"].(string)
		return ok
	case "image", "audio":
		_, dataOK := record["data"].(string)
		_, mimeTypeOK := record["mimeType"].(string)
		return dataOK && mimeTypeOK
	case "resource":
		_, ok := record["resource"].(map[string]any)
		return ok
	case "resource_link":
		_, uriOK := record["uri"].(string)
		_, nameOK := record["name"].(string)
		return uriOK && nameOK
	default:
		return false
	}
}

func isNonZeroWindCode(value any) bool {
	if code, ok := numericBusinessCode(value); ok {
		return code != 0
	}
	text, ok := value.(string)
	return ok && strings.TrimSpace(text) != ""
}

func hasMeaningfulEmbeddedError(value any) bool {
	if value == nil {
		return false
	}
	record, ok := value.(map[string]any)
	if !ok {
		return true
	}
	if message, ok := record["message"].(string); ok &&
		strings.TrimSpace(message) != "" {
		return true
	}
	switch code := record["code"].(type) {
	case string:
		return strings.TrimSpace(code) != ""
	case float64:
		return code != 0
	default:
		return false
	}
}

func numericBusinessCode(value any) (float64, bool) {
	switch code := value.(type) {
	case float64:
		return code, true
	case string:
		if code == "" {
			return 0, false
		}
		parsed, err := strconv.ParseFloat(strings.TrimSpace(code), 64)
		if err != nil {
			return 0, false
		}
		return parsed, true
	default:
		return 0, false
	}
}

func invalidUpstreamResult() error {
	return runner.Errorf(
		runner.ErrorUpstreamFailure,
		"Wind AIFin Market returned an invalid MCP result",
		nil,
	)
}