package memory

import (
	"context"
	"math"
	"sort"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/embedding"
)

// 本文件是记忆整理的**候选召回**（design/23 §5.1 T0-2）：分区内两两相似度
// ≥ recallThreshold 的对进入候选。
//
// ⚠️ 口径（不可走样）：相似度**只用来召回候选对**，永不直接当合并判据。
// design/23 §1.2 的实据是「必合」的 [12]/[15]=0.8330 与「必不合」的
// [3]/[6]=0.7100 只差 0.12，且全库最高相似度的一对（0.9319）恰是跨作用域
// ⇒ 单阈值双向失效，判定必须靠包含判定 + 规则否决（tidy_rules.go）。

// tidyPair 是一条被召回的候选对。
type tidyPair struct {
	a, b  MemoryItem
	score float64
	// mode 记录本轮相似度来源（embedding / lexical），进证据与日志。
	mode string
}

// recallPairs 在分区内召回候选对。条目少于 2 条时为空操作（无候选、不报错）。
//
// 两两比对是 O(n²·dim)：设计稿每分区扫描上限 5000 条（design §9），但真正
// 参与比对的最多 maxTidyCompareItems 条（最新优先），超限记 WARN 并登记到
// stats.Truncated —— 截断绝不静默。
//
// 返回的对按「分数降序 + id」确定性排序，使每分区候选上限的截断可复现
// （同一批数据两次运行产出同一批候选 ⇒ 幂等，design §4.1 A）。
func (t *Tidier) recallPairs(ctx context.Context, p tidyPartition, items []MemoryItem, stats *TidyStats) []tidyPair {
	if len(items) < 2 {
		return nil
	}
	if len(items) > maxTidyCompareItems {
		stats.Truncated = append(stats.Truncated, p.label()+"(compare)")
		t.log.Warn("memory tidy compare truncated by per-iteration guard",
			"partition", p.label(), "items", len(items), "limit", maxTidyCompareItems)
		items = items[:maxTidyCompareItems]
	}
	stats.Compared += len(items)

	sims, mode := t.similarityFunc(ctx, p, items)
	pairs := make([]tidyPair, 0)
	for i := 0; i < len(items); i++ {
		for j := i + 1; j < len(items); j++ {
			score := sims(i, j)
			if score < t.cfg.RecallThreshold {
				continue
			}
			pairs = append(pairs, tidyPair{a: items[i], b: items[j], score: score, mode: mode})
		}
	}
	stats.Pairs += len(pairs)
	sort.Slice(pairs, func(i, j int) bool {
		if pairs[i].score != pairs[j].score {
			return pairs[i].score > pairs[j].score
		}
		if pairs[i].a.ID != pairs[j].a.ID {
			return pairs[i].a.ID < pairs[j].a.ID
		}
		return pairs[i].b.ID < pairs[j].b.ID
	})
	return pairs
}

// similarityFunc 返回分区内的两两相似度函数与来源标签。
//
// 接入 embedder 时用它批量向量化后取余弦（与检索侧同一口径：余弦夹到
// [0,1]）。embedder 缺席或整批 Embed 失败时回退到零依赖的词项余弦
// —— 该回退保证整理在离线/无 embedding 配置时仍可跑（宁可比语义路径更保守），
// 且让测试无需网络即可确定性运行。两个来源的相似度**不可混用阈值**：
// 阈值只有一个（design §13 开放点 2 的 0.80），回退路径命中更少 ⇒ 更保守。
func (t *Tidier) similarityFunc(ctx context.Context, p tidyPartition, items []MemoryItem) (func(i, j int) float64, string) {
	if t.deps.Embedder != nil {
		vecs, err := t.embedItems(ctx, items)
		if err == nil && len(vecs) == len(items) {
			return func(i, j int) float64 {
				return clampUnit(cosineSim(vecs[i], vecs[j]))
			}, "embedding"
		}
		t.log.Warn("memory tidy embedding recall unavailable, falling back to lexical",
			"partition", p.label(), "error", err)
	}
	return func(i, j int) float64 {
		return lexicalSimilarity(items[i].Content, items[j].Content)
	}, "lexical"
}

// embedItems 分批把分区内条目向量化。单批失败即整体失败（调用方回退到
// 词项余弦），不半途改用混合表示 —— 混用会让召回范围不可预测。
func (t *Tidier) embedItems(ctx context.Context, items []MemoryItem) ([]embedding.Vector, error) {
	out := make([]embedding.Vector, 0, len(items))
	for start := 0; start < len(items); start += tidyEmbedBatch {
		end := min(start+tidyEmbedBatch, len(items))
		texts := make([]string, 0, end-start)
		for _, it := range items[start:end] {
			texts = append(texts, it.Content)
		}
		vecs, err := t.deps.Embedder.Embed(ctx, texts)
		if err != nil {
			return nil, err
		}
		if len(vecs) != len(texts) {
			return nil, ErrTidyEmbedCount
		}
		out = append(out, vecs...)
	}
	return out, nil
}

// lexicalSimilarity 是零依赖的相似度回退：两条内容词项集合的余弦。
// 词项来自 extractTidyFacts（CJK 二元组 + 拉丁词 + 数值 + emoji 符号），
// 因此「同义不同措辞」的中文近似表述仍能召回，而「红牛 vs 冰美式」
// 这类并列事实只共享极少词项（相似度低或中等），靠包含判定兜住。
func lexicalSimilarity(a, b string) float64 {
	fa, fb := extractTidyFacts(a), extractTidyFacts(b)
	if len(fa.terms) == 0 || len(fb.terms) == 0 {
		return 0
	}
	shared := 0
	for term := range fa.terms {
		if _, ok := fb.terms[term]; ok {
			shared++
		}
	}
	return float64(shared) / math.Sqrt(float64(len(fa.terms))*float64(len(fb.terms)))
}

// clampUnit 把余弦夹到 [0,1]：相似度是「召回门」的输入，负相关对
// （余弦 < 0）与零相关对在召回上等价。
func clampUnit(v float64) float64 {
	if v < 0 {
		return 0
	}
	if v > 1 {
		return 1
	}
	return v
}