package exec

import (
	"encoding/json"
	"fmt"
	"strings"

	"github.com/tidbcloud/AgentTeam/desktop/internal/resultbudget"
	"github.com/tidbcloud/AgentTeam/internal/workflow/run"
)

const (
	errorBudgetMetadataKey   = "error_budget"
	maxSerializedResultBytes = 64 * 1024
)

// applyActionResultBudget only type-erases the concrete Data DTO. It must not
// use the generic resultbudget projection here: that projection is allowed to
// shorten arbitrary MCP values, while an Exec action's Data object is a typed
// success contract with required identifiers. The action-aware marshal boundary
// below either returns the complete envelope or a typed failure.
func applyActionResultBudget[O any](result actionResult[O]) actionResult[any] {
	return actionResult[any]{Data: result.Data, Text: result.Text, Metadata: result.Metadata}
}

// marshalTypedActionResultForAction is the single built-in Exec success
// boundary. A complete typed envelope at or below 64 KiB is returned byte-for-
// byte unchanged. An oversize or unserializable result becomes a stable
// machine-readable Error; it is never replaced with a generic map that could
// satisfy JSON syntax while violating the action's output schema.
func marshalTypedActionResultForAction[O any](domain, action string, result actionResult[O]) ([]byte, error) {
	return marshalTypedActionResultForActionWithInput(domain, action, nil, result)
}

func marshalTypedActionResultForActionWithInput[O any](domain, action string, rawInput json.RawMessage, result actionResult[O]) ([]byte, error) {
	budgeted := applyActionResultBudget(result)
	encoded, err := json.Marshal(budgeted)
	if err != nil {
		if builtinActionIsMutationForInput(domain, action, rawInput) {
			return nil, mutationSerializationFailureError(domain, action, err)
		}
		return nil, &Error{
			Code:    CodeInternal,
			Problem: "Exec action result could not be serialized",
			Cause:   err,
		}
	}
	if len(encoded) > maxSerializedResultBytes {
		if builtinActionIsMutationForInput(domain, action, rawInput) {
			return nil, mutationResultBudgetExceededError(domain, action, rawInput, result.Data, len(encoded), nil)
		}
		return nil, resultBudgetExceededError(domain, action, len(encoded), nil)
	}
	return encoded, nil
}

func builtinActionIsMutation(domain, action string) bool {
	if domain == "workflow" {
		// Workflow's public action is a union whose nested operation determines
		// whether it is a mutation. Without the raw discriminator, fail closed.
		return false
	}
	policy, ok := builtinActionPolicyFor(domain, action)
	return ok && policy.Risk != builtinRiskReadOnly
}

func builtinActionIsMutationForInput(domain, action string, rawInput json.RawMessage) bool {
	if domain == "workflow" {
		return workflowOperationIsMutationForInput(action, rawInput)
	}
	return builtinActionIsMutation(domain, action)
}

func mutationResultBudgetExceededError(domain, action string, rawInput json.RawMessage, data any, serializedBytes int, cause error) *Error {
	receipt := mutationReceiptForAction(domain, action, rawInput, data)
	failure := resultBudgetExceededError(domain, action, serializedBytes, cause)
	failure.Next = ""
	failure.Retryable = false
	if receipt == nil || !receipt.Applied {
		failure.Hint = "Mutation result exceeded the response budget; inspect the authoritative resource before issuing another command."
		return failure
	}
	failure.Hint = mutationAppliedHint
	failure.mutationApplied = true
	failure.mutationReceipt = receipt
	return failure
}

// resultBudgetExceededError intentionally exposes only action identity and the
// reviewed continuation policy. The byte count is useful for operators but is
// not copied into the public envelope; raw serialization causes and payloads
// remain private.
func resultBudgetExceededError(domain, action string, serializedBytes int, cause error) *Error {
	if cause == nil {
		cause = fmt.Errorf("serialized result size %d exceeds %d", serializedBytes, maxSerializedResultBytes)
	}
	policy, ok := builtinActionPolicyFor(domain, action)
	continuation := "retry with a narrower request or an action-specific continuation"
	policyName := "narrow"
	if ok {
		continuation = policy.Continuation
		policyName = string(policy.ResultPolicy)
	}
	return &Error{
		Code:    CodeResultBudgetExceeded,
		Problem: fmt.Sprintf("Exec action %s %s result exceeds the result budget (64 KiB success boundary)", domain, action),
		Hint:    fmt.Sprintf("No partial result was returned (result_policy=%s); %s.", policyName, continuation),
		Next:    fmt.Sprintf("agentteam exec %s %s with a narrower request", domain, action),
		Cause:   cause,
	}
}

