package flow

import (
	"ai-agent/gateway"
	"ai-agent/workflow/consts/node"
	flowDto "ai-agent/workflow/model/dto/flow"
	"ai-agent/workflow/model/entity"
	"context"
	"fmt"
	"time"

	"github.com/cloudwego/eino/compose"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/util/gconv"
)

// init 注册自定义合并函数：处理 *flowDto.FlowExecutionInput 类型合并。
// 合并函数全局唯一且与图内容无关，放包级 init 注册一次（避免每次 BuildGraph 重注册全局状态）。
func init() {
	compose.RegisterValuesMergeFunc(func(values []*flowDto.FlowExecutionInput) (*flowDto.FlowExecutionInput, error) {
		if len(values) == 0 {
			return nil, nil
		}
		// 首次运行所有并行分支共享同一个 ConfigMap 指针，直接返回 values[0] 即可。
		// 但续跑（ReExecute）时各分支从 checkpoint 反序列化出独立的 ConfigMap 副本，
		// 只返回 values[0] 会丢失其他分支写入的 OutputResult（用户实测：node-8 成功的结果
		// 在汇合节点 node-7 变 null）。以第一个为基底，把其余分支中缺失的节点输出合并进来。
		base := values[0]
		for _, v := range values[1:] {
			if v == nil {
				continue
			}
			for nodeId, cfg := range v.ConfigMap {
				if cfg == nil {
					continue
				}
				baseCfg, ok := base.ConfigMap[nodeId]
				if !ok || baseCfg == nil {
					base.ConfigMap[nodeId] = cfg
					continue
				}
				if len(baseCfg.OutputResult) == 0 && len(cfg.OutputResult) > 0 {
					baseCfg.OutputResult = cfg.OutputResult
				}
			}
			// 合并已执行节点列表（按 NodeId 去重），续跑时被恢复分支的进度不丢失
			for _, en := range v.ExecutedNodes {
				dup := false
				for _, b := range base.ExecutedNodes {
					if b.NodeId == en.NodeId {
						dup = true
						break
					}
				}
				if !dup {
					base.ExecutedNodes = append(base.ExecutedNodes, en)
				}
			}
		}
		return base, nil
	})
}

// BuildGraph 根据 FlowInfo 构建完整的 Eino Graph 拓扑
func BuildGraph(ctx context.Context, flowContent *entity.FlowInfo) ([]entity.FlowNode, *compose.Graph[any, any]) {
	graph := compose.NewGraph[any, any](
		// 本地状态初始化
		compose.WithGenLocalState(func(ctx context.Context) *flowDto.NodeExecutionState {
			return &flowDto.NodeExecutionState{}
		}),
	)

	// 注册所有节点
	for _, item := range flowContent.Nodes {
		registerNodeToGraph(graph, item)
	}

	// 注册开始节点
	if flowContent.StartNodeId != "" {
		_ = graph.AddEdge(compose.START, flowContent.StartNodeId)
	}

	var nodeList []entity.FlowNode
	originalEndNodes := FindEndNodes(flowContent.StartNodeId, flowContent.Edges)
	for _, endID := range originalEndNodes {
		// 保存结果节点 ID 必须稳定：ReExecute 续跑重建图时复用同一 ID，
		// 否则 checkpoint 里 ConfigMap 存的是上次的旧 ID，续跑时新图按新 ID 查不到配置，
		// 报"节点信息为空"。一个末端节点对应一个保存结果节点，用 endID 派生即唯一且稳定。
		summaryNodeId := fmt.Sprintf("%s_%s", node.NodeTypeSystemSum, endID)
		summaryNode := entity.FlowNode{
			Id:       summaryNodeId,
			NodeCode: node.NodeTypeSystemSum,
			Name:     node.GetNodeTypeName(node.NodeTypeSystemSum),
		}
		nodeList = append(nodeList, summaryNode)
		flowContent.Nodes = append(flowContent.Nodes, summaryNode)

		registerNodeToGraph(graph, summaryNode)
		_ = graph.AddEdge(endID, summaryNodeId)
		_ = graph.AddEdge(summaryNodeId, compose.END)
	}

	// 构建边关系
	edgeMap := make(map[string][]entity.FlowEdge)
	for _, edge := range flowContent.Edges {
		edgeMap[edge.From] = append(edgeMap[edge.From], edge)
	}

	// 处理连线 & 分支
	for _, edges := range edgeMap {
		// 普通节点连线
		for _, e := range edges {
			_ = graph.AddEdge(e.From, e.To)
		}
	}
	return nodeList, graph
}

