package orchestration

import (
	"context"
	"io"
	"log/slog"
	"strings"
	"sync"
	"testing"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/restrict"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/db"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)

// 本文件在派发层锁定 design/22 §8-2：`subagent_spawn` 可以指名一个目标 persona，随后
// 子代理以那个角色身份运行 —— 该角色的指令成为前置 system 消息，该角色的记忆作用域
// 就是它读写的那一份（「她以该角色身份临时替班」）。一切都以真实 *loop.Loop 搭配按角色
// agent 池与真实记忆存储、在进程内走通，因此断言的是子代理运行实际看到了什么，而不是
// 接线。
//
// 本文件的另一半钉住向后兼容契约：省略 persona 必须逐字复现改动前的行为（子代理继承
// 派发回合的角色）；未知 persona 必须被拒绝，而不是静默降级；指名 persona 绝不得放宽
// 身份/权限车道。

const (
	// parentRole 是派发会话所绑定到的角色 —— 故意不是 persona 派发所要求的角色，
	// 因此一旦泄漏就会以错误的提示词暴露出来。
	parentRole = "xilian"
	// targetRole 是 persona 派发要求子代理替班的那个角色。
	targetRole = "yuki"
)

// parentPrompt 与 targetPrompt 是两个角色的指令；测试仅凭前置 system 消息就能判断循环
// 解析到了 agent 池里的哪一项。
const (
	parentPrompt = "你是昔涟。"
	targetPrompt = "你是由纪。"
)

// roleAdapter 记录循环构建的每次请求，并以固定最终文本作答，使测试能读到子代理运行实际
// 拿到的提示词与记忆块。
type roleAdapter struct {
	mu   sync.Mutex
	reqs []model.Request
}

func (a *roleAdapter) Name() string { return "mock" }

func (a *roleAdapter) Chat(_ context.Context, req model.Request, _ model.Options) (model.Response, error) {
	a.mu.Lock()
	a.reqs = append(a.reqs, req)
	a.mu.Unlock()
	return model.Response{Model: "mock-1", Content: "好的", StopReason: "end_turn"}, nil
}

func (a *roleAdapter) ChatStream(ctx context.Context, req model.Request, _ func(model.StreamEvent)) (model.Response, error) {
	return a.Chat(ctx, req, model.Options{})
}

// request 返回第 n 次记录的请求；运行产出的请求更少时判测试失败。请求按调用顺序记录，
// 即每个子代理运行一条。
func (a *roleAdapter) request(t *testing.T, n int) model.Request {
	t.Helper()
	a.mu.Lock()
	defer a.mu.Unlock()
	if len(a.reqs) <= n {
		t.Fatalf("adapter recorded %d requests, want at least %d", len(a.reqs), n+1)
	}
	return a.reqs[n]
}

// leadingSystem 返回请求的第一条消息，agent 的指令正落在这里
// （agent/request.go 首条 system 消息）。
func (a *roleAdapter) leadingSystem(t *testing.T, n int) string {
	t.Helper()
	msgs := a.request(t, n).Messages
	if len(msgs) == 0 {
		t.Fatalf("request %d carried no messages", n)
	}
	return messageText(msgs[0])
}

// carries 报告第 n 次请求里是否有任何消息包含 want。
func (a *roleAdapter) carries(t *testing.T, n int, want string) bool {
	t.Helper()
	for _, m := range a.request(t, n).Messages {
		if strings.Contains(messageText(m), want) {
			return true
		}
	}
	return false
}

// personaRig 是进程内端到端装置：一个在捕获式适配器之上、带双角色 agent 池的真实循环，
// 其上叠真实调度器（即生产形状），再往上是派发工具，另配真实记忆存储，使角色作用域能
// 针对已持久化的条目断言。turn 钩子记录循环对每次已结算运行解析出的角色。
type personaRig struct {
	tool     *SpawnTool
	sched    *Scheduler
	ad       *roleAdapter
	memStore *memory.Store

	mu    sync.Mutex
	turns []loop.TurnEnd
}

// newPersonaRig 组装装置。personaExists 是被测的角色解析器：nil 对应一个没有接 agent 池
// 的组合根。
func newPersonaRig(t *testing.T, personaExists func(string) bool) *personaRig {
	t.Helper()
	rig := &personaRig{ad: &roleAdapter{}, memStore: newPersonaMemoryStore(t)}
	pool := make(map[string]*agent.Agent, 2)
	for id, instructions := range map[string]string{parentRole: parentPrompt, targetRole: targetPrompt} {
		ag, err := agent.New(instructions, nil, nil, agent.WithModel(rig.ad, nil))
		if err != nil {
			t.Fatalf("agent.New(%s): %v", id, err)
		}
		pool[id] = ag
	}

	lp := loop.New(loop.Options{
		Agent:     pool[parentRole],
		Agents:    pool,
		PersonaID: parentRole,
		Memory:    rig.memStore,
		Retriever: memory.NewRetriever(rig.memStore),
		// 构建时不带 persona，与 cmd/cyrene 的构建方式完全一致：把回合的角色盖到
		// 每个条目上是循环的职责。
		Extractor: memory.NewManualExtractor("", "owner"),
		OnTurnEnd: func(end loop.TurnEnd) string {
			rig.mu.Lock()
			rig.turns = append(rig.turns, end)
			rig.mu.Unlock()
			return ""
		},
	})
	rig.sched = newScheduler(t, Options{MaxParallel: 2, Loop: lp, PersonaExists: personaExists})
	rig.tool = NewSpawnTool(rig.sched, WithSpawnWait(5*time.Second))
	return rig
}

