package orchestration

import (
	"context"
	"encoding/json"
	"errors"
	"sort"
	"strings"
	"testing"
	"time"

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

// This file locks the CONTRACT half of design/22 §8-2 (the behavior half lives in
// spawnpersona_test.go): the tool schema and result shape a model sees, the rejection
// of an unknown role, and the session env a sub-agent run is stamped with. The point of
// these tests is that a dispatch which names no persona stays byte-identical to the
// pre-change behavior, while a dispatch that names an unusable role fails loudly rather
// than degrading into the default role.

// TestSpawnPersonaRejectsUnknownRole is the "未知角色 → 明确报错，不静默回落" half: an id
// the process cannot run is refused before anything starts, for a single dispatch and
// for a batch (all-or-nothing, so no half-run batch is left behind), and a blank id is
// never treated as a role.
func TestSpawnPersonaRejectsUnknownRole(t *testing.T) {
	t.Parallel()
	cases := []struct {
		name    string
		exists  func(string) bool
		args    map[string]any
		wantErr bool
	}{
		{
			name:    "unknown persona",
			exists:  func(id string) bool { return id == targetRole },
			args:    map[string]any{"goal": "x", "persona": "nobody"},
			wantErr: true,
		},
		{
			name:    "no resolver wired",
			exists:  nil,
			args:    map[string]any{"goal": "x", "persona": targetRole},
			wantErr: true,
		},
		{
			name:    "batch with one unknown persona",
			exists:  func(id string) bool { return id == targetRole },
			args:    map[string]any{"tasks": []any{map[string]any{"goal": "a", "persona": targetRole}, map[string]any{"goal": "b", "persona": "nobody"}}},
			wantErr: true,
		},
		{
			name:   "known persona",
			exists: func(id string) bool { return id == targetRole },
			args:   map[string]any{"goal": "x", "persona": targetRole},
		},
		{
			name:   "blank persona is not a role",
			exists: nil,
			args:   map[string]any{"goal": "x", "persona": "   "},
		},
	}
	for _, tc := range cases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			rig := newPersonaRig(t, tc.exists)
			_, err := rig.tool.Execute(parentTurnCtx(), tc.args)
			if tc.wantErr {
				if !errors.Is(err, ErrUnknownPersona) {
					t.Fatalf("err = %v, want ErrUnknownPersona", err)
				}
				if got := len(rig.sched.Agents()); got != 0 {
					t.Fatalf("a rejected dispatch started %d sub-agents, want 0", got)
				}
				return
			}
			if err != nil {
				t.Fatalf("Execute: %v", err)
			}
		})
	}
}

// TestSchedulerSpawnCarriesPersonaEnv pins the library half of the same contract for
// callers that bypass the tool: Spawn stamps the target role onto the run's env (so
// the loop resolves that role's prompt and memory scope), a blank persona leaves the
// inherited env untouched, and Spawn itself refuses an unknown role.
func TestSchedulerSpawnCarriesPersonaEnv(t *testing.T) {
	t.Parallel()
	s := newScheduler(t, Options{PersonaExists: func(id string) bool { return id == targetRole }})
	envs := make(chan loop.SessionEnv, 4)
	s.runner = envRunner{envs: envs}

	sub, err := s.Spawn(parentTurnCtx(), AgentSpec{Goal: "g", PersonaID: targetRole})
	if err != nil {
		t.Fatalf("Spawn: %v", err)
	}
	// The recorded binding names the target role too, so the progress events and the
	// ws subagent_event frame report the role the sub-agent answered as.
	if sub.RoleID != targetRole || sub.PersonaID != targetRole {
		t.Fatalf("sub-agent binding = %q/%q, want the target role on both", sub.RoleID, sub.PersonaID)
	}
	if _, err := s.Wait(context.Background(), []string{sub.ID}); err != nil {
		t.Fatalf("Wait: %v", err)
	}
	if env := <-envs; env.PersonaID != targetRole || env.ActorEnv != restrict.ActorOwner || !env.MemWritable {
		t.Fatalf("run env = %+v, want the target persona on the unchanged owner lane", env)
	}

	legacy, err := s.Spawn(parentTurnCtx(), AgentSpec{Goal: "g"})
	if err != nil {
		t.Fatalf("Spawn (no persona): %v", err)
	}
	if _, err := s.Wait(context.Background(), []string{legacy.ID}); err != nil {
		t.Fatalf("Wait (no persona): %v", err)
	}
	if env := <-envs; env.PersonaID != parentRole {
		t.Fatalf("persona-less run env = %+v, want the inherited role %q untouched", env, parentRole)
	}

	if _, err := s.Spawn(context.Background(), AgentSpec{Goal: "g", PersonaID: "nobody"}); !errors.Is(err, ErrUnknownPersona) {
		t.Fatalf("Spawn with an unknown persona: err = %v, want ErrUnknownPersona", err)
	}
}