// ApplyErrorBudget is the single Exec failure-return boundary. Classification
// happens first so raw causes remain daemon-only, then the public envelope is
// bounded and gains explicit truncation metadata. The allowlisted canonical
// Workflow contract fields (stage/path/serverCode/diagnostics with field/index)
// survive deterministic bounded projection; raw causes are never exposed.
func ApplyErrorBudget(err error) ErrorEnvelope {
	classified := classify(err)
	if classified.Code == CodeValidation && strings.TrimSpace(classified.Problem) == "invalid input" && classified.Cause != nil {
		classified.Problem = validationProblem(classified)
		if classified.Hint == "" {
			classified.Hint = "Check `agentteam exec <domain> <action> --input-schema` and fix the JSON field indicated in problem."
		}
	}
	if classified.mutationApplied && classified.mutationReceipt != nil && classified.mutationReceipt.Applied {
		// A receipt is an applied-state signal. Never carry server-provided retry
		// advice alongside it, even when the daemon returned a malformed or stale
		// hint/next pair around an otherwise valid closed receipt.
		classified.Hint = mutationAppliedHint
		classified.Next = ""
		classified.Retryable = false
	}
	envelope := ErrorEnvelope{Code: classified.Code.normalized(), Retryable: classified.Retryable}
	report := resultbudget.Report{}
	var fieldReport resultbudget.Report
	envelope.Problem, fieldReport = budgetErrorField(classified.Problem, maxProblemRunes)
	report = mergeBudgetReports(report, fieldReport)
	envelope.Hint, fieldReport = budgetErrorField(classified.Hint, maxHintRunes)
	report = mergeBudgetReports(report, fieldReport)
	envelope.Next, fieldReport = budgetErrorField(classified.Next, maxNextRunes)
	report = mergeBudgetReports(report, fieldReport)
	envelope.Stage, fieldReport = budgetErrorField(classified.Stage, maxStageRunes)
	report = mergeBudgetReports(report, fieldReport)
	envelope.Path, fieldReport = budgetErrorField(classified.Path, maxPathRunes)
	report = mergeBudgetReports(report, fieldReport)
	envelope.ServerCode, fieldReport = budgetErrorField(classified.ServerCode, maxServerCodeRunes)
	report = mergeBudgetReports(report, fieldReport)
	envelope.Diagnostics, report = budgetDiagnostics(classified.Diagnostics, report)
	envelope.Truncation = boundedServerTruncation(classified.Truncation)
	envelope.Truncated = classified.Truncated
	if report.Truncated {
		envelope.Metadata = map[string]any{errorBudgetMetadataKey: report}
	}
	attachMutationReceiptMetadata(&envelope, classified.mutationReceipt)
	return envelope
}

// boundedServerTruncation projects the server-declared diagnostic truncation
// metadata deterministically, capping the field list and UTF-8-safe values.
func boundedServerTruncation(truncation *Truncation) *Truncation {
	if truncation == nil {
		return nil
	}
	bounded := &Truncation{
		DiagnosticsTruncated: truncation.DiagnosticsTruncated,
		EnvelopeTruncated:    truncation.EnvelopeTruncated,
		OmittedCount:         truncation.OmittedCount,
		SerializedBytes:      truncation.SerializedBytes,
	}
	for _, field := range truncation.TruncatedFields {
		if capped := capField(field, maxDiagnosticFieldRunes); capped != "" {
			bounded.TruncatedFields = append(bounded.TruncatedFields, capped)
		}
		if len(bounded.TruncatedFields) >= 20 {
			break
		}
	}
	return bounded
}

