package orchestration

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/bus"
	"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"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/orchestration/audit"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/orchestration/comm"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/trace"
)

// runner 执行一个子代理的循环。*loop.Loop 满足它；测试注入 mock，使调度器无需模型
// 后端即可被演练。
type runner interface {
	Run(ctx context.Context, sessionID string, messages []model.Message) (*loop.Result, error)
}

// Scheduler 在后台创建并运行子代理：限制并发、施加按代理时限、经 agent 间总线发布
// 进度。它可安全并发使用。
//
// 生命周期为：Spawn 登记一个子代理并异步启动它；子代理的 goroutine 取得一个并行槽位
// （受 MaxParallel 把守）、转入 running、在按代理时限下于隔离上下文中运行注入的循环
// 模板、写入摘要并发布终结态进度事件。Abort/CancelAll 中断运行中或排队中的子代理。
// Join 阻塞到每个子代理终结，返回合并结果。
type Scheduler struct {
	maxParallel     int
	perAgentTimeout time.Duration
	runner          runner
	bus             bus.Bus
	comm            *comm.Probe
	memory          *memory.Store
	tracer          *trace.Tracer
	flow            FlowSpanSink
	approver        *audit.Approver
	log             *slog.Logger
	agentID         string
	// personaExists 回答「本进程能否以自身指令运行那个角色」
	// （Options.PersonaExists）；nil 表示不能指名任何角色。
	personaExists func(personaID string) bool

	mu      sync.Mutex
	agents  map[string]*SubAgent
	active  int
	allDone chan struct{}
	allOne  sync.Once
	sem     chan struct{}
	seq     atomic.Uint64
	prefix  string
}

// noopRunner 是未配置 *loop.Loop 模板时所用的兜底 runner。它使每个子代理以
// ErrNoRunner 快速失败，而不是悄悄什么都不做。
type noopRunner struct{}

// Run 返回包装后的 ErrNoRunner。
func (noopRunner) Run(context.Context, string, []model.Message) (*loop.Result, error) {
	return nil, fmt.Errorf("orchestration: run sub-agent: %w", ErrNoRunner)
}

// New 依据 opts 构建 Scheduler 并套用默认值。MaxParallel 非正时选用
// DefaultMaxParallel；Loop 为 nil 时替换为空操作 runner；Comm 为 nil 时用 Bus（若存在）
// 配合 AgentID（默认 "orchestrator"）构建；Logger 为 nil 时替换为标注了来源
// "orchestration" 的丢弃式 logger。
func New(opts Options) *Scheduler {
	maxParallel := opts.MaxParallel
	if maxParallel <= 0 {
		maxParallel = DefaultMaxParallel
	}
	agentID := opts.AgentID
	if agentID == "" {
		agentID = defaultAgentID
	}

	s := &Scheduler{
		maxParallel:     maxParallel,
		perAgentTimeout: opts.PerAgentTimeout,
		runner:          wrapLoop(opts.Loop),
		bus:             opts.Bus,
		memory:          opts.Memory,
		tracer:          opts.Tracer,
		flow:            opts.Flow,
		approver:        opts.Approver,
		log:             sourcedLog(opts.Logger),
		agentID:         agentID,
		personaExists:   opts.PersonaExists,
		agents:          make(map[string]*SubAgent),
		allDone:         make(chan struct{}),
		sem:             make(chan struct{}, maxParallel),
		prefix:          randPrefix(),
	}
	if opts.Comm != nil {
		s.comm = opts.Comm
	} else if opts.Bus != nil {
		s.comm = comm.New(opts.Bus, agentID)
	}
	return s
}

// Bus 返回注入的事件总线（可能为 nil）。
func (s *Scheduler) Bus() bus.Bus { return s.bus }

// Comm 返回注入的（或自动构建的）comm 探针（可能为 nil）。
func (s *Scheduler) Comm() *comm.Probe { return s.comm }

