package memory

import (
	"context"
	"time"
)

// 本文件是记忆整理**时间预算闸**（切片 MT-A）：它是「上限」之外的第二道闸。
//
// 为什么需要它：抬高护栏（compare 400 → 2000）之后，一个巨型分区可以在墙钟上限
// 内烧掉几分钟。同一台开发机上实测（临时库、单分区、CompareLimit=2000）：
//
//	n=2000 语义比对（1024 维）≈ 23.5 秒   ← 纯 O(n²·dim) 计算
//	n=2000 词项余弦回退        ≈ 40.5 秒   ← 每对都重算词项集合，CPU 更贵
//
// 而真实的语义召回是**远端调用**：真机 89.9 秒 / 744 条 ⇒ 约 117 ms/条，2000 条
// 就是 234 秒，已经越过试算的 3 分钟墙钟上限。没有这道闸，结果是「整个请求超时、
// 一条候选都拿不到」；有它，本轮在 deadline 前收尾，把**已完成分区的结果**如实
// 交出去，并把这些分区标成 (budget) 截断。
//
// 它**不改变**心跳轮：心跳跑在进程生命周期 ctx 上（没有 deadline ⇒ 没有预算），
// 只有带墙钟上限的调用方（同步 / 异步试算）才受它约束。

const (
	// tidyBudgetMargin 是时间预算相对 ctx deadline 的默认收尾余量：留出投影响应、
	// 写注册表的时间，使本轮被**预算**截断，而不是被硬超时打断（后者会连已算完的
	// 分区一起丢掉）。
	tidyBudgetMargin = 10 * time.Second

	// tidyBudgetMarginDivisor 是短预算下的余量折算：总预算本身短于 tidyBudgetMargin
	// 时，余量退到预算的 1/divisor。否则「调用方只给了 5 秒」会被判成「一上来就
	// 预算用尽」的假截断（3 分钟的试算预算不受影响）。
	tidyBudgetMarginDivisor = 10

	// tidyBudgetCheckEveryPairs 是两两比对循环里检查预算的间隔（按已比较的对数）。
	// 每对都问一次时钟会把 O(n²) 循环本身拖慢；按批检查足以在预算用尽后很快收手。
	tidyBudgetCheckEveryPairs = 4096
)

// tidyBudget 是本轮的时间预算（由调用方的 ctx deadline 派生）。零值 = 无预算
// （ctx 没有 deadline，如心跳轮 —— 那时它的行为与从前逐字一致）。
type tidyBudget struct {
	deadline time.Time
	margin   time.Duration
	enabled  bool
	// now 是判定预算的时钟：生产是 time.Now，测试可注入（因此「预算用尽 ⇒ 诚实
	// 截断」这条路可以确定性地验，而不必真的等 3 分钟）。零值时钟视为无预算时钟，
	// 由 newTidyBudget 填上。
	now func() time.Time
}

// newTidyBudget 从 ctx 的 deadline 派生本轮预算。ctx 为 nil 或无 deadline 时
// 预算关闭。
func (t *Tidier) newTidyBudget(ctx context.Context) tidyBudget {
	if ctx == nil {
		return tidyBudget{}
	}
	deadline, ok := ctx.Deadline()
	if !ok {
		return tidyBudget{}
	}
	return tidyBudget{
		deadline: deadline,
		margin:   tidyBudgetMarginFor(deadline.Sub(t.now())),
		enabled:  true,
		now:      t.now,
	}
}

// tidyBudgetMarginFor 按剩余预算算出收尾余量（见 tidyBudgetMarginDivisor）。
func tidyBudgetMarginFor(total time.Duration) time.Duration {
	switch {
	case total <= 0:
		return 0
	case total < tidyBudgetMargin:
		return total / tidyBudgetMarginDivisor
	default:
		return tidyBudgetMargin
	}
}

// exhausted 报告预算是否已用尽。它按**本轮整理器的时钟**判断（生产 = time.Now；
// 测试可注入），而不是另开一个时间源 —— 预算讲的是「还能烧多久」，与判定用的
// 时钟同源才不会互相打架。
func (b tidyBudget) exhausted() bool {
	if !b.enabled || b.now == nil {
		return false
	}
	return b.now().After(b.deadline.Add(-b.margin))
}

// budgetTruncationLabel 渲染时间预算截断的分区标签。它与 (scan)/(compare)/
// (candidates) 同一形状（进响应的 truncated 数组），只是原因不同：那些是
// **配置的配额**被打满，这个是**墙钟预算**用尽。
func budgetTruncationLabel(p tidyPartition) string {
	return p.label() + "(budget)"
}