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），不参与身份判定，也不改变
// 循环的任何行为。
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 归属键与网关侧断言使用。
func ParallelLegFrom(ctx context.Context) string {
	if ctx == nil {
		return ""
	}
	s, _ := ctx.Value(parallelLegKey{}).(string)
	return s
}