// BuildGraphFromFlowContent 根据前端保存的工作流JSON，自动构建执行图并编译
func BuildGraphFromFlowContent(ctx context.Context, flowContent *entity.FlowInfo) ([]entity.FlowNode, compose.Runnable[any, any], error) {
	nodeList, graph := BuildGraph(ctx, flowContent)
	// BuildGraph 已把 summary(保存结果)节点追加进 flowContent.Nodes，此时是全部已注册节点的完整集合。
	// 方案: 每个业务节点正常完成后自动暂停并落 checkpoint(编译期 WithInterruptAfterNodes),
	// 崩溃恢复(BuildExecution(false)续跑)即跳过已完成的同步节点, 不再重跑/重复计费。
	// 详见根目录《工作流节点断点续跑技术设计.md》。Start 型空载节点不为它落 cp; 只接 END 的节点
	// Eino 不落暂停(无下游续跑点), 列了也无副作用。
	interruptAfter := make([]string, 0, len(flowContent.Nodes))
	for _, n := range flowContent.Nodes {
		if n.NodeCode == node.NodeTypeStart {
			continue
		}
		interruptAfter = append(interruptAfter, n.Id)
	}
	compile, err := graph.Compile(ctx,
		compose.WithGraphName("auto_build_workflow"),
		compose.WithCheckPointStore(NewDbCheckPointStore()),
		compose.WithNodeTriggerMode(compose.AllPredecessor),
		compose.WithInterruptAfterNodes(interruptAfter),
	)
	return nodeList, compile, err
}

// buildConfigMap 由 FlowInfo + 图节点列表构建 ConfigMap：先放流程配置节点，再放图中补充节点
// （保存结果节点等），供节点执行时按 nodeId 查配置。nodeList 取自 BuildGraph 返回值。
func buildConfigMap(flowContent *entity.FlowInfo, nodeList []entity.FlowNode) map[string]*entity.FlowNode {
	configMap := make(map[string]*entity.FlowNode)
	for _, cfg := range ExtractFlowNodeFrom(flowContent) {
		configMap[cfg.Id] = cfg
	}
	for i := range nodeList {
		configMap[nodeList[i].Id] = &nodeList[i]
	}
	return configMap
}