// Spawn 依据 spec 创建子代理、登记它、发布一条「已派发」进度事件，并在后台启动其
// 执行。它立即返回，因此前台会话不被阻塞；子代理异步运行并经总线报告进度。spec.Goal
// 为空时返回 ErrMissingGoal；spec.PersonaID 指名了本进程没有对应 agent 的角色时返回
// 包装后的 ErrUnknownPersona（design/22 §8-2）。
//
// 子代理的 Context 是 spec.Context 的深拷贝加上作为末尾一条 user 消息的任务本身
// （裁决 2026-10-04），其 WorkingMemory 是全新隔离的窗口，因此它绝不与前台会话共享
// 状态。
func (s *Scheduler) Spawn(ctx context.Context, spec AgentSpec) (*SubAgent, error) {
	if strings.TrimSpace(spec.Goal) == "" {
		return nil, fmt.Errorf("orchestration: spawn: %w", ErrMissingGoal)
	}
	if err := s.checkPersona(spec.PersonaID); err != nil {
		return nil, err
	}
	if ctx == nil {
		ctx = context.Background()
	}

	sub := &SubAgent{
		ID:            s.nextID(),
		Goal:          spec.Goal,
		RoleID:        effectiveRole(spec),
		PersonaID:     spec.PersonaID,
		Mode:          spec.Mode,
		AttachInfo:    spec.AttachInfo,
		ParentSession: spec.ParentSession,
		FlowOwner:     spec.FlowOwner,
		Data:          spec.Data,
		ToolWhitelist: cloneStrings(spec.ToolWhitelist),
		Context:       taskContext(spec),
		WorkingMemory: memory.NewWorkingMemory(0),
		timeout:       effectiveTimeout(spec.Timeout, s.perAgentTimeout),
		status:        StatusPending,
		done:          make(chan struct{}),
		finishedCh:    make(chan struct{}),
	}

	s.mu.Lock()
	if s.active == 0 {
		s.allDone = make(chan struct{})
		s.allOne = sync.Once{}
	}
	s.active++
	s.agents[sub.ID] = sub
	s.mu.Unlock()

	s.publish(EventDispatched, sub.progress())
	s.launch(ctx, sub)
	return sub, nil
}

// Run 显式启动每个尚未启动的已登记子代理。它幂等，供偏好延迟执行的调用方使用；Spawn
// 已自动启动，因此对刚派发的批次它通常是空操作。
func (s *Scheduler) Run(ctx context.Context) {
	if ctx == nil {
		ctx = context.Background()
	}
	s.mu.Lock()
	subs := make([]*SubAgent, 0, len(s.agents))
	for _, sub := range s.agents {
		subs = append(subs, sub)
	}
	s.mu.Unlock()
	for _, sub := range subs {
		s.launch(ctx, sub)
	}
}

// Join 阻塞到每个已派发的子代理终结（或 ctx 被取消），返回合并结果。ctx 取消时它
// 返回部分合并结果以及包装后的 ctx 错误。
func (s *Scheduler) Join(ctx context.Context) (*MergedResult, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	s.mu.Lock()
	ch := s.allDone
	s.mu.Unlock()

	select {
	case <-ch:
		return s.merge(), nil
	case <-ctx.Done():
		return s.merge(), ctx.Err()
	}
}

// Abort 取消给定 id 的子代理。它正在运行时中断其循环；排队中时不经运行即释放；
// 已经结束时该调用是空操作。未知 id 返回 ErrAgentNotFound。
func (s *Scheduler) Abort(ctx context.Context, id string) error {
	if ctx == nil {
		ctx = context.Background()
	}
	s.mu.Lock()
	sub, ok := s.agents[id]
	s.mu.Unlock()
	if !ok {
		return fmt.Errorf("orchestration: abort: %w", ErrAgentNotFound)
	}
	sub.cancelRun()
	return nil
}

// CancelAll 取消每个已派发的子代理，无论排队中还是运行中。它是防失控的逃生口：返回后
// 每个子代理都收敛到中止态。
func (s *Scheduler) CancelAll(ctx context.Context) error {
	if ctx == nil {
		ctx = context.Background()
	}
	s.mu.Lock()
	subs := make([]*SubAgent, 0, len(s.agents))
	for _, sub := range s.agents {
		subs = append(subs, sub)
	}
	s.mu.Unlock()
	for _, sub := range subs {
		sub.cancelRun()
	}
	return nil
}

