package einograph

import (
	"context"

	"github.com/cloudwego/eino/compose"
)

// 节点名称常量
const (
	NodeInputAdapt  = "input_adapt"  // 输入适配
	NodeBuildPrompt = "build_prompt" // 构建提示词
	NodeGenText     = "gen_text"     // 大模型生成原始文案
	NodeGenImage    = "gen_image"    // 生成图片 + HTML
)

// ==================== 构建流式图（已加入优化节点） ====================
func BuildCreationGraph(ctx context.Context) (compose.Runnable[any, any], error) {
	g := compose.NewGraph[any, any]()

	// 1. 输入适配
	err := g.AddLambdaNode(
		NodeInputAdapt,
		compose.InvokableLambda(InputAdaptLambda),
		compose.WithOutputKey("input"),
	)
	if err != nil {
		return nil, err
	}

	// 2. 构建提示词
	err = g.AddLambdaNode(
		NodeBuildPrompt,
		compose.InvokableLambda(BuildPromptLambda),
		compose.WithInputKey("input"),
		compose.WithOutputKey("messages"),
	)
	if err != nil {
		return nil, err
	}

	// 3. LLM 生成原始文案
	chatModel, err := NewChatModel(ctx)
	if err != nil {
		return nil, err
	}
	err = g.AddChatModelNode(
		NodeGenText,
		chatModel,
		compose.WithInputKey("messages"),
		compose.WithOutputKey("text_result"),
	)
	if err != nil {
		return nil, err
	}

	// 5. 生成图片 + HTML（使用优化后的文案）
	err = g.AddLambdaNode(
		NodeGenImage,
		compose.InvokableLambda(GenerateImageLambda),
		compose.WithInputKey("text_result"),
		compose.WithOutputKey("output"),
	)
	if err != nil {
		return nil, err
	}

	// ✅ 正确连线：优化节点必须放在生成文案之后、生成图片之前
	_ = g.AddEdge(compose.START, NodeInputAdapt)
	_ = g.AddEdge(NodeInputAdapt, NodeBuildPrompt)
	_ = g.AddEdge(NodeBuildPrompt, NodeGenText)
	_ = g.AddEdge(NodeGenText, NodeGenImage)
	_ = g.AddEdge(NodeGenImage, compose.END)

	// 编译
	return g.Compile(ctx,
		compose.WithGraphName("xiaohongshu_content_creation"),
		compose.WithNodeTriggerMode(compose.AnyPredecessor),
	)
}