package proxy

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"sync"
	"time"
)

// 会话亲和（借鉴 magpie 的 cache-economics 思路，做服务端可用的简化版）：
// 同一会话粘住上次应答的渠道，让上游的提示词缓存（整段上下文每轮重发）被再次
// 命中，而不是被加权随机打散后全价重算。
//
// 会话身份 = 密钥 + 模型 + 请求文本稳定前缀的摘要。编码代理每轮重发全量上下文，
// 系统提示与最早的用户消息不变，取前缀即可区分会话；不需要解析各家客户端的
// 会话头。命中亲和的渠道仍必须在候选集里（健康、未冷却、认识该模型）——
// 冷却或下线后自然回退到加权随机。
//
// TTL 默认 5 分钟（对齐 Anthropic 提示词缓存窗口），每次成功应答刷新；
// 上游不报缓存的（OpenAI 自动前缀缓存）同样受益：粘住就是省。

const (
	// affinityMax 亲和表上限：超过先清过期，仍超则不再新增（防御性，正常流量到不了）。
	affinityMax = 8192
	// affinityPrefix 会话身份摘要取请求文本的前缀字节数。
	affinityPrefix = 512
)

type affinityTable struct {
	mu sync.Mutex
	m  map[string]affinityEntry
}

type affinityEntry struct {
	channelID uint64
	until     time.Time
}

func newAffinityTable() *affinityTable {
	return &affinityTable{m: map[string]affinityEntry{}}
}

// get 返回该会话上次应答的渠道；过期条目顺手清理。
func (t *affinityTable) get(key string) (uint64, bool) {
	t.mu.Lock()
	defer t.mu.Unlock()
	e, ok := t.m[key]
	if !ok {
		return 0, false
	}
	if time.Now().After(e.until) {
		delete(t.m, key)
		return 0, false
	}
	return e.channelID, true
}

// put 记录（或刷新）会话→渠道的绑定。
func (t *affinityTable) put(key string, channelID uint64, ttl time.Duration) {
	t.mu.Lock()
	defer t.mu.Unlock()
	if len(t.m) >= affinityMax {
		now := time.Now()
		for k, e := range t.m {
			if now.After(e.until) {
				delete(t.m, k)
			}
		}
		if len(t.m) >= affinityMax {
			return // 表满且全是活的：放弃本次记录，旧绑定继续服务
		}
	}
	t.m[key] = affinityEntry{channelID: channelID, until: time.Now().Add(ttl)}
}

// affinityKey 会话身份：密钥 + 模型 + 请求文本稳定前缀的摘要。
func affinityKey(keyID uint64, model, firstText string) string {
	if len(firstText) > affinityPrefix {
		firstText = firstText[:affinityPrefix]
	}
	h := sha256.Sum256([]byte(firstText))
	return fmt.Sprintf("%d|%s|%s", keyID, model, hex.EncodeToString(h[:8]))
}