// Agents 按派发顺序返回每个已登记子代理的快照。随时读取返回的快照都是安全的。
func (s *Scheduler) Agents() []SubAgentResult {
	s.mu.Lock()
	subs := make([]*SubAgent, 0, len(s.agents))
	for _, sub := range s.agents {
		subs = append(subs, sub)
	}
	s.mu.Unlock()

	out := make([]SubAgentResult, 0, len(subs))
	for _, sub := range subs {
		out = append(out, sub.Snapshot())
	}
	return out
}

// Approve 经注入的 *audit.Approver 触发对某个子代理敏感操作的「主代理代审」。某项
// 操作是否敏感、是否获准由上层策略决定；本方法只记录审计条目。未注入 approver 时
// 返回 ErrNoApprover。
func (s *Scheduler) Approve(ctx context.Context, actor, action, describe string) (*audit.Approval, error) {
	if s.approver == nil {
		return nil, fmt.Errorf("orchestration: approve: %w", ErrNoApprover)
	}
	return s.approver.ReviewAndApprove(ctx, actor, action, describe)
}

// launch 恰好一次地启动子代理的后台执行。它由 Spawn（自动启动）与 Run（显式启动）
// 调用；launchOnce 保证即使两者都发生也只运行一次。
func (s *Scheduler) launch(ctx context.Context, sub *SubAgent) {
	sub.launchOnce.Do(func() {
		go s.run(ctx, sub)
	})
}

// run 驱动一个子代理：取得并行槽位、转入 running、在按代理时限下于子代理的隔离上下文
// 中运行循环，然后以摘要与终结态进度事件收尾。它总是恰好到达一次 finish，因此 active
// 只递减一次，Join 恰在每个子代理终结时解除阻塞。
func (s *Scheduler) run(ctx context.Context, sub *SubAgent) {
	if ctx == nil {
		ctx = context.Background()
	}

	select {
	case s.sem <- struct{}{}:
		defer func() { <-s.sem }()
	case <-ctx.Done():
		s.finish(sub, StatusAborted, canceledSummary(), ctx.Err())
		return
	case <-sub.done:
		s.finish(sub, StatusAborted, canceledSummary(), context.Canceled)
		return
	}

	if sub.aborted() || !sub.transitionRunning() {
		s.finish(sub, StatusAborted, canceledSummary(), context.Canceled)
		return
	}
	s.publish(EventStarted, sub.progress())

	runCtx, cancel := s.subCtx(ctx, sub)
	// 子代理不得把内心独白泄漏进调用方的聊天流
	// （design/05 §8: 少说话、多做事; design/04 §2: 不污染前台）：静音抑制
	// UI 流/容量广播，而进度仍经总线与泳道到达前端。flow 分支把本次 Run 发射的
	// 每个 span 都归属到派发**那一路**的图（`{会话}[#{路}]`，见 SubAgent.flowOwner），
	// 而不是另开第二张图、也不是落进会话主干那张。
	runCtx = loop.WithMutedBroadcast(runCtx)
	runCtx = loop.WithFlowBranch(runCtx, sub.flowOwner())
	// 被派发的目标角色决定本次运行每一步的提示词与记忆作用域（design/22 §8-2）。
	// 留空则保持继承来的 env：既有行为，即子代理按派发回合的角色运行。
	runCtx = withPersonaEnv(runCtx, sub.PersonaID)
	// 工具白名单（design/05 §7 裁决 02.4.1）随运行上下文传递：名单非空时
	// 只放行所列工具；空名单＝全开，与挂载前的行为一致。
	runCtx = loop.WithToolWhitelist(runCtx, sub.ToolWhitelist)
	// 派给该子代理的任务目标随运行上下文传递（design/11 §4 的口径 A）：敏感操作代审
	// 据此**结构化**识别「这次调用由子代理发起」，并拿到判「有没有越界」所需的任务，
	// 而不是从提示词、模型文本或会话 id 形态去猜（项目红线）。
	runCtx = loop.WithSubagentTask(runCtx, sub.Goal)
	// 派发标记挂在原始 ctx 上（不是 runCtx）：它是这一批派发/收束的
	// orchestration 节点（裁决 11.4.2），与子代理泳道根并列成链。
	orchSpan := s.orchSpan(ctx, sub)
	runCtx, span := s.subSpan(runCtx, sub)
	sub.setCancel(cancel)
	defer cancel()

	res, err := s.scopedRunner(sub).Run(runCtx, sub.ID, sub.Context)
	if err != nil {
		s.spanErr(sub, span, err)
		s.finishOrchSpan(sub, orchSpan, err)
		s.abortOrFail(sub, err)
		return
	}
	s.spanOK(sub, span, res)
	s.finishOrchSpan(sub, orchSpan, nil)
	sub.setResult(res)
	s.finish(sub, StatusDone, doneSummary(res), nil)
}

