package service

import (
	"crypto/sha256"
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/dto"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/pkg/cachex"
	"github.com/gin-gonic/gin"
	"github.com/samber/hot"
)

type poolBinding struct {
	PolicyFingerprint string `json:"policy_fingerprint"`
	Channel           int    `json:"channel"`
	Group             string `json:"group"`
}

var poolBindingOnce sync.Once
var poolBindingCache *cachex.HybridCache[poolBinding]

func poolBindings() *cachex.HybridCache[poolBinding] {
	poolBindingOnce.Do(func() {
		poolBindingCache = cachex.NewHybridCache[poolBinding](cachex.HybridCacheConfig[poolBinding]{Namespace: cachex.Namespace("new-api:pool-affinity:v1"), Redis: common.RDB, RedisEnabled: func() bool { return common.RedisEnabled && common.RDB != nil }, RedisCodec: cachex.JSONCodec[poolBinding]{}, Memory: func() *hot.HotCache[string, poolBinding] {
			return hot.NewHotCache[string, poolBinding](hot.LRU, 100000).WithTTL(time.Hour).Build()
		}})
	})
	return poolBindingCache
}

func PreferredPoolChannel(c *gin.Context, modelName string) (*model.Channel, string) {
	// Resolve the existing session-key rules, but isolate the binding by owner,
	// token and original model rather than the currently selected group.
	GetPreferredChannelByAffinity(c, modelName, "pool")
	meta, ok := getChannelAffinityMeta(c)
	if !ok || RequestPolicy(c).SessionMode == "off" {
		return nil, ""
	}
	key := common.GenerateHMAC(fmt.Sprintf("%d:%d:%d:%s:%s", c.GetInt("id"), c.GetInt("token_id"), len(modelName), modelName, meta.CacheKey))
	c.Set("pool_affinity_key", key)
	c.Set("pool_affinity_ttl", meta.TTLSeconds)
	// Snapshot the routing policy before selection. An older in-flight request
	// may finish after a manual switch, but its binding must not override the
	// new policy on the next request.
	policy := RequestRoutingPolicy(c)
	var fingerprintPool []model.TokenGroupEntry
	if policy.Mode != "smart" {
		fingerprintPool = ResolveRequestPool(c, modelName)
	}
	// Automatic rankings change with live metrics and availability. They must
	// not invalidate a healthy session binding; explicit policy edits still do.
	policyJSON, err := common.Marshal(struct {
		Policy *model.TokenRoutingPolicy
		Pool   []model.TokenGroupEntry
	}{policy, fingerprintPool})
	if err != nil {
		return nil, ""
	}
	policyFingerprint := fmt.Sprintf("%x", sha256.Sum256(policyJSON))
	c.Set("pool_affinity_policy_fingerprint", policyFingerprint)
	binding, found, err := poolBindings().Get(key)
	if err != nil {
		return nil, ""
	}
	if !found {
		// Adopt only a successful old binding attributable to this exact owner,
		// token, model and session. Never trust a legacy unscoped cache alone.
		appendRequestDetails(c, nil, model.NewLogOther())
		correlation := c.GetString("log_correlation_key")
		if correlation == "" || model.LOG_DB == nil {
			return nil, ""
		}
		var previous model.Log
		if model.LOG_DB.Where("user_id = ? AND token_id = ? AND model_name = ? AND correlation_key = ? AND type = ? AND created_at > ?", c.GetInt("id"), c.GetInt("token_id"), modelName, correlation, model.LogTypeConsume, time.Now().Unix()-int64(meta.TTLSeconds)).Order("id desc").First(&previous).Error != nil {
			return nil, ""
		}
		var other struct {
			Admin struct {
				Details struct {
					PolicyFingerprint string `json:"routing_policy_fingerprint"`
				} `json:"request_details"`
			} `json:"admin_info"`
			Result string `json:"model_guard_result"`
			Stream struct {
				Status    string `json:"status"`
				EndReason string `json:"end_reason"`
				Response  string `json:"response_status"`
			} `json:"stream_status"`
		}
		_ = common.UnmarshalJsonStr(previous.Other, &other)
		if other.Result != "" {
			return nil, ""
		}
		if previous.IsStream && (other.Stream.Status != "ok" || (other.Stream.EndReason != "done" && other.Stream.EndReason != "eof") || (other.Stream.Response != "" && other.Stream.Response != "completed")) {
			return nil, ""
		}
		binding = poolBinding{Channel: previous.ChannelId, Group: previous.Group, PolicyFingerprint: other.Admin.Details.PolicyFingerprint}
	}
	// Unversioned legacy bindings cannot prove they belong to this policy.
	// This check also prevents historical logs from restoring an old route.
	if binding.PolicyFingerprint != policyFingerprint {
		ClearPoolAffinity(c)
		return nil, ""
	}
	for _, entry := range ResolveRequestPool(c, modelName) {
		if entry.Group != binding.Group {
			continue
		}
		if requestGroupDegraded(c, modelName, entry.Group) {
			ClearPoolAffinity(c)
			recordDirectSelection(c, nil, binding.Channel, entry.Group, modelName, "pool_group_degraded")
			return nil, ""
		}
		if poolGroupSkipped(c, modelName, entry.Group) || (policy.Mode != "smart" && poolHigherPriorityRecovered(c, modelName, entry.Group)) {
			ClearPoolAffinity(c)
			return nil, ""
		}
		channel, err := model.CacheGetChannel(binding.Channel)
		if err == nil && channel.Status == common.ChannelStatusEnabled && model.IsChannelEnabledForGroupModel(entry.Group, modelName, channel.Id) {
			if ok, _ := model.ChannelSatisfiesFilters(channel, modelName, GetChannelConstraints(c).Filters); ok {
				if !claimPoolGroup(c, modelName, entry.Group) {
					return nil, ""
				}
				common.SetContextKey(c, constant.ContextKeyAutoGroup, entry.Group)
				MarkChannelAffinityUsed(c, entry.Group, channel.Id)
				recordDirectSelection(c, channel, channel.Id, entry.Group, modelName, "pool_affinity_hit")
				return channel, entry.Group
			}
		}
	}
	ClearPoolAffinity(c)
	return nil, ""
}

