package loop

import "context"

// loopSessionIDKey 是当前 Run 所驱动 session id 的 context key。
// 它在 Run 顶部播种，并在循环发出会话级 UI 事件
// （capacity_awaken）处读取，使事件携带正确的会话，
// 而不是被当作全局事件。
type loopSessionIDKey struct{}

// withLoopSession 返回携带 Run 会话 id 的 context。
func withLoopSession(ctx context.Context, sessionID string) context.Context {
	return context.WithValue(ctx, loopSessionIDKey{}, sessionID)
}

// loopSession 返回 ctx 上的 session id；未播种时返回 ""（为在
// Run 之外驱动循环的调用方提供的防御性兜底）。
func loopSession(ctx context.Context) string {
	s, _ := ctx.Value(loopSessionIDKey{}).(string)
	return s
}

// SessionFromContext 返回外层 Run 所驱动的 session id；未挂载时
// 返回 ""。它导出是为了让循环内执行的工具能把其启动的工作
// （派生的子代理、后台任务）归属到发起它的会话 —— 绝不
// 从模型输出推断。
func SessionFromContext(ctx context.Context) string { return loopSession(ctx) }

// loopTurnIDKey 是当前 Run 正在产出的 assistant turn 的 context key。
// gateway 在开始一个 turn 前挂载它（WithTurnID），使循环发出的
// 流式事件能被客户端关联到该 turn。
type loopTurnIDKey struct{}

// WithTurnID 把 turnID 挂到 ctx 上，使循环发出的流式事件携带
// 所属的 turn。它导出是因为 gateway 每个 turn 在 Run 之前调用一次；
// 未挂载的调用方会让事件的 TurnID 为空，
// 从而保持关联字段向后兼容。
func WithTurnID(ctx context.Context, turnID string) context.Context {
	return context.WithValue(ctx, loopTurnIDKey{}, turnID)
}

// TurnIDFrom 返回 ctx 上的 turn id；未挂载时返回 ""。它是 WithTurnID 的读侧，
// 供网关在 Run **之前**复用自己已经生成的那个 id（design/15 §2.5.9 的 ⚠️ 必改项）：
// 客户端按 `turnId` 两级分桶，而 `turn_begin` 由编排器在调用 Run 之前发出 ——
// 若 Run 内部再自己生成一个，两处 id 不同，分桶会立即失配。
func TurnIDFrom(ctx context.Context) string {
	if ctx == nil {
		return ""
	}
	return loopTurn(ctx)
}

// loopTurn 返回 ctx 上的 turn id；未挂载时返回 ""。
func loopTurn(ctx context.Context) string {
	s, _ := ctx.Value(loopTurnIDKey{}).(string)
	return s
}

// parallelLegKey 是「本 Run 驱动的是某条并发路（leg）」的 context key。
// 它**只**影响 checkpoint 归属（design/15 §2.5.9），不参与身份判定，也不改变
// 循环的任何行为。工具层的文件级互斥走**另一条独立通道**（toolLegIDKey，见
// 本文件下半部分）—— 两条通道必须分开：这条只有真并行路才有，而互斥身份
// 每条在途路都有。焊在一起会把普通回合的 checkpoint 键从 `sess` 改成
// `sess#turn-x`，连带改变崩溃恢复语义（checkpoint.go 的 checkpointAgentID）。
type parallelLegKey struct{}

// WithParallelLeg 把本回合所属**并发路**的标识挂到 ctx 上（design/15 §2.5.9）：
// 同一会话真并行时，两路各写各的 checkpoint 归属键 `{sessionID}#{legID}`，
// 使 LoadLatest 天然按路隔离 —— 否则两路会互相把对方的窗口当自己的续接
// （LangGraph `Send` 扇出的同族教训：并行分支必须有独立归属）。
//
// 空 legID 是空操作：单路、子代理、后台思考等未声明并行的运行不挂该键，
// 归属键逐字等于 sessionID，行为与加入本特性之前完全一致。
func WithParallelLeg(ctx context.Context, legID string) context.Context {
	if ctx == nil || legID == "" {
		return ctx
	}
	return context.WithValue(ctx, parallelLegKey{}, legID)
}

// ParallelLegFrom 返回 ctx 上挂的并发路标识；未挂载时返回 ""。它是
// WithParallelLeg 的读侧，供 checkpoint 归属键与网关侧断言使用。
//
// ⚠️ 它**只**服务 checkpoint 归属，绝不是工具互斥身份 —— 工具侧要读
// ToolLegIDFrom（那条通道每条在途路都有）。
func ParallelLegFrom(ctx context.Context) string {
	if ctx == nil {
		return ""
	}
	s, _ := ctx.Value(parallelLegKey{}).(string)
	return s
}

// toolLegIDKey 是「本 Run 属于哪一条**在途路**」的 context key，**只**服务
// 工具层的文件级互斥（design/15 §2.5.7、主人 2026-10-05 裁决 3）。
//
// 它与 parallelLegKey 刻意分成两条通道，因为二者的**受众与取舍完全不同**：
//   - 这条：每条在途路都挂（普通回合 / 并行路 / 队首续跑路），回答「此刻有没有
//     **另一条路**正在动这个文件」；它不能改变 checkpoint 归属，否则「1 普通回合
//     + 1 并行路」的会话里普通回合的快照键会被升成 `sess#turn-x`，崩溃恢复语义
//     随之改变（硬要求：恢复语义零变化）；
//   - 那条：只有真并行路才有，回答「这条路的 checkpoint 该写哪个键」。
//
// 把两者焊成一个（复用 ParallelLegFrom）是曾经的真实缺口：普通回合既不登记也不
// 被检查，它与并行路同文件写仍可能并发。
type toolLegIDKey struct{}

// WithToolLegID 把本路的互斥身份挂到 ctx 上（design/15 §2.5.7）。身份由**网关/
// 编排器**在**每一条在途路**上挂载，取值就是该路的 turnID（编排器生成、唯一，
// 普通回合也有）。它只回答「这条工具调用属于哪条在途路」，与 actor 元数据 /
// 策略引擎无关（项目红线：身份识别绝不走这条通道），也不影响 checkpoint 归属。
//
// 空 legID 是空操作：子代理、后台思考、未接编排器的调用方不挂该键 ⇒ 工具层读到
// 空串 ⇒ 完全不做互斥，行为与加入本特性之前逐字相同（兼容硬要求）。
func WithToolLegID(ctx context.Context, legID string) context.Context {
	if ctx == nil || legID == "" {
		return ctx
	}
	return context.WithValue(ctx, toolLegIDKey{}, legID)
}

// ToolLegIDFrom 返回 ctx 上挂的互斥身份；未挂载时返回 ""（⇒ 不登记也不检查）。
// 它是 WithToolLegID 的读侧，只被循环的工具执行点读一次（tools.go 的 invokeOne
// 把它桥接成 tool.WithLegID —— 工具侧凭 tool.LegIDFrom 取到）。
func ToolLegIDFrom(ctx context.Context) string {
	if ctx == nil {
		return ""
	}
	s, _ := ctx.Value(toolLegIDKey{}).(string)
	return s
}