// Package resultbudget applies final, protocol-level budgets before tool
// results are written back to runtime clients. It is transport-neutral: MCP
// and Exec adapt their result shapes at their respective return boundaries.
package resultbudget

import (
	"encoding/json"
	"fmt"
	"sort"
	"strings"
	"unicode/utf8"

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

const truncationNotice = "\n\n[truncated by result budget; use a narrower tool call, cursor, section, or explicit get tool for more.]"

// Limits defines product-level safety caps for a single MCP tool result.
type Limits struct {
	MaxContentBlocks   int
	MaxTextRunes       int
	MaxTextBlockRunes  int
	MaxStructuredBytes int
	MaxStructuredDepth int
	MaxListItems       int
	MaxMapKeys         int
	MaxStringRunes     int
}

// Report describes the budget cuts applied to one result.
type Report struct {
	Truncated    bool     `json:"truncated"`
	BudgetKind   string   `json:"budget_kind"`
	OmittedCount int      `json:"omitted_count,omitempty"`
	OmittedPaths []string `json:"omitted_paths,omitempty"`
}

// Structured applies the structured-value portion of the shared budget. It
// preserves JSON-like shape where possible and returns an explicit report when
// truncation occurred. The caller owns where that report is surfaced.
func Structured(value any, limits Limits) (any, Report) {
	state := &budgetState{limits: normalizeLimits(limits)}
	budgeted := state.budgetValue(normalizeJSONLike(value), "$", 0)
	if encoded, err := json.Marshal(budgeted); err == nil && len(encoded) > state.limits.MaxStructuredBytes {
		state.omit("$", "structured value exceeded byte budget")
		budgeted = map[string]any{
			"truncated":     true,
			"budget_kind":   "structured",
			"omitted_paths": []string{"$"},
			"message":       "structured value exceeded result budget; use a narrower request",
		}
	}
	if !state.truncated {
		return budgeted, Report{}
	}
	return budgeted, state.report()
}

// Text applies the text portion of the shared budget to one logical text
// output and returns an explicit report when truncation occurred.
func Text(value string, limits Limits) (string, Report) {
	state := &budgetState{limits: normalizeLimits(limits)}
	remaining := state.limits.MaxTextRunes
	budgeted := state.budgetTextContent(value, "$", &remaining)
	if !state.truncated {
		return budgeted, Report{}
	}
	return budgeted, state.report()
}

// DefaultLimits returns the standard final-boundary MCP result budget.
func DefaultLimits() Limits {
	return Limits{
		MaxContentBlocks:   20,
		MaxTextRunes:       24 * 1024,
		MaxTextBlockRunes:  8 * 1024,
		MaxStructuredBytes: 32 * 1024,
		MaxStructuredDepth: 8,
		MaxListItems:       100,
		MaxMapKeys:         100,
		MaxStringRunes:     4096,
	}
}

// ApplyToolResult applies the default MCP result budget to result.
func ApplyToolResult(result *types.MCPToolResult) *types.MCPToolResult {
	return ApplyToolResultWithLimits(result, DefaultLimits())
}

// ApplyToolResultWithLimits applies limits to result in place and returns it.
func ApplyToolResultWithLimits(result *types.MCPToolResult, limits Limits) *types.MCPToolResult {
	if result == nil {
		return nil
	}
	limits = normalizeLimits(limits)
	state := &budgetState{limits: limits}

	result.Content = state.budgetContent(result.Content)
	if result.StructuredContent != nil {
		if raw, ok := rawStructuredContent(result.StructuredContent); ok && rawStructuredContentWithinLimits(raw, limits) {
			result.StructuredContent = raw
		} else {
			normalized := normalizeJSONLike(result.StructuredContent)
			result.StructuredContent = state.budgetValue(normalized, "$.structuredContent", 0)
			if encoded, err := json.Marshal(result.StructuredContent); err == nil && len(encoded) > limits.MaxStructuredBytes {
				state.omit("$.structuredContent", "structuredContent exceeded byte budget")
				result.StructuredContent = map[string]any{
					"truncated":     true,
					"budget_kind":   "structuredContent",
					"omitted_paths": []string{"$.structuredContent"},
					"message":       "structuredContent exceeded MCP result budget; use a narrower tool call, cursor, section, or explicit get tool for more",
				}
			}
		}
	}
	if result.Meta != nil {
		budgeted := state.budgetStructuredToolField(result.Meta, "$._meta")
		if metadata, ok := budgeted.(map[string]any); ok {
			result.Meta = metadata
		} else {
			state.omit("$._meta", "invalid MCP result metadata")
			result.Meta = nil
		}
	}

	if state.truncated {
		result.StructuredContent = withReport(result.StructuredContent, state.report())
	}
	return result
}

func rawStructuredContent(value any) (json.RawMessage, bool) {
	raw, ok := value.(json.RawMessage)
	if !ok || !json.Valid(raw) {
		return nil, false
	}
	return append(json.RawMessage(nil), raw...), true
}

func rawStructuredContentWithinLimits(raw json.RawMessage, limits Limits) bool {
	if len(raw) > limits.MaxStructuredBytes {
		return false
	}
	var value any
	decoder := json.NewDecoder(strings.NewReader(string(raw)))
	decoder.UseNumber()
	if err := decoder.Decode(&value); err != nil {
		return false
	}
	state := &budgetState{limits: limits}
	state.budgetValue(value, "$.structuredContent", 0)
	return !state.truncated
}

// ApplyRawToolResult applies the default result budget to a raw JSON value that
// should contain an MCP tool result. If decoding as a tool result fails, the raw
// JSON value is budgeted as generic structured content.
func ApplyRawToolResult(raw json.RawMessage) json.RawMessage {
	return ApplyRawToolResultWithLimits(raw, DefaultLimits())
}

// ApplyRawToolResultWithLimits is ApplyRawToolResult with explicit limits.
func ApplyRawToolResultWithLimits(raw json.RawMessage, limits Limits) json.RawMessage {
	trimmed := strings.TrimSpace(string(raw))
	if trimmed == "" || trimmed == "null" {
		return raw
	}

	var object map[string]any
	if err := json.Unmarshal(raw, &object); err == nil && looksLikeToolResultObject(object) {
		state := &budgetState{limits: normalizeLimits(limits)}
		state.budgetToolResultObject(object)
		if state.truncated {
			object["structuredContent"] = withReport(object["structuredContent"], state.report())
		}
		if encoded, err := json.Marshal(object); err == nil {
			return encoded
		}
		return raw
	}

	var value any
	if err := json.Unmarshal(raw, &value); err != nil {
		return raw
	}
	state := &budgetState{limits: normalizeLimits(limits)}
	budgeted := state.budgetValue(value, "$", 0)
	if state.truncated {
		budgeted = withReport(budgeted, state.report())
	}
	if encoded, err := json.Marshal(budgeted); err == nil && len(encoded) <= state.limits.MaxStructuredBytes {
		return encoded
	}
	state.omit("$", "raw result exceeded byte budget")
	fallback := withReport(nil, state.report())
	encoded, err := json.Marshal(fallback)
	if err != nil {
		return raw
	}
	return encoded
}

func looksLikeToolResultObject(object map[string]any) bool {
	_, hasContent := object["content"]
	_, hasStructured := object["structuredContent"]
	_, hasIsError := object["isError"]
	return hasContent || hasStructured || hasIsError
}

type budgetState struct {
	limits       Limits
	truncated    bool
	omittedCount int
	omittedPaths []string
}

const (
	maxReportedPaths     = 20
	maxReportedPathRunes = 256
)

func (s *budgetState) budgetToolResultObject(object map[string]any) {
	for key, value := range object {
		switch key {
		case "content":
			if content, ok := value.([]any); ok {
				object[key] = s.budgetRawContent(content)
			} else {
				object[key] = s.budgetValue(normalizeJSONLike(value), "$.content", 0)
			}
		case "structuredContent":
			object[key] = s.budgetStructuredToolField(value, "$.structuredContent")
		case "isError":
			if _, ok := value.(bool); !ok {
				object[key] = s.budgetValue(normalizeJSONLike(value), "$.isError", 0)
			}
		case "_meta":
			object[key] = s.budgetStructuredToolField(value, "$._meta")
		default:
			s.omit("$."+key, "omitted non-standard MCP tool result top-level field")
			delete(object, key)
		}
	}
}

func (s *budgetState) budgetStructuredToolField(value any, path string) any {
	budgeted := s.budgetValue(normalizeJSONLike(value), path, 0)
	if encoded, err := json.Marshal(budgeted); err == nil && len(encoded) > s.limits.MaxStructuredBytes {
		s.omit(path, "structured field exceeded byte budget")
		return map[string]any{
			"truncated":     true,
			"budget_kind":   "structuredContent",
			"omitted_paths": []string{path},
			"message":       "structured field exceeded MCP result budget; use a narrower tool call, cursor, section, or explicit get tool for more",
		}
	}
	return budgeted
}

func (s *budgetState) budgetRawContent(content []any) []any {
	if len(content) == 0 {
		return content
	}
	maxBlocks := s.limits.MaxContentBlocks
	if len(content) > maxBlocks {
		s.omit("$.content", fmt.Sprintf("omitted %d content blocks", len(content)-maxBlocks))
		content = content[:maxBlocks]
	}

	remaining := s.limits.MaxTextRunes
	out := make([]any, len(content))
	for i, block := range content {
		path := fmt.Sprintf("$.content[%d]", i)
		object, ok := block.(map[string]any)
		if !ok {
			out[i] = s.budgetValue(block, path, 0)
			continue
		}
		budgeted := make(map[string]any, len(object))
		blockType, _ := object["type"].(string)
		for key, value := range object {
			if key == "text" && blockType == "text" {
				text, _ := value.(string)
				budgeted[key] = s.budgetTextContent(text, fmt.Sprintf("%s.text", path), &remaining)
				continue
			}
			budgeted[key] = s.budgetValue(value, path+"."+key, 1)
		}
		out[i] = budgeted
	}
	return out
}

func (s *budgetState) budgetTextContent(text, path string, remaining *int) string {
	if text == "" {
		return text
	}
	limit := s.limits.MaxTextBlockRunes
	if *remaining < limit {
		limit = *remaining
	}
	if limit <= 0 {
		s.omit(path, "text omitted after total text budget")
		return "[omitted by MCP result text budget]"
	}
	truncated, ok := truncateWithNotice(text, limit)
	if !ok {
		s.omit(path, "text block exceeded budget")
	} else {
		truncated = text
	}
	*remaining -= utf8.RuneCountInString(truncated)
	if *remaining < 0 {
		*remaining = 0
	}
	return truncated
}

func (s *budgetState) budgetContent(content []types.MCPToolContent) []types.MCPToolContent {
	if len(content) == 0 {
		return content
	}
	maxBlocks := s.limits.MaxContentBlocks
	if len(content) > maxBlocks {
		s.omit("$.content", fmt.Sprintf("omitted %d content blocks", len(content)-maxBlocks))
		content = content[:maxBlocks]
	}

	remaining := s.limits.MaxTextRunes
	out := make([]types.MCPToolContent, len(content))
	for i, block := range content {
		out[i] = block
		if block.Type != "text" || block.Text == "" {
			continue
		}
		out[i].Text = s.budgetTextContent(block.Text, fmt.Sprintf("$.content[%d].text", i), &remaining)
	}
	return out
}

func (s *budgetState) budgetValue(value any, path string, depth int) any {
	if depth > s.limits.MaxStructuredDepth {
		s.omit(path, "max structuredContent depth exceeded")
		return "[truncated by MCP structuredContent depth budget]"
	}
	switch v := value.(type) {
	case nil, bool, float64, int, int64, json.Number:
		return v
	case string:
		if utf8.RuneCountInString(v) <= s.limits.MaxStringRunes {
			return v
		}
		s.omit(path, "string exceeded structuredContent budget")
		return truncateRunes(v, s.limits.MaxStringRunes) + "… [truncated]"
	case []any:
		limit := s.limits.MaxListItems
		if len(v) > limit {
			s.omit(path, fmt.Sprintf("omitted %d list items", len(v)-limit))
			v = v[:limit]
		}
		out := make([]any, len(v))
		for i, item := range v {
			out[i] = s.budgetValue(item, fmt.Sprintf("%s[%d]", path, i), depth+1)
		}
		return out
	case map[string]any:
		keys := make([]string, 0, len(v))
		for key := range v {
			keys = append(keys, key)
		}
		sort.Strings(keys)
		if len(keys) > s.limits.MaxMapKeys {
			s.omit(path, fmt.Sprintf("omitted %d map keys", len(keys)-s.limits.MaxMapKeys))
			keys = keys[:s.limits.MaxMapKeys]
		}
		out := make(map[string]any, len(keys))
		for _, key := range keys {
			out[key] = s.budgetValue(v[key], path+"."+key, depth+1)
		}
		return out
	default:
		return s.budgetValue(normalizeJSONLike(v), path, depth)
	}
}

func (s *budgetState) omit(path, _ string) {
	s.truncated = true
	s.omittedCount++
	if len(s.omittedPaths) < maxReportedPaths {
		s.omittedPaths = append(s.omittedPaths, truncateRunes(path, maxReportedPathRunes))
	}
}

func (s *budgetState) report() Report {
	paths := append([]string(nil), s.omittedPaths...)
	return Report{
		Truncated:    true,
		BudgetKind:   "tool_result",
		OmittedCount: s.omittedCount,
		OmittedPaths: paths,
	}
}

func withReport(content any, report Report) any {
	if object, ok := content.(map[string]any); ok {
		copy := make(map[string]any, len(object)+4)
		for key, value := range object {
			copy[key] = value
		}
		copy["truncated"] = report.Truncated
		copy["budget_kind"] = report.BudgetKind
		if report.OmittedCount > 0 {
			copy["omitted_count"] = report.OmittedCount
		}
		if len(report.OmittedPaths) > 0 {
			copy["omitted_paths"] = report.OmittedPaths
		}
		return copy
	}
	result := map[string]any{
		"truncated":   report.Truncated,
		"budget_kind": report.BudgetKind,
	}
	if content != nil {
		result["result"] = content
	}
	if report.OmittedCount > 0 {
		result["omitted_count"] = report.OmittedCount
	}
	if len(report.OmittedPaths) > 0 {
		result["omitted_paths"] = report.OmittedPaths
	}
	return result
}

func normalizeLimits(limits Limits) Limits {
	defaults := DefaultLimits()
	if limits.MaxContentBlocks <= 0 {
		limits.MaxContentBlocks = defaults.MaxContentBlocks
	}
	if limits.MaxTextRunes <= 0 {
		limits.MaxTextRunes = defaults.MaxTextRunes
	}
	if limits.MaxTextBlockRunes <= 0 {
		limits.MaxTextBlockRunes = defaults.MaxTextBlockRunes
	}
	if limits.MaxStructuredBytes <= 0 {
		limits.MaxStructuredBytes = defaults.MaxStructuredBytes
	}
	if limits.MaxStructuredDepth <= 0 {
		limits.MaxStructuredDepth = defaults.MaxStructuredDepth
	}
	if limits.MaxListItems <= 0 {
		limits.MaxListItems = defaults.MaxListItems
	}
	if limits.MaxMapKeys <= 0 {
		limits.MaxMapKeys = defaults.MaxMapKeys
	}
	if limits.MaxStringRunes <= 0 {
		limits.MaxStringRunes = defaults.MaxStringRunes
	}
	return limits
}

func normalizeJSONLike(value any) any {
	encoded, err := json.Marshal(value)
	if err != nil {
		return fmt.Sprint(value)
	}
	var out any
	if err := json.Unmarshal(encoded, &out); err != nil {
		return string(encoded)
	}
	return out
}

func truncateWithNotice(text string, limit int) (string, bool) {
	if utf8.RuneCountInString(text) <= limit {
		return text, true
	}
	noticeRunes := utf8.RuneCountInString(truncationNotice)
	if limit <= noticeRunes {
		return truncateRunes(text, limit), false
	}
	return truncateRunes(text, limit-noticeRunes) + truncationNotice, false
}

func truncateRunes(text string, limit int) string {
	if limit <= 0 {
		return ""
	}
	if utf8.RuneCountInString(text) <= limit {
		return text
	}
	var b strings.Builder
	b.Grow(limit)
	count := 0
	for _, r := range text {
		if count >= limit {
			break
		}
		b.WriteRune(r)
		count++
	}
	return b.String()
}