package loop

import (
	"context"
	"log/slog"
	"sort"
	"sync"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/trace"
)

// graphUpdateSubject 是 FlowPublisher 发布每轮实时
// flow 图所用的 bus subject。gateway hub 订阅 ws.> 并将其
// 作为 WS 事件类型 graph_update 转发（design/19 §1）。
const graphUpdateSubject = "ws.graph_update"

// flowLogSource 是发布器日志的 source 标签（落盘失败记 WARN；AGENTS §10 禁裸日志）。
const flowLogSource = "runtime/loop/flow"

// allowedSessions 是用于控制哪些会话可发布
// flow 图的谓词签名（主人的 web 会话），使 IM/陌生人/子代理轮次绝不
// 泄露进管理面板。谓词为 nil 时允许所有会话。
//
// 谓词收的是**会话 id**（不是归属桶）：调用方先经 splitFlowOwner 拆掉路标识，
// 因此白名单永远只按会话判 —— 路的存在绝不放宽或收紧任何一档。
type allowedSessions func(sessionID string) bool

// RoundRecorder 把一条**已收尾轮次**的图落盘（design/19 §4 的 P3「多轮历史回放」；
// 切片 F2）。生产实现是 `internal/runtime/flowgraph` 的 `*SQLStore`。
//
// 语义约定（实现方必须遵守）：**每轮一条**（只在 EndRound 里调用，绝不在每个 span
// 上调用），失败**上抛**给发布器（发布器记 WARN 后继续）—— 图是可视化资产，落盘失败
// 绝不允许影响会话收尾，也绝不改变实时帧。
//
// 「每轮一条」在真并行下的含义是**每一路一条**：两路各有自己的图 id（见
// trace.RunGraph 的 id 口径），因此 `flow_graphs` 的 `(session_id, turn_id)` 唯一键
// 不会被两路撞上（design/15 §2.5.10-1）。
//
// 白名单不在实现方：EndRound 对发布会话范围外的会话**直接返回**，因此落盘出口天然
// 只被白名单内的会话触发（复用 flowSessionAllowed，不存在第二份白名单）。
type RoundRecorder interface {
	RecordRound(g *trace.FlowGraph) error
}

// FlowPublisherOption 配置 FlowPublisher 的可选行为（落盘出口 + 日志）。
type FlowPublisherOption func(*FlowPublisher)

// WithRoundRecorder 挂一个「每轮收尾时落一条」的落盘出口。传 nil 等价于不落盘。
func WithRoundRecorder(rec RoundRecorder) FlowPublisherOption {
	return func(p *FlowPublisher) { p.record = rec }
}

// WithFlowLogger 给发布器一个 logger：落盘失败要**可见**（记 WARN），但绝不把图内容
// 写进日志（图里可能含用户输入文本，只记 session/turn 标识）。
func WithFlowLogger(log *slog.Logger) FlowPublisherOption {
	return func(p *FlowPublisher) { p.log = log }
}