func ClearPoolAffinity(c *gin.Context) {
	if key := c.GetString("pool_affinity_key"); key != "" {
		ttl := c.GetInt("pool_affinity_ttl")
		if ttl <= 0 {
			ttl = 3600
		}
		// A tombstone prevents re-adopting the same historical binding after a
		// rejection. A later successful generation replaces it normally.
		_ = poolBindings().SetWithTTL(key, poolBinding{}, time.Duration(ttl)*time.Second)
	}
	ClearCurrentChannelAffinityCache(c)
}

func RecordPoolAffinity(c *gin.Context) {
	if !RequestPolicy(c).Successful {
		return
	}
	RecordPoolSuccess(c, c.GetString("original_model"), common.GetContextKeyString(c, constant.ContextKeyAutoGroup))
	key := c.GetString("pool_affinity_key")
	if key == "" {
		return
	}
	ttl := c.GetInt("pool_affinity_ttl")
	if ttl <= 0 {
		ttl = 3600
	}
	_ = poolBindings().SetWithTTL(key, poolBinding{Channel: c.GetInt("channel_id"), Group: common.GetContextKeyString(c, constant.ContextKeyAutoGroup), PolicyFingerprint: c.GetString("pool_affinity_policy_fingerprint")}, time.Duration(ttl)*time.Second)
}

func RequestRoutingPolicy(c *gin.Context) *model.TokenRoutingPolicy {
	value, _ := c.Get(model.TokenRoutingContextKey)
	policy, _ := value.(*model.TokenRoutingPolicy)
	return policy
}

