package pkulaw

import (
	"bytes"
	"encoding/json"
	"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 := uniqueStrictTextFallback(result.Content)
		if err != nil {
			return nil, err
		}
		candidate = fallback
	}

	compact, decoded, err := runner.DecodeUniqueJSONDocument(candidate)
	if err != nil ||
		isEmptyRootJSONValue(decoded) ||
		hasRejectedEmbeddedShape(decoded, true) {
		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 uniqueStrictTextFallback(
	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()
		}
	}
	if annotations, exists := item["annotations"]; exists {
		if _, ok := annotations.(map[string]any); !ok {
			return nil, invalidUpstreamResult()
		}
	}
	if metadata, exists := item["_meta"]; exists {
		if _, ok := metadata.(map[string]any); !ok {
			return nil, invalidUpstreamResult()
		}
	}
	if len(item) < 2 {
		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 isEmptyRootJSONValue(value any) bool {
	switch typed := value.(type) {
	case nil:
		return true
	case string:
		return strings.TrimSpace(typed) == ""
	case []any:
		return len(typed) == 0
	case map[string]any:
		return len(typed) == 0
	default:
		return false
	}
}

func hasRejectedEmbeddedShape(value any, root bool) bool {
	switch typed := value.(type) {
	case []any:
		for _, item := range typed {
			if hasRejectedEmbeddedShape(item, false) {
				return true
			}
		}
		return false
	case map[string]any:
		if looksLikeMCPEnvelope(typed, root) {
			return true
		}
		if resultType, ok := typed["resultType"].(string); ok &&
			resultType == "input_required" {
			return true
		}
		if requests, exists := typed["inputRequests"]; exists &&
			isMeaningfulErrorValue(requests) {
			return true
		}
		message, hasMessage := typed["message"].(string)
		if hasMessage && strings.TrimSpace(message) != "" {
			if code, exists := typed["code"]; exists &&
				isNonZeroErrorCode(code) {
				return true
			}
			if status, exists := typed["status"]; exists &&
				isRecognizableErrorStatus(status) {
				return true
			}
		}
		for key, item := range typed {
			normalizedKey := strings.ToLower(strings.TrimSpace(key))
			switch normalizedKey {
			case "error", "errors":
				if isMeaningfulErrorValue(item) {
					return true
				}
			case "errorcode", "error_code", "mcp_tool_error_code":
				if isNonZeroErrorCode(item) {
					return true
				}
			case "success":
				if succeeded, ok := item.(bool); ok && !succeeded {
					return true
				}
			case "status":
				if isRecognizableErrorStatus(item) {
					return true
				}
			}
		}
		for _, item := range typed {
			if hasRejectedEmbeddedShape(item, false) {
				return true
			}
		}
	}
	return false
}

func looksLikeMCPEnvelope(record map[string]any, root bool) bool {
	if version, ok := record["jsonrpc"].(string); ok &&
		version == "2.0" {
		return true
	}
	for key := range record {
		switch strings.ToLower(strings.TrimSpace(key)) {
		case "iserror", "structuredcontent":
			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
	case "tool_use":
		_, idOK := record["id"].(string)
		_, nameOK := record["name"].(string)
		_, inputOK := record["input"].(map[string]any)
		return idOK && nameOK && inputOK
	case "tool_result":
		_, ok := record["toolUseId"].(string)
		return ok
	default:
		return false
	}
}

func isMeaningfulErrorValue(value any) bool {
	switch typed := value.(type) {
	case nil:
		return false
	case bool:
		return typed
	case string:
		return strings.TrimSpace(typed) != ""
	case float64:
		return typed != 0
	case []any:
		return len(typed) != 0
	case map[string]any:
		return len(typed) != 0
	default:
		return true
	}
}

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

func isRecognizableErrorStatus(value any) bool {
	switch typed := value.(type) {
	case float64:
		return typed >= 400
	case string:
		raw := strings.TrimSpace(typed)
		if numeric, err := strconv.ParseFloat(raw, 64); err == nil {
			return numeric >= 400
		}
		fields := strings.Fields(raw)
		if len(fields) > 0 {
			if numeric, err := strconv.ParseFloat(fields[0], 64); err == nil {
				return numeric >= 400
			}
		}
		if isExplicitFailureStatus(raw) {
			return true
		}
	}
	return false
}

func isExplicitFailureStatus(value string) bool {
	switch strings.ToLower(strings.TrimSpace(value)) {
	case "error", "failed", "failure":
		return true
	default:
		return false
	}
}

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