// turn 返回循环报告的第 n 个已结算回合；装置观察到的回合更少时判测试失败。
func (r *personaRig) turn(t *testing.T, n int) loop.TurnEnd {
	t.Helper()
	r.mu.Lock()
	defer r.mu.Unlock()
	if len(r.turns) <= n {
		t.Fatalf("loop reported %d settled turns, want at least %d", len(r.turns), n+1)
	}
	return r.turns[n]
}

// parentTurnCtx 是派发回合所运行的上下文：父会话自己的角色，位于 owner 全权身份车道上，
// 即网关解析出的样子。persona 派发必须能改变前者而不触碰后者。
func parentTurnCtx() context.Context {
	env := loop.OwnerEnv()
	env.PersonaID = parentRole
	return loop.WithSessionEnv(context.Background(), env)
}

// newPersonaMemoryStore 为装置打开一个内存存储。
func newPersonaMemoryStore(t *testing.T) *memory.Store {
	t.Helper()
	ctx := context.Background()
	d, err := db.Open(ctx, ":memory:", slog.New(slog.NewTextHandler(io.Discard, nil)))
	if err != nil {
		t.Fatalf("db.Open: %v", err)
	}
	t.Cleanup(func() { _ = d.Close() })
	store, err := memory.NewStore(ctx, d, slog.New(slog.NewTextHandler(io.Discard, nil)))
	if err != nil {
		t.Fatalf("memory.NewStore: %v", err)
	}
	return store
}

// seedRoleMemory 在 roleID 下归档一条 persona 作用域的记忆。
func seedRoleMemory(t *testing.T, store *memory.Store, roleID, content string) {
	t.Helper()
	_, err := store.Add(context.Background(), memory.MemoryItem{
		PersonaID: roleID,
		RoleID:    roleID,
		Scope:     memory.ScopePersona,
		Kind:      memory.KindRelationship,
		Content:   content,
	})
	if err != nil {
		t.Fatalf("seed %s memory: %v", roleID, err)
	}
}

// storedPersonaOf 返回内容包含给定字符串的那条记忆所归档到的 persona；不存在时返回 ""。
func storedPersonaOf(t *testing.T, store *memory.Store, content string) string {
	t.Helper()
	items, err := store.List(context.Background())
	if err != nil {
		t.Fatalf("memory list: %v", err)
	}
	for _, it := range items {
		if strings.Contains(it.Content, content) {
			return it.PersonaID
		}
	}
	return ""
}

