package eino

import (
	"context"

	"github.com/cloudwego/eino/components/prompt"
	"github.com/cloudwego/eino/compose"
	"github.com/cloudwego/eino/schema"

	"github.com/sundynix/sundynix-dispatcher/internal/llm"
	"github.com/sundynix/sundynix-shared/contract"
)

// memoryFetcher 召回某用户与本次输入相关的偏好记忆（经 MCP memory_get 工具）。
type memoryFetcher func(ctx context.Context, userID, query string) string

// buildGraph 编译这套"记忆增强"图：
//
//	START → recall(召回画像→写State) → prompt(注入system) → model(流式) → END
//
// 返回可流式执行的 Runnable。
func buildGraph(ctx context.Context, pool *llm.Pool, fetch memoryFetcher) (compose.Runnable[*contract.Task, *schema.Message], error) {
	g := compose.NewGraph[*contract.Task, *schema.Message](
		compose.WithGenLocalState(func(context.Context) *AgentState { return &AgentState{} }),
	)

	// 1) recall：取 user_id → memory_get 召回画像 → 写入 State，并输出模板变量。
	if err := g.AddLambdaNode("recall", compose.InvokableLambda(
		func(ctx context.Context, t *contract.Task) (map[string]any, error) {
			uid, _ := t.Meta[contract.MetaUserID].(string)
			profile := fetch(ctx, uid, string(t.Graph))
			_ = compose.ProcessState(ctx, func(_ context.Context, s *AgentState) error {
				s.UserID, s.Profile, s.Input = uid, profile, string(t.Graph)
				return nil
			})
			if profile == "" {
				profile = "（暂无该用户的偏好记忆）"
			}
			return map[string]any{"profile": profile, "query": string(t.Graph)}, nil
		})); err != nil {
		return nil, err
	}

	// 2) prompt：把画像注入 system message，用户输入作为 user message。
	tpl := prompt.FromMessages(schema.FString,
		schema.SystemMessage("你在与特定用户对话。关于该用户的已知信息：\n{profile}\n请据此个性化作答并保持其偏好。"),
		schema.UserMessage("{query}"),
	)
	if err := g.AddChatTemplateNode("prompt", tpl); err != nil {
		return nil, err
	}

	// 3) model：LLM Pool 适配为 ChatModel 节点，流式产出。
	if err := g.AddChatModelNode("model", newPoolModel(pool)); err != nil {
		return nil, err
	}

	if err := g.AddEdge(compose.START, "recall"); err != nil {
		return nil, err
	}
	if err := g.AddEdge("recall", "prompt"); err != nil {
		return nil, err
	}
	if err := g.AddEdge("prompt", "model"); err != nil {
		return nil, err
	}
	if err := g.AddEdge("model", compose.END); err != nil {
		return nil, err
	}

	return g.Compile(ctx)
}