// abortOrFail 对子代理运行后的错误归类：期限超出或取消产出中止态（超时/停止），其余
// 都是真正的失败。超时提示点名的是实际施加于该子代理的期限（按调用覆盖或调度器默认），
// 而不是全局那个。
func (s *Scheduler) abortOrFail(sub *SubAgent, err error) {
	switch {
	case errors.Is(err, context.DeadlineExceeded):
		s.finish(sub, StatusAborted, timeoutSummaryFor(sub.timeout), err)
	case errors.Is(err, context.Canceled):
		s.finish(sub, StatusAborted, canceledSummary(), err)
	default:
		s.finish(sub, StatusFailed, failedSummary(err), err)
	}
}

// finish 恰好一次地给子代理收尾：记录终结状态与摘要、发布终结态进度事件、递减 active
// 计数，并在最后一个子代理收尾时关闭 allDone channel。
func (s *Scheduler) finish(sub *SubAgent, status Status, summary string, err error) {
	sub.finishOnce.Do(func() {
		sub.setOutcome(status, summary, err)
		s.publish(eventFor(status), sub.progress())
		sub.markFinished()

		s.mu.Lock()
		s.active--
		if s.active == 0 {
			s.allOne.Do(func() { close(s.allDone) })
		}
		s.mu.Unlock()
	})
}

// subCtx 派生运行上下文并施加该子代理的生效时限：按调用的 AgentSpec.Timeout
// （design/05 §7 裁决 02.4.4 按调用覆盖）优先于调度器的 PerAgentTimeout；两者都未配置
// 即除取消外无期限。
func (s *Scheduler) subCtx(parent context.Context, sub *SubAgent) (context.Context, context.CancelFunc) {
	if sub.timeout > 0 {
		return context.WithTimeout(parent, sub.timeout)
	}
	return context.WithCancel(parent)
}

// subSpan 在配置了 tracer 时开启一个子代理 span，把它串成 ctx 中已有 span 的子节点。
// 子代理被归属到某张 flow 图（`SubAgent.flowOwner`，即 `{会话}[#{路}]`）时，该 span
// 还会作为子代理泳道根折进那张图，使分支渲染在派发它的**那一路**图内，而不是自成
// 一张图、也不会错落进另一路的图（design/19 §2）。它返回携带该 span 的 ctx 与 span
// 自身（tracer 为 nil 时 span 为 nil）。
func (s *Scheduler) subSpan(ctx context.Context, sub *SubAgent) (context.Context, *trace.Span) {
	if s.tracer == nil {
		return ctx, nil
	}
	sctx, span := s.tracer.StartSpan(ctx, "subagent."+sub.ID, trace.SpanKindSubagent)
	if owner := sub.flowOwner(); owner != "" && s.flow != nil {
		_ = span.SetProp("session", owner)
		s.flow.SpanStarted(owner, span)
	}
	// 目标角色随分支节点传播，使流程图面板展示子代理以哪个角色作答
	// （design/22 §8-2 可观测性），而不只是子代理 id。
	if sub.PersonaID != "" {
		_ = span.SetProp("persona", sub.PersonaID)
	}
	return sctx, span
}

