package proxy

import (
	"sync"
	"time"
)

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

const (
	// affinityShards 分片数（2 的幂，取低位做下标）。
	affinityShards = 16
	// affinityPerShard 每片的条目上限：超过先清过期，仍满则不再新增（防御性）。
	affinityPerShard = 8192
	// affinityPrefix 取请求体的前多少字节参与摘要。
	affinityPrefix = 512
)

type affinityTable struct {
	shards [affinityShards]affinityShard
}

type affinityShard struct {
	mu sync.Mutex
	m  map[uint64]affinityEntry
}

type affinityEntry struct {
	channelID uint64
	until     time.Time
}

func newAffinityTable() *affinityTable {
	t := &affinityTable{}
	for i := range t.shards {
		t.shards[i].m = map[uint64]affinityEntry{}
	}
	return t
}

func (t *affinityTable) shard(key uint64) *affinityShard {
	return &t.shards[key&(affinityShards-1)]
}

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

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

// len 当前条目数（测试与诊断用）。
func (t *affinityTable) len() int {
	n := 0
	for i := range t.shards {
		t.shards[i].mu.Lock()
		n += len(t.shards[i].m)
		t.shards[i].mu.Unlock()
	}
	return n
}

// affinityKey 会话身份摘要：密钥 + 模型 + 请求体前 affinityPrefix 字节。
// 内联 FNV-1a，避免在热路径上分配哈希器或拼接字符串。
func affinityKey(keyID uint64, model string, prefix []byte) uint64 {
	const (
		offset64 = 14695981039346656037
		prime64  = 1099511628211
	)
	h := uint64(offset64)
	mix := func(b byte) {
		h ^= uint64(b)
		h *= prime64
	}
	for i := 0; i < 8; i++ {
		mix(byte(keyID >> (8 * i)))
	}
	for i := 0; i < len(model); i++ {
		mix(model[i])
	}
	if len(prefix) > affinityPrefix {
		prefix = prefix[:affinityPrefix]
	}
	for _, b := range prefix {
		mix(b)
	}
	return h
}