// budgetDiagnostics deterministically bounds the typed diagnostic list: count
// capped and every field rune-capped, with explicit truncation metadata when any
// cut occurs. Raw cause text is never projected here because diagnostics come
// only from the allowlisted server contract or typed handler fields.
func budgetDiagnostics(diagnostics []Diagnostic, prior resultbudget.Report) ([]Diagnostic, resultbudget.Report) {
	if len(diagnostics) == 0 {
		return diagnostics, prior
	}
	report := prior
	if len(diagnostics) > maxDiagnostics {
		report = mergeBudgetReports(report, resultbudget.Report{
			Truncated: true, BudgetKind: "tool_result",
			OmittedCount: len(diagnostics) - maxDiagnostics,
			OmittedPaths: []string{"$.diagnostics"},
		})
		diagnostics = diagnostics[:maxDiagnostics]
	}
	out := make([]Diagnostic, 0, len(diagnostics))
	for index, diagnostic := range diagnostics {
		// Identity fields (code/path/field) are validated byte-for-byte and fail
		// closed: a malformed/spaced value is never trimmed into an accepted
		// identity that unlocks public diagnostics.
		if _, ok := run.PublicFailureCode(diagnostic.Code); !ok {
			report = mergeBudgetReports(report, resultbudget.Report{
				Truncated: true, BudgetKind: "tool_result",
				OmittedCount: 1, OmittedPaths: []string{fmt.Sprintf("$.diagnostics[%d]", index)},
			})
			continue
		}
		if diagnostic.Path != "" {
			if _, ok := run.PublicFailurePath(diagnostic.Path); !ok {
				report = mergeBudgetReports(report, resultbudget.Report{
					Truncated: true, BudgetKind: "tool_result",
					OmittedCount: 1, OmittedPaths: []string{fmt.Sprintf("$.diagnostics[%d].path", index)},
				})
				continue
			}
		}
		if diagnostic.Field != "" {
			if _, ok := run.PublicDiagnosticField(diagnostic.Field); !ok {
				report = mergeBudgetReports(report, resultbudget.Report{
					Truncated: true, BudgetKind: "tool_result",
					OmittedCount: 1, OmittedPaths: []string{fmt.Sprintf("$.diagnostics[%d].field", index)},
				})
				continue
			}
		}
		bounded := Diagnostic{
			Severity: diagnostic.Severity,
			Code:     diagnostic.Code,
			Path:     diagnostic.Path,
			Field:    diagnostic.Field,
			Index:    diagnostic.Index,
			Line:     diagnostic.Line,
			Column:   diagnostic.Column,
			NodeID:   diagnostic.NodeID,
			EdgeID:   diagnostic.EdgeID,
			Message:  diagnostic.Message,
		}
		var fieldReport resultbudget.Report
		pathPrefix := fmt.Sprintf("$.diagnostics[%d]", index)
		bounded.Severity, fieldReport = budgetErrorField(bounded.Severity, maxDiagnosticFieldRunes)
		report = mergeBudgetReports(report, fieldReport)
		bounded.Code, fieldReport = budgetErrorField(bounded.Code, maxDiagnosticFieldRunes)
		report = mergeBudgetReports(report, fieldReport)
		bounded.Path, fieldReport = budgetErrorField(bounded.Path, maxPathRunes)
		report = mergeBudgetReports(report, fieldReport)
		bounded.Field, fieldReport = budgetErrorField(bounded.Field, maxDiagnosticFieldRunes)
		report = mergeBudgetReports(report, fieldReport)
		bounded.NodeID, fieldReport = budgetErrorField(bounded.NodeID, maxDiagnosticFieldRunes)
		report = mergeBudgetReports(report, fieldReport)
		bounded.EdgeID, fieldReport = budgetErrorField(bounded.EdgeID, maxDiagnosticFieldRunes)
		report = mergeBudgetReports(report, fieldReport)
		bounded.Message, fieldReport = budgetErrorField(bounded.Message, maxProblemRunes)
		if fieldReport.Truncated {
			report = mergeBudgetReports(report, resultbudget.Report{
				Truncated: true, BudgetKind: "tool_result",
				OmittedCount: 1, OmittedPaths: []string{pathPrefix + ".message"},
			})
		}
		out = append(out, bounded)
	}
	return out, report
}

func validationProblem(err *Error) string {
	if err == nil {
		return "invalid input"
	}
	cause := strings.TrimSpace(fmt.Sprint(err.Cause))
	if cause == "" {
		return strings.TrimSpace(err.Problem)
	}
	for _, prefix := range []string{"exec: invalid input: ", "instance fails schema: "} {
		cause = strings.TrimPrefix(cause, prefix)
	}
	return "invalid input: " + cause
}

func budgetErrorField(value string, maxRunes int) (string, resultbudget.Report) {
	value = strings.TrimSpace(strings.ToValidUTF8(value, "�"))
	limits := resultbudget.DefaultLimits()
	limits.MaxTextRunes = maxRunes
	limits.MaxTextBlockRunes = maxRunes
	return resultbudget.Text(value, limits)
}

func mergeBudgetReports(a, b resultbudget.Report) resultbudget.Report {
	if !a.Truncated {
		return b
	}
	if !b.Truncated {
		return a
	}
	a.OmittedCount += b.OmittedCount
	a.OmittedPaths = append(a.OmittedPaths, b.OmittedPaths...)
	if len(a.OmittedPaths) > 20 {
		a.OmittedPaths = a.OmittedPaths[:20]
	}
	return a
}