// FlowPublisher 从循环的 trace span 累积**按路**的 flow 图，
// 并在每个 span 开始/结束时发布 ws.graph_update，给前端
// 一张实时的 input→think→tool→output 图（design/19 §1）。它容忍 nil
// （nil 发布者为空操作）且并发安全；同一会话的图
// 在同一把互斥锁下串行化，因为并行的工具 span 会并发完成。
//
// 归属：**一个 Run 拥有一个路 bucket**。单路时 bucket 就是会话 id（行为与引入
// 「路」之前逐字相同）；真并行时每一路各有自己的 bucket（`{会话}#{路}`，见
// FlowOwnerKey），因此两路的节点绝不互相长在对方的图上、也绝不互相重置
// （design/15 §2.5.10-1 登记的那条「graph_update 互相顶掉」）。声明自己是另一
// 会话分支的 Run（loop.WithFlowBranch）绝不调用 StartRound/EndRound，
// 而是把其 span 发布到所属路的 bucket 下，因此子代理渲染为
// 父图的一条泳道，而不是第二张图（design/19 §2）。
type FlowPublisher struct {
	mu      sync.Mutex
	publish func(subject string, payload any)
	allow   allowedSessions
	// graphs 是「归属桶 → 该桶当前这张图」。键见 FlowOwnerKey。
	graphs map[string]*trace.RunGraph
	// seq 是**每会话**的建图序号：每次 StartRound 加一，并写进图 id
	// （`{会话}-{序号}[#{路}]`）。用序号而不是轮次号做 id 的第二个分量，是为了让
	// 「同一轮里的第二条路」也拿到一个全新的 id —— 否则两路会撞同一个
	// `flow_graphs.turn_id`（唯一键），后收尾的那一路把前一路整行顶掉。
	seq map[string]int
	// rounds 是每会话当前的**轮次号**（= 该轮第一条路的建图序号）。同一轮的各条路
	// 共享它，客户端据此只保留最新一轮的泳道。
	rounds map[string]int
	// roundOf 是「归属桶 → 它属于哪一轮」。
	roundOf map[string]int
	// inflight 是每会话**在飞的桶**集合。新一轮的判据是「上一轮的路全部收尾」
	// （集合为空），因此它是轮次号推进的唯一依据，也是内存不随路数无界增长的
	// 依据（新一轮起时丢掉更早轮次的桶）。
	inflight map[string]map[string]bool
	// record 是每轮收尾的落盘出口（design/19 §4 P3；切片 F2）。nil = 不落盘
	// （保持 F1 之前的纯内存行为，且测试可只测发布）。
	record RoundRecorder
	// log 用于把落盘失败记成 WARN（带 source）。nil 时静默（测试/未接线）。
	log *slog.Logger
}

// NewFlowPublisher 构建一个 FlowPublisher。publish 非 nil 时，会在
// 每次 span 状态迁移时以 graphUpdateSubject 与 *trace.FlowGraph
// 快照调用。allow 非 nil 时，控制哪些 session id 可发布（返回
// false 则整个会话被丢弃）。两者都可以为 nil（nil publish 为空操作）。
//
// 可变项经 Functional Options 传入（不改变既有两参调用）：落盘出口（F2）与日志。
func NewFlowPublisher(publish func(subject string, payload any), allow func(sessionID string) bool, opts ...FlowPublisherOption) *FlowPublisher {
	p := &FlowPublisher{
		publish:  publish,
		allow:    allow,
		graphs:   make(map[string]*trace.RunGraph),
		seq:      make(map[string]int),
		rounds:   make(map[string]int),
		roundOf:  make(map[string]int),
		inflight: make(map[string]map[string]bool),
	}
	for _, opt := range opts {
		if opt != nil {
			opt(p)
		}
	}
	return p
}

// StartRound 为 owner 开始一张新图，用用户文本
// 播种一个 input 节点并发布初始帧。
//
// owner 是**归属桶**（`{会话}[#{路}]`，见 FlowOwnerKey）：单路调用方传的就是会话
// id（引入本切片之前的所有调用点原样可用），真并行路由经 flowOwner 传
// `{会话}#{路}`。两条路各起各的图 ⇒ 后起的那一路绝不重置前一路的图。
//
// 本入口**没有 ctx 通道**，因此这一路的身份只能取**归属桶里的路标识**：真并行路的桶
// 正是网关按该路的 turnId 拼的，两者逐字同值；单路的桶里没有路标识 ⇒ 图上该键缺席。
// 要覆盖**全部编排回合**（普通单路回合与队首续跑路也算）请用 StartRoundWithTurnID，
// 它是 `Run` 走的入口。
func (p *FlowPublisher) StartRound(owner, inputText string) {
	_, lane := splitFlowOwner(owner)
	p.startRound(owner, lane, inputText)
}

// StartRoundWithTurnID 与 StartRound 同一件事，但这一路的**权威回合标识**取自 ctx
// （`TurnIDFrom`，design/15 §2.5.9）：编排器给**每一路**都盖它 —— 普通单路回合、队首
// 续跑路、真并行路都有，因此它是「这张图属于哪一回合」的唯一正当来源，也是落盘
// `flow_graphs.lane_turn_id` 与消息面 `toolCalls[]`（按 `chat_messages.turn_id` 等值
// 回看）能对上的前提。
//
// ctx 上没有权威标识时（后台思考、未接编排器的调用方）**保持缺席**：快照上的 `turnId`
// 键与落盘行的 `lane_turn_id` 都是空 ⇒ 读侧查不到这一轮的图。那是契约允许的诚实空态，
// **绝不**拿会话 id、时间戳之类去臆造一个身份。
//
// 归属桶（owner）的语义一字未改：真并行两路仍各起各的图（桶是 `{会话}#{路}`），单路的
// 桶仍是会话 id。本方法只决定**这张图属于哪一回合**，不决定它是不是另一张图。
func (p *FlowPublisher) StartRoundWithTurnID(ctx context.Context, owner, inputText string) {
	p.startRound(owner, TurnIDFrom(ctx), inputText)
}

