package resultbudget

import (
	"encoding/json"
	"strings"
	"testing"

	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

func TestApplyToolResultBudgetsTextAndStructuredContent(t *testing.T) {
	items := make([]any, 0, 5)
	for i := 0; i < 5; i++ {
		items = append(items, strings.Repeat("s", 20))
	}
	result := &types.MCPToolResult{
		Content: []types.MCPToolContent{{Type: "text", Text: strings.Repeat("x", 80)}},
		StructuredContent: map[string]any{
			"items": items,
			"body":  strings.Repeat("y", 80),
		},
	}

	ApplyToolResultWithLimits(result, Limits{
		MaxContentBlocks:   1,
		MaxTextRunes:       40,
		MaxTextBlockRunes:  40,
		MaxStructuredBytes: 4096,
		MaxStructuredDepth: 4,
		MaxListItems:       2,
		MaxMapKeys:         10,
		MaxStringRunes:     10,
	})

	if got := len([]rune(result.Content[0].Text)); got > 40 {
		t.Fatalf("text runes = %d, want <= 40", got)
	}
	structured, ok := result.StructuredContent.(map[string]any)
	if !ok {
		t.Fatalf("structuredContent = %T, want map", result.StructuredContent)
	}
	if structured["truncated"] != true {
		t.Fatalf("truncated metadata missing: %#v", structured)
	}
	gotItems, ok := structured["items"].([]any)
	if !ok || len(gotItems) != 2 {
		t.Fatalf("items = %#v, want 2 budgeted items", structured["items"])
	}
	if body, _ := structured["body"].(string); len([]rune(body)) > len("ssssssssss… [truncated]") {
		t.Fatalf("body was not string-budgeted: %q", body)
	}
}

func TestApplyToolResultPreservesUntruncatedRawStructuredJSON(t *testing.T) {
	raw := json.RawMessage(`{"z":1.2300,"a":{"items":[true,null]}}`)
	result := &types.MCPToolResult{
		Content:           []types.MCPToolContent{{Type: "text", Text: string(raw)}},
		StructuredContent: raw,
	}

	ApplyToolResult(result)

	structured, ok := result.StructuredContent.(json.RawMessage)
	if !ok {
		t.Fatalf("structuredContent = %T, want json.RawMessage", result.StructuredContent)
	}
	if string(structured) != string(raw) {
		t.Fatalf("structuredContent = %s, want byte-exact %s", structured, raw)
	}
	encoded, err := json.Marshal(result.StructuredContent)
	if err != nil {
		t.Fatalf("marshal structuredContent: %v", err)
	}
	if string(encoded) != string(raw) {
		t.Fatalf("encoded structuredContent = %s, want %s", encoded, raw)
	}
}

func TestApplyRawToolResultBudgetsProviderJSON(t *testing.T) {
	raw := json.RawMessage(`{"content":[{"type":"text","text":"` + strings.Repeat("x", 60) + `"}],"structuredContent":{"rows":[1,2,3,4]}}`)
	budgeted := ApplyRawToolResultWithLimits(raw, Limits{
		MaxContentBlocks:   1,
		MaxTextRunes:       30,
		MaxTextBlockRunes:  30,
		MaxStructuredBytes: 4096,
		MaxStructuredDepth: 4,
		MaxListItems:       2,
		MaxMapKeys:         10,
		MaxStringRunes:     20,
	})

	var result types.MCPToolResult
	if err := json.Unmarshal(budgeted, &result); err != nil {
		t.Fatalf("unmarshal budgeted result: %v", err)
	}
	if got := len([]rune(result.Content[0].Text)); got > 30 {
		t.Fatalf("text runes = %d, want <= 30", got)
	}
	structured, ok := result.StructuredContent.(map[string]any)
	if !ok || structured["truncated"] != true {
		t.Fatalf("truncation metadata missing: %#v", result.StructuredContent)
	}
}

func TestApplyRawToolResultPreservesNonTextContentBlockFields(t *testing.T) {
	raw := json.RawMessage(`{"content":[{"type":"image","mimeType":"image/png","data":"` + strings.Repeat("x", 60) + `"}]}`)

	budgeted := ApplyRawToolResultWithLimits(raw, Limits{
		MaxContentBlocks:   1,
		MaxTextRunes:       30,
		MaxTextBlockRunes:  30,
		MaxStructuredBytes: 4096,
		MaxStructuredDepth: 4,
		MaxListItems:       2,
		MaxMapKeys:         10,
		MaxStringRunes:     10,
	})

	var result map[string]any
	if err := json.Unmarshal(budgeted, &result); err != nil {
		t.Fatalf("unmarshal budgeted result: %v", err)
	}
	content, ok := result["content"].([]any)
	if !ok || len(content) != 1 {
		t.Fatalf("content = %#v, want one content block", result["content"])
	}
	block, ok := content[0].(map[string]any)
	if !ok {
		t.Fatalf("content[0] = %T, want object", content[0])
	}
	if got := block["type"]; got != "image" {
		t.Fatalf("type = %v, want image", got)
	}
	if got := block["mimeType"]; got != "image/png" {
		t.Fatalf("mimeType = %v, want image/png", got)
	}
	data, ok := block["data"].(string)
	if !ok || !strings.Contains(data, "[truncated]") {
		t.Fatalf("data was not preserved and budgeted: %#v", block["data"])
	}
}

func TestApplyRawToolResultOmitsExtraTopLevelFields(t *testing.T) {
	raw := json.RawMessage(`{"content":[{"type":"text","text":"ok"}],"debug":"` + strings.Repeat("x", 60) + `","other":{"rows":[1,2,3,4]}}`)

	budgeted := ApplyRawToolResultWithLimits(raw, Limits{
		MaxContentBlocks:   1,
		MaxTextRunes:       30,
		MaxTextBlockRunes:  30,
		MaxStructuredBytes: 4096,
		MaxStructuredDepth: 4,
		MaxListItems:       2,
		MaxMapKeys:         10,
		MaxStringRunes:     10,
	})

	var result map[string]any
	if err := json.Unmarshal(budgeted, &result); err != nil {
		t.Fatalf("unmarshal budgeted result: %v", err)
	}
	if _, ok := result["debug"]; ok {
		t.Fatalf("debug top-level field was preserved: %#v", result)
	}
	if _, ok := result["other"]; ok {
		t.Fatalf("other top-level field was preserved: %#v", result)
	}
	structured, ok := result["structuredContent"].(map[string]any)
	if !ok || structured["truncated"] != true {
		t.Fatalf("truncation metadata missing: %#v", result["structuredContent"])
	}
	paths, ok := structured["omitted_paths"].([]any)
	if !ok {
		t.Fatalf("omitted_paths missing: %#v", structured)
	}
	var sawDebug bool
	for _, path := range paths {
		if path == "$.debug" {
			sawDebug = true
		}
	}
	if !sawDebug {
		t.Fatalf("omitted_paths = %#v, want $.debug", paths)
	}
}

func TestApplyRawToolResultReportsGenericJSONTruncation(t *testing.T) {
	raw := json.RawMessage(`{"rows":[1,2,3,4]}`)

	budgeted := ApplyRawToolResultWithLimits(raw, Limits{
		MaxContentBlocks:   1,
		MaxTextRunes:       30,
		MaxTextBlockRunes:  30,
		MaxStructuredBytes: 4096,
		MaxStructuredDepth: 4,
		MaxListItems:       2,
		MaxMapKeys:         10,
		MaxStringRunes:     20,
	})

	var result map[string]any
	if err := json.Unmarshal(budgeted, &result); err != nil {
		t.Fatalf("unmarshal budgeted generic JSON: %v", err)
	}
	if result["truncated"] != true {
		t.Fatalf("truncation metadata missing: %#v", result)
	}
	rows, ok := result["rows"].([]any)
	if !ok || len(rows) != 2 {
		t.Fatalf("rows = %#v, want 2 budgeted rows", result["rows"])
	}
}