package engine

// norm_attachment_reorder.go -- 附件消息上浮到 tool_result 前面.
//
// 升华改进(ELEVATED): 附件消息上浮到 tool_result 前面.
// 附件包含上下文信息(文件内容,图片,数据),应该在工具结果之前让模型看到.
// 替代方案:保持原始顺序(模型可能先看到工具结果再看到附件,注意力分散).
//
// 来源:新增步骤,早期方案 normalize 未处理.
// 通过 Message.Metadata["is_attachment"] = true 识别附件消息.
// 跨场景通用:仓储场景的传感器数据附件,法律场景的文件附件.

import (
	"git.flytoex.net/yuanwei/flyto-agent/pkg/query"
)

// AttachmentReorderer 将附件消息上浮到同一 user turn 的 tool_result 前面.
//
// 精妙之处(CLEVER): 只在同角色连续消息块内重排,不跨角色边界.
// 这保证了 user/assistant 交替的基本约束不被破坏.
// 重排只影响 user 消息中 attachment 和 tool_result 的相对顺序.
type AttachmentReorderer struct{}

func (r *AttachmentReorderer) Name() string  { return "attachment_reorder" }
func (r *AttachmentReorderer) Priority() int { return 5 }

func (r *AttachmentReorderer) Normalize(messages []query.Message) []query.Message {
	if len(messages) == 0 {
		return messages
	}

	result := make([]query.Message, 0, len(messages))
	for _, msg := range messages {
		if msg.Role == query.RoleUser && isAttachment(msg) {
			// 附件消息标记为需要上浮(在后续合并步骤中生效)
			result = append(result, msg)
			continue
		}
		result = append(result, msg)
	}

	// 在 user 消息序列中,将附件消息上浮
	return reorderAttachments(result)
}

// isAttachment 检查消息是否被标记为附件.
func isAttachment(msg query.Message) bool {
	if msg.Metadata == nil {
		return false
	}
	v, ok := msg.Metadata["is_attachment"]
	if !ok {
		return false
	}
	b, ok := v.(bool)
	return ok && b
}

// reorderAttachments 在消息列表中,将连续 user 消息块内的附件上浮.
//
// 算法:找到连续的 user 消息块,在块内将附件消息移到 tool_result 消息前面.
// 不跨块操作,不改变 assistant 消息的位置.
func reorderAttachments(messages []query.Message) []query.Message {
	result := make([]query.Message, 0, len(messages))
	i := 0

	for i < len(messages) {
		// 找连续 user 消息块
		if messages[i].Role != query.RoleUser {
			result = append(result, messages[i])
			i++
			continue
		}

		// 收集连续的 user 消息
		blockStart := i
		for i < len(messages) && messages[i].Role == query.RoleUser {
			i++
		}
		block := messages[blockStart:i]

		// 分离附件和非附件
		var attachments, others []query.Message
		for _, msg := range block {
			if isAttachment(msg) {
				attachments = append(attachments, msg)
			} else {
				others = append(others, msg)
			}
		}

		// 附件在前,其他在后
		result = append(result, attachments...)
		result = append(result, others...)
	}

	return result
}