package loop

import (
	"context"
	"errors"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/persist"
)

// 本文件是 design/15 §2.5.9「checkpoint 归属键由会话升为**会话 + 路**」的回归：
// 真并行时两路各写各的归属键，LoadLatest 因此天然按路隔离 —— 否则两路合写同一个
// key，后写的那份会把前一路的窗口整段顶掉（LangGraph `Send` 扇出的同族教训）。
//
// 单路口径是**逐字不变**的：ctx 上不挂路标识时归属键仍等于 sessionID —— 包括那种
// 「挂了工具互斥身份、但没挂并发路标识」的 ctx（普通回合的实际形状，见
// TestCheckpointOwnerIgnoresToolLegID：两条通道必须保持解耦）。

// keyedCheckpointStore 是按 AgentID 分桶的内存 checkpoint 存储，用来证明归属键
// 真的按路隔离（判据是**键**，而不是「循环怎么调用 store」）。
type keyedCheckpointStore struct {
	saved  []*persist.Checkpoint
	latest map[string]*persist.Checkpoint
}

func newKeyedCheckpointStore() *keyedCheckpointStore {
	return &keyedCheckpointStore{latest: make(map[string]*persist.Checkpoint)}
}

func (s *keyedCheckpointStore) Save(_ context.Context, cp *persist.Checkpoint) error {
	s.saved = append(s.saved, cp)
	s.latest[cp.AgentID] = cp
	return nil
}

func (s *keyedCheckpointStore) LoadLatest(_ context.Context, agentID string) (*persist.Checkpoint, error) {
	cp, ok := s.latest[agentID]
	if !ok {
		return nil, persist.ErrNotFound
	}
	return cp, nil
}

// TestCheckpointOwnerSingleLegUnchanged 钉住单路（未挂路标识）的归属键逐字等于
// sessionID：行 id 与 AgentID 都与引入并发路之前完全一致。
func TestCheckpointOwnerSingleLegUnchanged(t *testing.T) {
	t.Parallel()
	store := newKeyedCheckpointStore()
	l := New(Options{Checkpoint: store})
	l.saveCheckpoint(context.Background(), "sess", 2, []model.Message{model.TextMessage(model.RoleUser, "hi")}, agent.StopReasonToolCalls)

	if len(store.saved) != 1 {
		t.Fatalf("saved = %d, want 1", len(store.saved))
	}
	if got := store.saved[0].AgentID; got != "sess" {
		t.Fatalf("单路 AgentID = %q, want %q（逐字不变）", got, "sess")
	}
	if got := store.saved[0].ID; got != "sess:2" {
		t.Fatalf("单路 checkpoint id = %q, want %q", got, "sess:2")
	}
}

// TestCheckpointOwnerIgnoresToolLegID 是**守门断言**（design/15 §2.5.7 与 §2.5.9
// 解耦的守卫）：checkpoint 归属键只认 WithParallelLeg，**绝不**认工具互斥身份
// （WithToolLegID）。
//
// 它的存在是为了让「两个用途又被焊回去」变成一条会红的用例：网关给**每一条**在途路
// 都盖互斥身份（普通回合也有），一旦 checkpointAgentID 改读那条通道，单条普通回合的
// 归属键就会从 `sess` 升成 `sess#turn-1` —— 崩溃恢复语义随之改变（旧快照读不到、
// 凭空多出按路隔离）。
func TestCheckpointOwnerIgnoresToolLegID(t *testing.T) {
	t.Parallel()
	store := newKeyedCheckpointStore()
	l := New(Options{Checkpoint: store})
	// 普通回合的实际形状：有互斥身份（本路 turnID），**没有**并发路标识。
	ctx := WithToolLegID(context.Background(), "turn-1")
	window := []model.Message{model.TextMessage(model.RoleUser, "hi")}
	l.saveCheckpoint(ctx, "sess", 3, window, agent.StopReasonToolCalls)

	if len(store.saved) != 1 {
		t.Fatalf("saved = %d, want 1", len(store.saved))
	}
	if got := store.saved[0].AgentID; got != "sess" {
		t.Fatalf("互斥身份绝不能改变归属键：AgentID = %q, want %q（逐字不变）", got, "sess")
	}
	if got := store.saved[0].ID; got != "sess:3" {
		t.Fatalf("互斥身份绝不能改变归属键：checkpoint id = %q, want %q", got, "sess:3")
	}

	// 续接侧同样只看归属键：LoadLatest 请求的键仍是 sessionID，快照照样能接上。
	reader := &recordingLoadStore{inner: store}
	l.checkpoint = reader
	if _, step, resumed := l.loadCheckpoint(ctx, "sess", window); !resumed || step != 4 {
		t.Fatalf("普通回合 resumed=%v step=%d, want true/4（有互斥身份也必须照旧续接）", resumed, step)
	}
	if len(reader.calls) != 1 || reader.calls[0] != "sess" {
		t.Fatalf("LoadLatest 键 = %v, want [sess]", reader.calls)
	}
}