// subagentEventDone 是子代理运行落定时记在其 span 上的日志事件类型，使面板节点带上
// 结果的元数据（状态/步数/回复体积与指纹，绝无正文；design/19 §1 节点实时日志）。
const subagentEventDone = "subagent.done"

// spanOK 把子代理 span 以 ok 收尾，记录**结果的体积与指纹**（不是正文）与一条同口径
// 的结果事件，随后把已关闭的 span 折回父 flow。
//
// 🔴 子代理的回复正文不进 span：它本来就要留在会话库（那是主人要读的对话），
// span 与 stderr JSONL 只是落盘的诊断面 —— 把整段回复再抄一遍，等于给一次
// stderr tail 就能拷走全部对话开了口子（M40 修掉的最严重一条泄漏）。这里留
// text_digest（字节/rune/sha256 前缀）与事件里的 status/steps/chars/sha256，
// 足以分辨「这次子代理答了多久、答了多少、是不是同一份」，却拿不到一个字。
func (s *Scheduler) spanOK(sub *SubAgent, span *trace.Span, res *loop.Result) {
	if span == nil {
		return
	}
	if res != nil {
		digest := trace.DigestOf(res.Text)
		_ = span.SetProp("text_digest", digest)
		_ = span.Record(trace.Event{Type: subagentEventDone, Props: map[string]any{
			"status": string(StatusDone),
			"steps":  res.Steps,
			"chars":  digest.Runes,
			"sha256": digest.SHA256,
		}})
	}
	span.Finish(trace.SpanStatusOK)
	s.foldSpanFinish(sub, span)
}

// spanErr 把子代理 span 以 error 收尾，记录原因，随后把已关闭的 span 折回父 flow。
func (s *Scheduler) spanErr(sub *SubAgent, span *trace.Span, err error) {
	if span == nil {
		return
	}
	_ = span.SetProp("error", err.Error())
	_ = span.Record(trace.Event{Type: subagentEventDone, Props: map[string]any{
		"status": string(StatusFailed),
		"error":  compactText(err.Error(), summaryTextMax),
	}})
	span.Finish(trace.SpanStatusError)
	s.foldSpanFinish(sub, span)
}

// foldSpanFinish 把一个刚关闭的子代理 span 折进其所属的流程图，把泳道根翻转为终态。
// 该 span 已把归属桶带在 "session" 属性上，因此发布器知道哪张图拥有它
// （真并行时就是派发它的那一路的图，见 SubAgent.flowOwner）。
func (s *Scheduler) foldSpanFinish(sub *SubAgent, span *trace.Span) {
	if span == nil || s.flow == nil || sub == nil {
		return
	}
	if owner := sub.flowOwner(); owner != "" {
		s.flow.SpanFinished(owner, span)
	}
}

// publish 经总线广播一条进度事件。未配置 comm 探针时是空操作。发布使用后台上下文，
// 因此被取消的运行仍能在子代理结束前送出它的终结态事件。
func (s *Scheduler) publish(name string, data Progress) {
	if s.comm == nil {
		return
	}
	if err := s.comm.PublishEvent(context.Background(), comm.Event{Name: name, Data: data}); err != nil {
		s.log.Warn("publish progress", "event", name, "agent_id", data.AgentID, "error", err)
	}
}

// nextID 用调度器的随机前缀与一个原子序号生成唯一的子代理 id。可安全并发使用。
func (s *Scheduler) nextID() string {
	return fmt.Sprintf("%s-%d", s.prefix, s.seq.Add(1))
}

// eventFor 把终结状态映射为对应的进度事件名。
func eventFor(status Status) string {
	switch status {
	case StatusDone:
		return EventDone
	case StatusFailed:
		return EventFailed
	case StatusAborted:
		return EventAborted
	default:
		return EventStarted
	}
}