// TestSpawnToolRunsAsTargetPersona 是 design/22 §8-2 的验收，进程内端到端：指名 persona
// 的派发以那个角色的指令与那个角色的记忆运行，而派发会话自己的角色保持不变。同一装置
// 随后证明省略 persona 的路径未变，即改动前的行为：子代理继承派发回合的角色。
func TestSpawnToolRunsAsTargetPersona(t *testing.T) {
	t.Parallel()
	rig := newPersonaRig(t, func(id string) bool { return id == parentRole || id == targetRole })
	// 每个角色一条记忆，使注入的块能点名该次运行读的是哪个作用域。
	seedRoleMemory(t, rig.memStore, targetRole, "由纪的暗号是香菜")
	seedRoleMemory(t, rig.memStore, parentRole, "昔涟的暗号是乌冬")

	raw, err := rig.tool.Execute(parentTurnCtx(), map[string]any{
		"goal": "替我确认一下口味", "persona": targetRole,
	})
	if err != nil {
		t.Fatalf("Execute: %v", err)
	}
	got := decodeSpawnResult(t, raw)
	if len(got.Agents) != 1 || got.Agents[0].Status != string(StatusDone) {
		t.Fatalf("agents = %+v, want one settled sub-agent", got.Agents)
	}
	// 报告的角色就是目标 persona，因此工具结果、进度事件与 ws subagent_event 帧都点名
	// 它真正以之作答的角色。
	if got.Agents[0].Role != targetRole {
		t.Fatalf("reported role = %q, want the target persona %q", got.Agents[0].Role, targetRole)
	}

	// 1) 模型实际收到的提示词是目标角色的指令 —— 既不是派发角色的，也不是共享默认值。
	if prompt := rig.ad.leadingSystem(t, 0); prompt != targetPrompt {
		t.Fatalf("sub-agent system prompt = %q, want the target persona's %q", prompt, targetPrompt)
	}
	// 2) 循环为本次运行解析出的角色（它决定记忆、情绪与思考旋钮的作用域）是目标 persona，
	// 尽管派发的 env 说的是 xilian。
	if turn := rig.turn(t, 0); turn.PersonaID != targetRole {
		t.Fatalf("sub-agent turn persona = %q, want %q", turn.PersonaID, targetRole)
	}
	// 3) 记忆读取跟随目标角色：只有由纪的记忆被注入。
	if !rig.ad.carries(t, 0, "由纪的暗号是香菜") {
		t.Fatalf("sub-agent did not recall the target persona's memory: %+v", rig.ad.request(t, 0).Messages)
	}
	if rig.ad.carries(t, 0, "昔涟的暗号是乌冬") {
		t.Fatalf("sub-agent read the dispatching persona's memory: %+v", rig.ad.request(t, 0).Messages)
	}
	// 4) 身份车道原样传递：指名 persona 从来不是一种权限
	// （项目红线：身份识别靠元数据 + 策略引擎）。
	env := rig.turn(t, 0).Env
	if env.ActorEnv != restrict.ActorOwner || !env.MemWritable {
		t.Fatalf("sub-agent env = %+v, want the dispatching owner lane unchanged", env)
	}

	// 向后兼容：不带 persona 的派发仍以派发角色（xilian）运行、读昔涟的记忆 —— 逐字节
	// 等同改动前的行为。
	raw, err = rig.tool.Execute(parentTurnCtx(), map[string]any{
		"goal": "老路子", "role": "legacy-label",
	})
	if err != nil {
		t.Fatalf("Execute without persona: %v", err)
	}
	got = decodeSpawnResult(t, raw)
	if prompt := rig.ad.leadingSystem(t, 1); prompt != parentPrompt {
		t.Fatalf("persona-less sub-agent prompt = %q, want the spawning role's %q", prompt, parentPrompt)
	}
	if turn := rig.turn(t, 1); turn.PersonaID != parentRole {
		t.Fatalf("persona-less sub-agent persona = %q, want the inherited %q", turn.PersonaID, parentRole)
	}
	if rig.ad.carries(t, 1, "由纪的暗号是香菜") {
		t.Fatalf("persona-less sub-agent read another role's memory: %+v", rig.ad.request(t, 1).Messages)
	}
	if !rig.ad.carries(t, 1, "昔涟的暗号是乌冬") {
		t.Fatalf("persona-less sub-agent lost the inherited role's memory: %+v", rig.ad.request(t, 1).Messages)
	}
	// 未指名 persona 时，旧的角色标签仍原样往返。
	if got.Agents[0].Role != "legacy-label" {
		t.Fatalf("legacy role = %q, want it untouched when persona is omitted", got.Agents[0].Role)
	}
}

// TestSpawnedSubAgentMemoryWriteScopeIsTargetPersona 钉住 design/22 §8-2 的写侧：子代理
// 回合提取出的记忆归档到目标 persona 下，而不是派发它的那个会话的角色下。它同时钉住
// persona 绝不放宽身份车道：在受限（只读）env 下，同一次派发即使指名了角色也依然什么
// 都不写。
func TestSpawnedSubAgentMemoryWriteScopeIsTargetPersona(t *testing.T) {
	t.Parallel()
	rig := newPersonaRig(t, func(string) bool { return true })

	sub, err := rig.sched.Spawn(parentTurnCtx(), AgentSpec{
		Goal:      "记住主人的口味",
		RoleID:    targetRole,
		PersonaID: targetRole,
		Context:   []model.Message{model.TextMessage(model.RoleUser, "记住 我讨厌香菜")},
	})
	if err != nil {
		t.Fatalf("Spawn: %v", err)
	}
	if _, err := rig.sched.Wait(context.Background(), []string{sub.ID}); err != nil {
		t.Fatalf("Wait: %v", err)
	}
	if got := storedPersonaOf(t, rig.memStore, "我讨厌香菜"); got != targetRole {
		t.Fatalf("extracted memory filed under %q, want the target persona %q", got, targetRole)
	}

	// 受限 env：派发回合来自群聊（只读）。指名 persona 绝不得在车道禁止的地方开出一条
	// 写入路径。
	restricted := loop.WithSessionEnv(context.Background(), loop.SessionEnv{
		ActorEnv:    restrict.ActorGroupMember,
		MemWritable: false,
		PersonaID:   parentRole,
	})
	sub, err = rig.sched.Spawn(restricted, AgentSpec{
		Goal:      "群聊里的她被派去记住东西",
		PersonaID: targetRole,
		Context:   []model.Message{model.TextMessage(model.RoleUser, "记住 我喜欢吃辣")},
	})
	if err != nil {
		t.Fatalf("Spawn (restricted): %v", err)
	}
	if _, err := rig.sched.Wait(context.Background(), []string{sub.ID}); err != nil {
		t.Fatalf("Wait (restricted): %v", err)
	}
	if got := storedPersonaOf(t, rig.memStore, "我喜欢吃辣"); got != "" {
		t.Fatalf("a restricted sub-agent wrote memory under %q, want no write at all", got)
	}
	env := rig.turn(t, 1).Env
	if env.MemWritable || env.ActorEnv != restrict.ActorGroupMember {
		t.Fatalf("restricted sub-agent env = %+v, want the read-only group lane preserved", env)
	}
}