// TestCheckpointOwnerIsolatedPerLeg 钉住归属键按路隔离：同一会话的两条路各写各的
// 键，且各自 LoadLatest 只命中自己的那份 —— 任何一路都绝不会把对方的崩溃快照
// 当成自己的续接。
func TestCheckpointOwnerIsolatedPerLeg(t *testing.T) {
	t.Parallel()
	store := newKeyedCheckpointStore()
	l := New(Options{Checkpoint: store})
	ctxA := WithParallelLeg(context.Background(), "turn-a")
	ctxB := WithParallelLeg(context.Background(), "turn-b")

	windowA := []model.Message{
		model.TextMessage(model.RoleUser, "leg-a 的问题"),
		model.TextMessage(model.RoleTool, "tool x result: a"),
	}
	l.saveCheckpoint(ctxA, "sess", 2, windowA, agent.StopReasonToolCalls)

	if got := store.saved[0].AgentID; got != "sess#turn-a" {
		t.Fatalf("并行路 AgentID = %q, want %q", got, "sess#turn-a")
	}
	if got := store.saved[0].ID; got != "sess#turn-a:2" {
		t.Fatalf("并行路 checkpoint id = %q, want %q", got, "sess#turn-a:2")
	}

	// 同会话的另一条路：它自己的键上没有快照 ⇒ 全新开始（绝不续接 A 的窗口）。
	msgs := []model.Message{model.TextMessage(model.RoleUser, "leg-b 的问题")}
	gotB, stepB, resumedB := l.loadCheckpoint(ctxB, "sess", msgs)
	if resumedB || stepB != 1 {
		t.Fatalf("另一条路 resumed=%v step=%d, want false/1（绝不读对方的快照）", resumedB, stepB)
	}
	if len(gotB) != len(msgs) || messageText(gotB[0]) != "leg-b 的问题" {
		t.Fatalf("另一条路的窗口被对方的快照顶掉了: %+v", gotB)
	}

	// 而 A 自己仍然续接得到（归属键保持有效）。
	gotA, stepA, resumedA := l.loadCheckpoint(ctxA, "sess", windowA)
	if !resumedA || stepA != 3 {
		t.Fatalf("A resumed=%v step=%d, want true/3", resumedA, stepA)
	}
	if len(gotA) != len(windowA) {
		t.Fatalf("A 的窗口 = %d 条, want %d", len(gotA), len(windowA))
	}
}

// TestCheckpointOwnerSingleLegDoesNotSeeParallelSnapshot 钉住另一个方向：单路
// 回合（归属键 = sessionID）读不到并行路写的快照 —— 两边的键是两个空间。
func TestCheckpointOwnerSingleLegDoesNotSeeParallelSnapshot(t *testing.T) {
	t.Parallel()
	store := newKeyedCheckpointStore()
	l := New(Options{Checkpoint: store})
	l.saveCheckpoint(WithParallelLeg(context.Background(), "turn-a"), "sess", 2, []model.Message{
		model.TextMessage(model.RoleUser, "leg-a 的问题"),
	}, agent.StopReasonToolCalls)

	// 单路那次 LoadLatest 请求的键是 "sess"；存储里只有 "sess#turn-a"。
	reader := &recordingLoadStore{inner: store}
	l.checkpoint = reader
	msgs := []model.Message{model.TextMessage(model.RoleUser, "单路的问题")}
	if _, step, resumed := l.loadCheckpoint(context.Background(), "sess", msgs); resumed || step != 1 {
		t.Fatalf("单路 resumed=%v step=%d, want false/1", resumed, step)
	}
	if len(reader.calls) != 1 || reader.calls[0] != "sess" {
		t.Fatalf("LoadLatest 键 = %v, want [sess]", reader.calls)
	}
}

// recordingLoadStore 记录 LoadLatest 的入参键，同时把读请求转发给内存存储。
type recordingLoadStore struct {
	inner *keyedCheckpointStore
	calls []string
}

func (r *recordingLoadStore) Save(ctx context.Context, cp *persist.Checkpoint) error {
	return r.inner.Save(ctx, cp)
}

func (r *recordingLoadStore) LoadLatest(ctx context.Context, agentID string) (*persist.Checkpoint, error) {
	r.calls = append(r.calls, agentID)
	cp, err := r.inner.LoadLatest(ctx, agentID)
	if errors.Is(err, persist.ErrNotFound) {
		return nil, err
	}
	return cp, err
}

// TestParallelLegContextRoundTrip 钉住 ctx 通道本身：空 legID 是空操作
// （单路 / 子代理 / 后台思考不挂载 ⇒ 归属键不变），nil ctx 不 panic。
func TestParallelLegContextRoundTrip(t *testing.T) {
	t.Parallel()
	ctx := WithParallelLeg(context.Background(), "turn-a")
	if got := ParallelLegFrom(ctx); got != "turn-a" {
		t.Fatalf("ParallelLegFrom = %q, want turn-a", got)
	}
	if got := ParallelLegFrom(WithParallelLeg(context.Background(), "")); got != "" {
		t.Fatalf("空 legID 应是不挂载的 ctx 副本，得到 %q", got)
	}
	if got := ParallelLegFrom(nil); got != "" {
		t.Fatalf("nil ctx 应返回空串，得到 %q", got)
	}
	// TurnIDFrom 是 WithTurnID 的读侧：编排器先生成、loop 复用，两者必须同值。
	turnCtx := WithTurnID(context.Background(), "turn-xyz")
	if got := TurnIDFrom(turnCtx); got != "turn-xyz" {
		t.Fatalf("TurnIDFrom = %q, want turn-xyz", got)
	}
	if got := TurnIDFrom(context.Background()); got != "" {
		t.Fatalf("未挂载 TurnIDFrom = %q, want 空串", got)
	}
}