package runtime

import (
	"github.com/tidbcloud/AgentTeam/internal/knowledgeqa/protocol"

	"encoding/json"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestInferenceAccountingProjectsArbitraryUsageOntoBoundedVocabulary(t *testing.T) {
	rawUsage := map[string]any{
		"input_tokens": 7,
		"prompt_tokens_details": map[string]any{
			"cached_tokens":           3,
			"provider_private_detail": strings.Repeat("x", protocol.MaxAccountingUsageJSONBytes),
		},
		"provider_private_payload": strings.Repeat("y", protocol.MaxAccountingUsageJSONBytes),
	}
	rawJSON, err := json.Marshal(rawUsage)
	require.NoError(t, err)
	require.Greater(t, len(rawJSON), protocol.MaxAccountingUsageJSONBytes)

	result, err := normalizeInferenceAccounting(InferenceResult{
		Usage: rawUsage, SettlementRef: "usage_1",
	}, "attempt-step-01")

	require.NoError(t, err)
	require.Equal(t, map[string]any{
		"input_tokens":          float64(7),
		"prompt_tokens_details": map[string]any{"cached_tokens": float64(3)},
	}, result.Usage)
	encoded, err := json.Marshal(result.Usage)
	require.NoError(t, err)
	require.LessOrEqual(t, len(encoded), protocol.MaxAccountingUsageJSONBytes)
}

func TestInferenceAccountingClosedVocabularyFitsUsageEnvelope(t *testing.T) {
	rawUsage := make(map[string]any)
	for _, key := range []string{
		"input_tokens", "output_tokens", "total_tokens", "prompt_tokens", "completion_tokens",
		"cached_tokens", "audio_tokens", "reasoning_tokens", "cache_creation_input_tokens",
		"cache_read_input_tokens", "prompt_cache_hit_tokens", "prompt_cache_miss_tokens",
	} {
		rawUsage[key] = maxInferenceUsageSafeInteger
	}
	for _, key := range []string{
		"input_tokens_details", "output_tokens_details", "prompt_tokens_details", "completion_tokens_details",
	} {
		rawUsage[key] = map[string]any{
			"cached_tokens":                maxInferenceUsageSafeInteger,
			"audio_tokens":                 maxInferenceUsageSafeInteger,
			"reasoning_tokens":             maxInferenceUsageSafeInteger,
			"accepted_prediction_tokens":   maxInferenceUsageSafeInteger,
			"rejected_prediction_tokens":   maxInferenceUsageSafeInteger,
			"cache_creation_tokens":        maxInferenceUsageSafeInteger,
			"cache_read_tokens":            maxInferenceUsageSafeInteger,
			"provider_private_usage_field": strings.Repeat("x", protocol.MaxAccountingUsageJSONBytes),
		}
	}

	result, err := normalizeInferenceAccounting(InferenceResult{
		Usage: rawUsage, SettlementRef: "usage_1",
	}, "attempt-step-01")

	require.NoError(t, err)
	encoded, err := json.Marshal(result.Usage)
	require.NoError(t, err)
	require.LessOrEqual(t, len(encoded), protocol.MaxAccountingUsageJSONBytes)
	require.Len(t, result.Usage, 16)
	require.Equal(t, float64(maxInferenceUsageSafeInteger), result.Usage["input_tokens"])
	require.Contains(t, string(encoded), "9007199254740991")
}

func TestInferenceAccountingAcceptsExactSafeIntegerBoundary(t *testing.T) {
	for _, test := range []struct {
		name  string
		value any
	}{
		{name: "int64", value: maxInferenceUsageSafeInteger},
		{name: "float64", value: float64(maxInferenceUsageSafeInteger)},
		{name: "json number", value: json.Number("9007199254740991")},
	} {
		t.Run(test.name, func(t *testing.T) {
			result, err := normalizeInferenceAccounting(InferenceResult{
				Usage: map[string]any{"input_tokens": test.value}, SettlementRef: "usage_1",
			}, "attempt-step-01")

			require.NoError(t, err)
			require.Equal(t, float64(maxInferenceUsageSafeInteger), result.Usage["input_tokens"])
			encoded, err := json.Marshal(result.Usage)
			require.NoError(t, err)
			require.JSONEq(t, `{"input_tokens":9007199254740991}`, string(encoded))
		})
	}
}

func TestInferenceAccountingRejectsSafeIntegerBoundaryPlusOne(t *testing.T) {
	for _, test := range []struct {
		name  string
		value any
	}{
		{name: "int64", value: maxInferenceUsageSafeInteger + 1},
		{name: "float64", value: float64(maxInferenceUsageSafeInteger + 1)},
		{name: "json number", value: json.Number("9007199254740992")},
	} {
		t.Run(test.name, func(t *testing.T) {
			result, err := normalizeInferenceAccounting(InferenceResult{
				Usage: map[string]any{
					"input_tokens": test.value, "output_tokens": 3,
				},
				SettlementRef: "usage_1",
			}, "attempt-step-01")

			var runtimeErr *Error
			require.ErrorAs(t, err, &runtimeErr)
			require.Equal(t, CodeMaaSAccountingRecovery, runtimeErr.Code)
			require.Equal(t, map[string]any{"output_tokens": float64(3)}, result.Usage)
		})
	}
}

func TestInferenceAccountingUsesBoundedRecoveryReferenceForOversizedProducerRef(t *testing.T) {
	rawRef := strings.Repeat("private-usage-ref", 32)
	require.Greater(t, len(rawRef), protocol.MaxAccountingSettlementReferenceBytes)

	result, err := normalizeInferenceAccounting(InferenceResult{
		Usage: map[string]any{"input_tokens": 7}, SettlementRef: rawRef,
	}, "attempt-step-01")

	var runtimeErr *Error
	require.ErrorAs(t, err, &runtimeErr)
	require.Equal(t, CodeMaaSAccountingRecovery, runtimeErr.Code)
	require.Equal(t, accountingRecoveryReference("attempt-step-01"), result.SettlementRef)
	require.LessOrEqual(t, len(result.SettlementRef), protocol.MaxAccountingSettlementReferenceBytes)
	require.NotContains(t, result.SettlementRef, rawRef)
	require.Equal(t, map[string]any{"input_tokens": float64(7)}, result.Usage)
}

func TestInferenceAccountingRejectsInvalidRecognizedUsageWithDurableProjection(t *testing.T) {
	result, err := normalizeInferenceAccounting(InferenceResult{
		Usage: map[string]any{
			"input_tokens":  strings.Repeat("7", protocol.MaxAccountingUsageJSONBytes),
			"output_tokens": 3,
		},
		SettlementRef: "usage_1",
	}, "attempt-step-01")

	var runtimeErr *Error
	require.ErrorAs(t, err, &runtimeErr)
	require.Equal(t, CodeMaaSAccountingRecovery, runtimeErr.Code)
	require.Equal(t, "usage_1", result.SettlementRef)
	require.Equal(t, map[string]any{"output_tokens": float64(3)}, result.Usage)
}