// startRound 是两条入口的共用实现：turnID 是这一路的权威回合标识（空 = 缺席，绝不
// 臆造），owner 是归属桶。
//
// 轮次号只在**上一轮的路全部收尾**之后推进（inflight 为空），因此「同一条消息
// 引出的两路」共享同一个轮次号；新一轮开始时丢掉更早轮次的桶，内存因此不随路数
// 无界增长（图的历史在落盘里，见 flowgraph 包）。
func (p *FlowPublisher) startRound(owner, turnID, inputText string) {
	session, _ := splitFlowOwner(owner)
	if session == "" || owner == "" || !p.allowed(session) {
		return
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	round := p.beginRoundLocked(session, owner)
	g := trace.NewRunGraph(session)
	g.SetTurnID(turnID)
	g.SetSeq(p.seq[session])
	g.Start(round, time.Now())
	g.SeedInput(inputText)
	p.graphs[owner] = g
	p.roundOf[owner] = round
	if p.inflight[session] == nil {
		p.inflight[session] = make(map[string]bool, 2)
	}
	p.inflight[session][owner] = true
	p.publishLocked(g)
}

// beginRoundLocked 分配（或沿用）本轮次号并清掉更早轮次的桶。调用方必须持有 p.mu。
func (p *FlowPublisher) beginRoundLocked(session, owner string) int {
	p.seq[session]++
	if len(p.inflight[session]) == 0 {
		p.rounds[session] = p.seq[session]
		p.dropEarlierRoundsLocked(session)
	}
	return p.rounds[session]
}

// dropEarlierRoundsLocked 丢掉该会话**更早轮次**留下的桶与图。图的历史由落盘承载
// （`flow_graphs`），实时桶只需覆盖当前这一轮 —— 否则每一条并行路（turnId 每回合
// 全新、bucket 键因此永不复用）都会在内存里永久留一张图。调用方必须持有 p.mu。
func (p *FlowPublisher) dropEarlierRoundsLocked(session string) {
	current := p.rounds[session]
	for owner, round := range p.roundOf {
		if round >= current {
			continue
		}
		if s, _ := splitFlowOwner(owner); s != session {
			continue
		}
		delete(p.roundOf, owner)
		delete(p.graphs, owner)
	}
}

// EndRound 用回复播种一个 output 节点并发布最终帧，
// 以此收尾该路的图。它保留该图，使错过最终帧的客户端能重放最后一份快照 ——
// **重放的出口是 `LastGraph` / `LatestLanes`**（网关只读查询，契约 v0.71/v0.77），
// 而不是又发一帧：实时帧向来只推给当时在线的连接，刷新/晚开的页面只能靠只读拉取
// 补齐（`graph_update` 本身不重放）。
//
// 收尾也是**落盘**的唯一时机（切片 F2；design/19 §4 的 P3「多轮历史回放」）：每轮
// 落**一条**（绝不是每 span 一条），真并行下**每一路各一条**；且只在该会话处于发布
// 会话白名单内时发生（本方法对白名单外的会话直接返回 ⇒ 落盘出口天然只被白名单会话
// 触发）。深拷贝与写库在**放锁之后**做：一次库写绝不把其他会话的 span 发布挡住；
// 失败只记 WARN，绝不影响本轮收尾。
func (p *FlowPublisher) EndRound(owner, outputText string) {
	session, _ := splitFlowOwner(owner)
	if session == "" || owner == "" || !p.allowed(session) {
		return
	}
	p.mu.Lock()
	g := p.graphs[owner]
	if g == nil {
		p.mu.Unlock()
		return
	}
	g.SeedOutput(outputText)
	p.publishLocked(g)
	var snap *trace.FlowGraph
	if p.record != nil {
		snap = g.Snapshot().Clone()
	}
	// 本路退出在飞集合（**只删自己的桶**）：集合空了才说明整轮收尾，下一轮才会
	// 拿到新的轮次号 —— 因此轮次号在同轮并行的两路之间保持同一个值。
	delete(p.inflight[session], owner)
	p.mu.Unlock()

	if snap != nil {
		if err := p.record.RecordRound(snap); err != nil {
			p.logPersistError(session, snap.ID, err)
		}
	}
}

// MarkInputSteered 给本路的 input 节点打上 `steered:true` 标记并发布一帧
// （design/15 §5 / 裁决 13.4.5「注入时给 input 节点贴 steered:true 标记」）。
// 它让链路图如实表达「这一路的输入在回合中途被追加过」——图仍然只有一条主干
// （I1：同会话恒只有一个在飞回合），而不是被画成第二轮。
//
// nil 接收者、空归属、发布白名单外的会话、或该路还没有图时都是空操作；标记
// 只并入既有 input 节点（Upsert 的合并语义），绝不新建节点、绝不改动它的标签。
func (p *FlowPublisher) MarkInputSteered(owner string) {
	session, _ := splitFlowOwner(owner)
	if p == nil || owner == "" || session == "" || !p.allowed(session) {
		return
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	g := p.graphs[owner]
	if g == nil {
		return
	}
	g.Upsert(trace.NodeSpec{
		ID:    "input",
		Kind:  "input",
		Props: map[string]any{"steered": true},
	})
	p.publishLocked(g)
}

// SpanStarted 为一个刚开始的 span 记录一个运行中的占位
// 节点。state_transition span 会被图构建器省略。
func (p *FlowPublisher) SpanStarted(owner string, sp *trace.Span) {
	if sp == nil {
		return
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	g := p.bucketLocked(owner)
	if g == nil {
		return
	}
	if g.AddSpanStart(sp.ID, sp.Name, sp.Kind, sp.ParentID, sp.StartedAt) {
		p.publishLocked(g)
	}
}

// SpanFinished 把一个刚结束的 span 作为最终节点（状态 done/error）
// 并入该路的图并发布。必须在 span 的 Finish 之后调用。
func (p *FlowPublisher) SpanFinished(owner string, sp *trace.Span) {
	if sp == nil {
		return
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	g := p.bucketLocked(owner)
	if g == nil {
		return
	}
	g.AddSpan(sp.Snapshot())
	p.publishLocked(g)
}

// bucketLocked 取归属桶当前那张图；归属非法或该路还没有图时返回 nil。调用方必须
// 持有 p.mu。白名单判据与建图处同一处（bucket 拆出的会话），因此**不可能**让
// 白名单外的会话经 span 折叠进图。
func (p *FlowPublisher) bucketLocked(owner string) *trace.RunGraph {
	session, _ := splitFlowOwner(owner)
	if owner == "" || session == "" || !p.allowed(session) {
		return nil
	}
	return p.graphs[owner]
}

// LastGraph 返回 sessionID **最近一路**链路图的一份只读深拷贝快照，供
// 管理面板在「轮次已经结束、实时帧已经错过」时补拉（契约 v0.71 的
// `GET /api/v1/admin/flow/graph`；design/19 §4 的 P3「多轮历史回放」的
// 最小一步）。
//
// 「最近」= 该会话所有在内存的路里**开始最晚**的那一路（同刻用图 id 定序）。
// 真并行时它只是最近的一路，**全貌**是 `LatestLanes`（契约 v0.77 的 `lanes[]`）。
//
// 它是**只读**的：不发布、不改图、不动轮次计数 —— 因此补拉绝不会在
// 管理面板里凭空造出一帧 `graph_update`，也绝不改任何会话的实时图。
//
// 该会话还没有图时返回 nil：从未跑过带 span 的轮次、或该会话本就**不在
// 发布会话范围**（`StartRound` 对白名单外的会话直接返回 ⇒ 从不建图，
// 隔离渠道会话因此天然查不到）。调用方**必须**把 nil 读作「确实没有图」，
// 绝不允许据此补画任何节点或假图。
//
// 返回的是深拷贝（`trace.FlowGraph.Clone`）：读者在锁外编码时，
// 活图可能正被下一帧 span 就地合并，共享底层 map/切片会撕裂。
func (p *FlowPublisher) LastGraph(sessionID string) *trace.FlowGraph {
	if p == nil || sessionID == "" {
		return nil
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	latest := p.latestLocked(sessionID)
	if latest == nil {
		return nil
	}
	return latest.Clone()
}

// LatestLanes 返回 sessionID **最近一轮每一路**的图快照（每条路一份深拷贝），
// 按路的开始时间升序（最早的在前）。契约 v0.77 的 `lanes[]` 就是它：真并行时
// 同一会话同一轮有多条路，`LastGraph` 只给出其中最新的一条，管理面板要靠本方法
// 才看得全（design/19 §2「并行分支独立显示」）。
//
// 没有图时返回空切片（不是 nil）：调用方据此渲染诚实空态，绝不补画。
// 与 LastGraph 一样是**只读**的，且返回深拷贝（读者在锁外编码）。
func (p *FlowPublisher) LatestLanes(sessionID string) []*trace.FlowGraph {
	if p == nil || sessionID == "" {
		return nil
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	round := p.rounds[sessionID]
	owners := make([]string, 0, 2)
	for owner, r := range p.roundOf {
		if r != round {
			continue
		}
		if s, _ := splitFlowOwner(owner); s == sessionID {
			owners = append(owners, owner)
		}
	}
	out := make([]*trace.FlowGraph, 0, len(owners))
	for _, owner := range owners {
		g := p.graphs[owner]
		if g == nil {
			continue
		}
		out = append(out, g.Snapshot().Clone())
	}
	sortLanes(out)
	return out
}

// latestLocked 取该会话开始最晚的那一路的活图（没有时 nil）。调用方必须持有 p.mu。
func (p *FlowPublisher) latestLocked(sessionID string) *trace.FlowGraph {
	var best *trace.FlowGraph
	for owner, g := range p.graphs {
		if s, _ := splitFlowOwner(owner); s != sessionID {
			continue
		}
		snap := g.Snapshot()
		if best == nil || laneNewer(snap, best) {
			best = snap
		}
	}
	return best
}

// laneNewer 报告 a 是否比 b 更「新」：先比开始时刻，同刻用图 id 定序（图 id 里含
// 会话级单调序号，因此定序不依赖时间分辨率）。
func laneNewer(a, b *trace.FlowGraph) bool {
	if a.StartedAt.Equal(b.StartedAt) {
		return a.ID > b.ID
	}
	return a.StartedAt.After(b.StartedAt)
}

// sortLanes 按「开始时刻升序、同刻按图 id」把一路路图排序，使一个轮次里的多条路
// 有一个确定的展示次序（绝不让展示次序随 map 遍历序抖动）。
func sortLanes(lanes []*trace.FlowGraph) {
	sort.SliceStable(lanes, func(i, j int) bool {
		if lanes[i].StartedAt.Equal(lanes[j].StartedAt) {
			return lanes[i].ID < lanes[j].ID
		}
		return lanes[i].StartedAt.Before(lanes[j].StartedAt)
	})
}

// publishLocked 发布 g 的一份快照。调用方必须持有 p.mu。
func (p *FlowPublisher) publishLocked(g *trace.RunGraph) {
	if p.publish == nil {
		return
	}
	p.publish(graphUpdateSubject, g.Snapshot())
}

// allowed 报告某会话的图是否应当发布。
func (p *FlowPublisher) allowed(sessionID string) bool {
	if p.allow == nil {
		return true
	}
	return p.allow(sessionID)
}

// logPersistError 把一次落盘失败记成带来源的 WARN。它只记 session/turn 标识，
// **绝不打印图内容**（图里可能含用户输入文本）；logger 未接线时静默。
func (p *FlowPublisher) logPersistError(sessionID, turnID string, err error) {
	if p.log == nil {
		return
	}
	logger.WithSource(p.log, flowLogSource).Warn("flow round persist failed",
		"session", sessionID, "turn", turnID, "err", err)
}

// lastUserText 返回 messages 中最新用户消息的文本，用于
// 标注该轮的 input 节点。没有用户文本时返回 ""。
func lastUserText(messages []model.Message) string {
	for i := len(messages) - 1; i >= 0; i-- {
		if messages[i].Role != model.RoleUser {
			continue
		}
		if t := messageText(messages[i]); t != "" {
			return t
		}
	}
	return ""
}