// envRunner records the session env each sub-agent run starts with. The memory scope
// a run uses is decided from exactly this value, so it is the boundary under test for
// callers that go through Scheduler.Spawn directly.
type envRunner struct {
	envs chan loop.SessionEnv
}

func (r envRunner) Run(ctx context.Context, _ string, _ []model.Message) (*loop.Result, error) {
	r.envs <- loop.SessionEnvFrom(ctx)
	return &loop.Result{Text: "ok", Steps: 1}, nil
}

// TestSpawnToolResultShapeIsUnchangedWithoutPersona pins the strict backward
// compatibility of the model-visible result: a dispatch that names no persona still
// returns exactly the same JSON keys as before this feature — no "persona" field is
// added anywhere — so nothing about the omitted-persona contract moved.
func TestSpawnToolResultShapeIsUnchangedWithoutPersona(t *testing.T) {
	t.Parallel()
	s, sp := spawnToolFor(t, Options{}, time.Second)
	s.runner = &mockRunner{result: &loop.Result{Text: "done", Steps: 2, StopReason: "stop"}}

	raw, err := sp.Execute(context.Background(), map[string]any{
		"goal": "g", "role": "cyrene", "mode": "engineer",
	})
	if err != nil {
		t.Fatalf("Execute: %v", err)
	}
	b, err := json.Marshal(raw)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}
	var top map[string]any
	if err := json.Unmarshal(b, &top); err != nil {
		t.Fatalf("unmarshal: %v", err)
	}
	if got := sortedKeys(top); strings.Join(got, ",") != "agents,status,summary" {
		t.Fatalf("result keys = %v, want the unchanged agents/status/summary", got)
	}
	agents, _ := top["agents"].([]any)
	if len(agents) != 1 {
		t.Fatalf("agents = %v, want one entry", top["agents"])
	}
	agent, _ := agents[0].(map[string]any)
	want := "goal,id,mode,role,status,steps,summary,text"
	if got := sortedKeys(agent); strings.Join(got, ",") != want {
		t.Fatalf("agent keys = %v, want %s (no new client-visible field)", got, want)
	}
}

// sortedKeys returns m's keys in sorted order, for a stable comparison in tests.
func sortedKeys(m map[string]any) []string {
	out := make([]string, 0, len(m))
	for k := range m {
		out = append(out, k)
	}
	sort.Strings(out)
	return out
}

// TestSpawnSchemaKeepsPersonaOptional pins the tool contract: persona is an optional
// string at both the single-goal and the tasks level, nothing new is required, and the
// pre-existing arguments keep their meaning — so a model (or an old client) that never
// sends persona sees byte-identical validation behavior.
func TestSpawnSchemaKeepsPersonaOptional(t *testing.T) {
	t.Parallel()
	props, err := NewSpawnTool(nil).Parameters()
	if err != nil {
		t.Fatalf("Parameters: %v", err)
	}
	// The required set is unchanged: goal or tasks, nothing else.
	if got := requiredAnyOf(t, props); strings.Join(got, ",") != "goal,tasks" {
		t.Fatalf("anyOf required = %v, want goal/tasks only", got)
	}
	top := stringProp(t, props, "persona")
	if top == nil {
		t.Fatal("the top-level persona property is missing")
	}
	if top["type"] != "string" {
		t.Fatalf("persona type = %v, want string", top["type"])
	}
	items, _ := props["properties"].(map[string]any)["tasks"].(map[string]any)["items"].(map[string]any)
	taskProps, _ := items["properties"].(map[string]any)
	if _, ok := taskProps["persona"]; !ok {
		t.Fatal("a tasks entry cannot name a persona")
	}
	if req, _ := items["required"].([]any); len(req) != 1 || req[0] != "goal" {
		t.Fatalf("tasks required = %v, want goal only", req)
	}
}

// requiredAnyOf flattens the anyOf required lists of a tool schema.
func requiredAnyOf(t *testing.T, props map[string]any) []string {
	t.Helper()
	anyOf, _ := props["anyOf"].([]any)
	out := make([]string, 0, len(anyOf))
	for _, entry := range anyOf {
		m, _ := entry.(map[string]any)
		req, _ := m["required"].([]any)
		for _, r := range req {
			s, _ := r.(string)
			out = append(out, s)
		}
	}
	return out
}

// stringProp returns the named top-level property of a tool schema property block.
func stringProp(t *testing.T, props map[string]any, name string) map[string]any {
	t.Helper()
	properties, _ := props["properties"].(map[string]any)
	prop, _ := properties[name].(map[string]any)
	return prop
}