package runner

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"strings"
)

type cacheAffinityKeys struct {
	Primary string
	Lookup  []string
	Record  []string
}

func buildCacheAffinityKey(kind string, modelType string, body map[string]any) string {
	return buildCacheAffinityKeys(kind, modelType, body).Primary
}

func buildCacheAffinityKeys(kind string, modelType string, body map[string]any) cacheAffinityKeys {
	if len(body) == 0 {
		return cacheAffinityKeys{}
	}
	for _, key := range []string{"cacheAffinityKey", "cache_affinity_key"} {
		if value := strings.TrimSpace(stringFromAny(body[key])); value != "" {
			key := "explicit:" + sha256Text(value)
			return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}}
		}
	}
	for _, key := range []string{"sessionId", "session_id", "conversationId", "conversation_id"} {
		if value := strings.TrimSpace(stringFromAny(body[key])); value != "" {
			key := key + ":" + sha256Text(value)
			return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}}
		}
	}
	if !cacheAffinityPromptHashSupported(kind, modelType) {
		return cacheAffinityKeys{}
	}
	basePayload := map[string]any{
		"kind":      kind,
		"modelType": modelType,
	}
	if tools, ok := body["tools"]; ok {
		basePayload["tools"] = tools
	}
	if instructions, ok := body["instructions"]; ok {
		basePayload["instructions"] = instructions
	}
	messageHashes := cacheAffinityMessageHashes(body)
	if len(messageHashes) > 0 {
		baseHash := stableJSONHash(basePayload)
		keys := make([]string, 0, len(messageHashes))
		for index := range messageHashes {
			keyPayload := map[string]any{
				"base":          baseHash,
				"messageHashes": messageHashes[:index+1],
				"prefixLength":  index + 1,
			}
			keys = append(keys, "prompt_lcp:"+stableJSONHash(keyPayload))
		}
		if len(keys) == 0 {
			return cacheAffinityKeys{}
		}
		lookup := make([]string, 0, len(keys))
		for index := len(keys) - 1; index >= 0; index-- {
			lookup = append(lookup, keys[index])
		}
		return cacheAffinityKeys{
			Primary: keys[len(keys)-1],
			Lookup:  lookup,
			Record:  keys,
		}
	}
	if input, ok := body["input"]; ok {
		basePayload["input"] = input
	}
	if len(basePayload) <= 2 {
		return cacheAffinityKeys{}
	}
	key := "prompt:" + stableJSONHash(basePayload)
	return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}}
}

func cacheAffinityPromptHashSupported(kind string, modelType string) bool {
	switch strings.TrimSpace(kind) {
	case "chat.completions", "responses":
		return true
	}
	switch strings.TrimSpace(modelType) {
	case "chat", "text_generate", "responses":
		return true
	default:
		return false
	}
}

func cacheAffinityMessageHashes(body map[string]any) []string {
	if hashes := stringListFromAny(body["messageHashes"]); len(hashes) > 0 {
		return hashes
	}
	if hashes := stringListFromAny(body["message_hashes"]); len(hashes) > 0 {
		return hashes
	}
	messages, ok := body["messages"].([]any)
	if !ok || len(messages) == 0 {
		return nil
	}
	hashes := make([]string, 0, len(messages))
	for _, message := range messages {
		raw, err := json.Marshal(message)
		if err != nil || len(raw) == 0 {
			continue
		}
		hashes = append(hashes, sha256Text(string(raw)))
	}
	return hashes
}

func stableJSONHash(value any) string {
	raw, err := json.Marshal(value)
	if err != nil || len(raw) == 0 {
		return ""
	}
	return sha256Text(string(raw))
}

func sha256Text(value string) string {
	sum := sha256.Sum256([]byte(value))
	return hex.EncodeToString(sum[:])
}