package runtime

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

	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"math"
	"strings"
)

// Keep usage exact across the JSON/float64 consumers between MeshX, AgentTeam, and
// signed metrics. Values and aggregates outside this range use accounting
// recovery instead of being rounded.
const maxInferenceUsageSafeInteger int64 = 1<<53 - 1

func enforceInferenceAccountingContract(
	result InferenceResult,
	childIdempotencyKey string,
	inferenceErr error,
) (InferenceResult, error) {
	result, accountingErr := normalizeInferenceAccounting(result, childIdempotencyKey)
	return result, prioritizeInferenceAccountingError(inferenceErr, accountingErr)
}

func prioritizeInferenceAccountingError(inferenceErr, accountingErr error) error {
	if accountingErr == nil {
		return inferenceErr
	}
	if inferenceErr == nil {
		return accountingErr
	}
	return &Error{
		Code:    CodeMaaSAccountingRecovery,
		Message: "Knowledge Q&A model accounting result is outside its legal contract",
		Cause:   errors.Join(accountingErr, inferenceErr),
	}
}

// normalizeInferenceAccounting is the producer boundary for the only
// accounting data allowed to leave an Inference implementation. Usage is
// projected onto a fixed numeric vocabulary, so arbitrary provider fields can
// never grow the attempt aggregate. An invalid recognized value or an
// oversized settlement reference terminates the attempt through accounting
// recovery after returning a bounded, durable projection.
func normalizeInferenceAccounting(
	result InferenceResult,
	childIdempotencyKey string,
) (InferenceResult, error) {
	usage, usageValid := normalizeInferenceUsage(result.Usage)
	if encoded, err := json.Marshal(usage); err != nil || len(encoded) > protocol.MaxAccountingUsageJSONBytes {
		usage = nil
		usageValid = false
	}
	result.Usage = usage
	result.SettlementRef = strings.TrimSpace(result.SettlementRef)
	settlementValid := len(result.SettlementRef) <= protocol.MaxAccountingSettlementReferenceBytes
	if !settlementValid {
		result.SettlementRef = accountingRecoveryReference(childIdempotencyKey)
	}
	if usageValid && settlementValid {
		return result, nil
	}
	return result, maasAccountingRecovery()
}

func accountingRecoveryReference(childIdempotencyKey string) string {
	digest := sha256.Sum256([]byte(strings.TrimSpace(childIdempotencyKey)))
	return "kqa_request_sha256_" + hex.EncodeToString(digest[:])
}

func normalizeInferenceUsage(usage map[string]any) (map[string]any, bool) {
	if len(usage) == 0 {
		return nil, true
	}
	result := make(map[string]any)
	valid := true
	for key, value := range usage {
		switch {
		case inferenceUsageScalarAllowed(key):
			number, ok := normalizeInferenceUsageNumber(value)
			if !ok {
				valid = false
				continue
			}
			result[key] = number
		case inferenceUsageDetailsAllowed(key):
			details, ok := value.(map[string]any)
			if !ok {
				valid = false
				continue
			}
			normalized := make(map[string]any)
			for detailKey, detailValue := range details {
				if !inferenceUsageDetailScalarAllowed(detailKey) {
					continue
				}
				number, ok := normalizeInferenceUsageNumber(detailValue)
				if !ok {
					valid = false
					continue
				}
				normalized[detailKey] = number
			}
			if len(normalized) > 0 {
				result[key] = normalized
			}
		}
	}
	if len(result) == 0 {
		result = nil
	}
	return result, valid
}

func normalizeInferenceUsageNumber(value any) (float64, bool) {
	switch number := value.(type) {
	case int:
		if number < 0 || uint64(number) > uint64(maxInferenceUsageSafeInteger) {
			return 0, false
		}
		return float64(number), true
	case int64:
		if number < 0 || number > maxInferenceUsageSafeInteger {
			return 0, false
		}
		return float64(number), true
	case float64:
		if math.IsNaN(number) || math.IsInf(number, 0) || number < 0 ||
			number > float64(maxInferenceUsageSafeInteger) || math.Trunc(number) != number {
			return 0, false
		}
		return number, true
	case json.Number:
		integer, err := number.Int64()
		if err != nil || integer < 0 || integer > maxInferenceUsageSafeInteger {
			return 0, false
		}
		return float64(integer), true
	default:
		return 0, false
	}
}

func inferenceUsageScalarAllowed(key string) bool {
	switch key {
	case "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":
		return true
	default:
		return false
	}
}

func inferenceUsageDetailsAllowed(key string) bool {
	switch key {
	case "input_tokens_details", "output_tokens_details", "prompt_tokens_details", "completion_tokens_details":
		return true
	default:
		return false
	}
}

func inferenceUsageDetailScalarAllowed(key string) bool {
	switch key {
	case "cached_tokens", "audio_tokens", "reasoning_tokens", "accepted_prediction_tokens",
		"rejected_prediction_tokens", "cache_creation_tokens", "cache_read_tokens":
		return true
	default:
		return false
	}
}