func ResolveRequestPool(c *gin.Context, modelName string) []model.TokenGroupEntry {
	policy := RequestRoutingPolicy(c)
	if policy == nil {
		return nil
	}
	userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
	if policy.Mode == "smart" {
		entries := WithImageGenerationGroup(c, modelName, resolveAutomaticPool(c, modelName))
		if len(policy.ExcludedGroups) == 0 {
			return entries
		}
		filtered := make([]model.TokenGroupEntry, 0, len(entries))
		for _, entry := range entries {
			if !policy.ExcludesGroup(entry.Group) {
				filtered = append(filtered, entry)
			}
		}
		return filtered
	}
	inherited := GetUserAutoGroup(userGroup)
	if policy.Mode == "user" {
		inherited = []string{userGroup}
	}
	return WithImageGenerationGroup(c, modelName, policy.Resolve(modelName, inherited, func(group string) bool { return IsUserSelectableGroup(userGroup, group) }))
}

func selectTokenRoutingPool(param *RetryParam) (*model.Channel, string, error) {
	entries := ResolveRequestPool(param.Ctx, param.ModelName)
	constraints := GetChannelConstraints(param.Ctx)
	filters := append([]dto.ChannelFilter{}, constraints.Filters...)
	var excluded []int
	for _, value := range param.Ctx.GetStringSlice("use_channel") {
		id, _ := strconv.Atoi(value)
		excluded = append(excluded, id)
	}
	filters = append(filters, dto.ChannelFilter{Kind: dto.FilterAttemptedChannel, ExcludedChannelIDs: excluded})
	// Retry once in the current group, including single-channel groups. Preserve
	// every security/protocol filter; remove only the request's attempted list.
	if retryGroup := param.Ctx.GetString("pool_retry_group"); retryGroup != "" && poolFailures(param.Ctx)[retryGroup] == 1 && !poolGroupSkipped(param.Ctx, param.ModelName, retryGroup) && !requestGroupDegraded(param.Ctx, param.ModelName, retryGroup) {
		param.Ctx.Set("pool_retry_group", "")
		for _, entry := range entries {
			if entry.Group != retryGroup {
				continue
			}
			channel, err := model.GetRandomSatisfiedChannel(retryGroup, param.ModelName, 0, constraints.Filters)
			if err != nil {
				return nil, retryGroup, err
			}
			if channel != nil {
				common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroup, retryGroup)
				recordDirectSelection(param.Ctx, channel, channel.Id, retryGroup, param.ModelName, "pool_same_group_retry")
				return channel, retryGroup, nil
			}
		}
	}
	for start := 0; start < len(entries); {
		end := start + 1
		for end < len(entries) && entries[end].Priority == entries[start].Priority {
			end++
		}
		groups := entries[start:end]
		var names []string
		for _, entry := range groups {
			names = append(names, entry.Group)
		}
		offset := nextTierIndex(fmt.Sprintf("pool:%d:%s:%s", param.Ctx.GetInt("token_id"), param.ModelName, strings.Join(names, ",")), len(groups))
		for i := range len(groups) {
			group := groups[(offset+i)%len(groups)].Group
			if poolGroupSkipped(param.Ctx, param.ModelName, group) {
				recordDirectSelection(param.Ctx, nil, 0, group, param.ModelName, "pool_group_cooldown")
				continue
			}
			if requestGroupDegraded(param.Ctx, param.ModelName, group) {
				recordDirectSelection(param.Ctx, nil, 0, group, param.ModelName, "pool_group_degraded")
				continue
			}
			trace := &model.SelectionTrace{}
			channel, err := model.GetRandomSatisfiedChannel(group, param.ModelName, 0, filters, trace)
			if err != nil {
				return nil, group, err
			}
			if channel != nil {
				if !claimPoolGroup(param.Ctx, param.ModelName, group) {
					continue
				}
				common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroup, group)
				recordDirectSelection(param.Ctx, channel, channel.Id, group, param.ModelName, "model_pool_priority")
				return channel, group, nil
			}
		}
		start = end
	}
	if param.Ctx.GetBool("pool_degradation_unavailable") {
		return nil, param.TokenGroup, fmt.Errorf("group quality temporarily unavailable")
	}
	return nil, param.TokenGroup, nil
}