// registerNodeToGraph 将单个节点注册到图中（包含通用包装逻辑）
func registerNodeToGraph(graph *compose.Graph[any, any], flowNode entity.FlowNode) {
	// 通用包装：全程入参都是 *FlowExecutionInput
	wrapLambda := func(lambda func(ctx context.Context, input any) (any, error)) func(ctx context.Context, input any) (any, error) {
		return func(ctx context.Context, input any) (any, error) {
			startTime := time.Now()

			// 构建节点执行入参（含中断恢复）
			execInput, realInput, err := BuildNodeExecutionInput(ctx, input, flowNode)
			if err != nil {
				return nil, err
			}

			flowNodeDesc := flowNode.Desc
			if g.IsEmpty(flowNodeDesc) {
				flowNodeDesc = flowNode.Name
			}

			// 上报节点执行进度（WebSocket场景下推送进度给前端）
			if reporter := GetProgressReporter(ctx); reporter != nil {
				reporter.ReportStart(flowNode.Id, flowNodeDesc, nodeReportIndex(execInput, flowNode.Id, 1), len(execInput.ConfigMap))
			}

			// 上传入参到OSS
			ossResult, err := gateway.Upload(ctx, fmt.Sprintf("nodeInput:%v.txt", time.Now().UnixMilli()), gconv.Bytes(gconv.String(realInput)))
			if err != nil {
				return nil, HandleFailedNodeExecution(ctx, execInput, 0, flowNode, err, 0)
			}

			// 创建节点执行记录
			nodeExecutionId, err := CreateNodeExecutionRecord(ctx, execInput, flowNode, ossResult)
			if err != nil {
				return nil, HandleFailedNodeExecution(ctx, execInput, 0, flowNode, err, 0)
			}
			realInput.NodeExecutionId = nodeExecutionId

			// 执行节点
			_, err = lambda(ctx, realInput)
			durationMs := time.Since(startTime).Milliseconds()

			if err != nil {
				// 执行失败处理
				return nil, HandleFailedNodeExecution(ctx, execInput, nodeExecutionId, flowNode, err, durationMs)
			}

			// 执行成功处理
			if err = HandleSuccessfulNodeExecution(ctx, execInput, realInput, nodeExecutionId, flowNode, durationMs); err != nil {
				return nil, HandleFailedNodeExecution(ctx, execInput, nodeExecutionId, flowNode, err, durationMs)
			}

			// 上报节点执行进度（WebSocket场景下推送进度给前端）
			if reporter := GetProgressReporter(ctx); reporter != nil {
				reporter.ReportComplete(flowNode.Id, flowNodeDesc, nodeReportIndex(execInput, flowNode.Id, 0), len(execInput.ConfigMap))
			}

			// 返回整个 execInput，让下一个节点继续用
			return execInput, nil
		}
	}

	switch flowNode.NodeCode {
	case node.NodeTypeStart:
		_ = graph.AddLambdaNode(flowNode.Id, compose.InvokableLambda(wrapLambda(StartLambda)))
	case node.NodeTypeSystemSum:
		_ = graph.AddLambdaNode(flowNode.Id, compose.InvokableLambda(wrapLambda(SummaryLambda)))
	case node.NodeTypeModel:
		_ = graph.AddLambdaNode(flowNode.Id, compose.InvokableLambda(wrapLambda(ModelLambda)))
	case node.NodeTypeForm:
		_ = graph.AddLambdaNode(flowNode.Id, compose.InvokableLambda(wrapLambda(FormLambda)))
	case node.NodeTypeDataMerge:
		_ = graph.AddLambdaNode(flowNode.Id, compose.InvokableLambda(wrapLambda(DataMergeLambda)))
	case node.NodeTypeSubFlow:
		_ = graph.AddLambdaNode(flowNode.Id, compose.InvokableLambda(wrapLambda(SubFlowLambda)))
	case node.NodeTypeHttp:
		_ = graph.AddLambdaNode(flowNode.Id, compose.InvokableLambda(wrapLambda(HttpLambda)))
	case node.NodeTypeScriptTranscribe:
		_ = graph.AddLambdaNode(flowNode.Id, compose.InvokableLambda(wrapLambda(ScriptTranscribeLambda)))
	}
}

// nodeReportIndex 计算节点在进度上报中的序号：节点已在已执行列表则取其位置，
// 否则按当前已执行数 + offset（start 上报时节点尚未入列 offset=1；complete 后已入列命中 IndexOf）
func nodeReportIndex(execInput *flowDto.FlowExecutionInput, nodeId string, offset int) int {
	if idx := IndexOf(execInput.ExecutedNodes, nodeId); idx != -1 {
		return idx
	}
	return len(execInput.ExecutedNodes) + offset
}

// IndexOf 返回元素第一次出现的下标，不存在返回 -1
func IndexOf(slice []flowDto.ExecutedNode, target string) int {
	for i, v := range slice {
		if v.NodeId == target {
			return i + 1
		}
	}
	return -1
}