package flow

import (
	"ai-agent/workflow/consts/node"
	flowDto "ai-agent/workflow/model/dto/flow"
	"ai-agent/workflow/model/entity"
	"ai-agent/workflow/service/flow/values"
	"context"
	"encoding/json"
	"fmt"

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

func SubFlowLambda(ctx context.Context, input any) (any, error) {
	// 1. 类型断言（和其他节点保持一致的入参结构）
	nodeExecInput, ok := input.(*flowDto.NodeExecutionInput)
	if !ok {
		return nil, fmt.Errorf("子流程节点入参类型错误，期望*flowDto.NodeExecutionInput，实际%T", input)
	}
	// 2. 解析子流程配置
	subFlowConfig := nodeExecInput.Config.SubConfig
	if subFlowConfig == nil {
		return nil, fmt.Errorf("子流程节点缺少配置")
	}
	getRes, err := FlowUserService.Get(ctx, &flowDto.GetFlowUserReq{
		Id: subFlowConfig.WorkflowId,
	})
	if err != nil {
		return nil, err
	}
	// 3. 引入参数解析：把首页表单值/上游引用值/静态默认值写入子流程开始节点 outputConfig。
	//    须在 BuildGraph / ExtractFlowNodeFrom 之前执行，batchInputs 深拷贝的才是注入后的开始节点。
	injectSubFlowFields(nodeExecInput.Global, getRes.FlowContent, subFlowConfig.Fields)
	// 4. 并发数：从主流程开始节点 outputConfig 的 maxConcurrency 字段读取（前端把子流程节点生成次数表单字段聚合到主流程开始节点），读不到再用子流程节点配置兜底
	maxConcurrency := mainFlowMaxConcurrency(nodeExecInput.Global, subFlowConfig.MaxConcurrency)
	// 4. 编译子流程Graph（复用现有 BuildGraphFromFlowContent 逻辑）
	nodeList, subGraph := BuildGraph(ctx, getRes.FlowContent)
	// 4. 构建子流程Workflow（绑定START/END，和示例对齐）
	innerWorkflow := compose.NewWorkflow[*flowDto.FlowExecutionInput, *flowDto.FlowExecutionInput]()
	// 挂载子图节点并绑定全局START
	innerWorkflow.AddGraphNode("sub_flow_graph", subGraph).AddInput(compose.START)
	// 绑定子图输出到全局END
	innerWorkflow.End().AddInput("sub_flow_graph")
	// 生成次数（批量条数）：maxConcurrency<=0 时按 1 次兜底
	batchCount := maxConcurrency
	if batchCount <= 0 {
		batchCount = 1
	}
	// 5. 构建BatchNode（批量执行子流程，复用示例逻辑）
	batchNode := batch.NewBatchNode(&batch.NodeConfig[*flowDto.FlowExecutionInput, *flowDto.FlowExecutionInput]{
		Name:           fmt.Sprintf("sub_flow_batch_%s", nodeExecInput.Config.Id),
		InnerTask:      innerWorkflow,
		MaxConcurrency: batchCount,
	})

	// 6. 提取批量输入：按生成次数生成 N 份（每份独立克隆 ConfigMap，避免并发执行时节点输出写串）
	configMap := buildConfigMap(getRes.FlowContent, nodeList)
	batchInputs := make([]*flowDto.FlowExecutionInput, 0, batchCount)
	for j := 0; j < batchCount; j++ {
		// 缓存作用域：拼本(外层)子流程节点 Id + 份号 j，供 async/segment 缓存键隔离各份
		// （同一 exec 下 N 份内层节点 id 相同，无作用域会互相命中/覆盖 done 结果，见《工作流子流程批量缓存隔离设计.md》）。
		// 嵌套子流程继承父 scope（Global.SubFlowScope）递归叠加，保证跨 exec/嵌套路径唯一；
		// 跨 launch 续跑按同序重放 batch → 各份 scope 稳定，仍能命中自己那份结果。
		scope := nodeExecInput.Global.SubFlowScope + fmt.Sprintf("[%s#%d]", nodeExecInput.Config.Id, j)
		batchInputs = append(batchInputs, &flowDto.FlowExecutionInput{
			NodeGroupId:  nodeExecInput.Global.NodeGroupId,
			ExecutionId:  nodeExecInput.Global.ExecutionId,
			FlowId:       nodeExecInput.Global.FlowId,
			ConfigMap:    cloneConfigMap(configMap),
			SessionId:    nodeExecInput.Global.SessionId,
			SubFlowScope: scope,
		})
	}
	// 7. 执行批量子流程
	batchOutput, err := batchNode.Invoke(ctx, batchInputs)
	if err != nil {
		return nil, fmt.Errorf("执行子流程BatchNode失败: %v", err)
	}
	// 8. 展平每份子流程执行的节点输出，写回当前节点 OutputResult 供下游引用
	var outputRes []map[string]any
	for _, single := range batchOutput {
		if single == nil {
			continue
		}
		outputRes = append(outputRes, collectFlowNodeResults(single)...)
	}
	g.Log().Info(ctx, fmt.Sprintf("子流程执行完成，共 %d 次，输出 %d 条", batchCount, len(outputRes)))
	nodeExecInput.Config.OutputResult = outputRes
	return nodeExecInput, nil
}

// injectSubFlowFields 将子流程节点引入参数（subConfig.Fields）解析后写入子流程开始节点
// outputConfig，使子流程启动时能读到首页表单值/上游引用值/静态默认值。
// 每个字段的取值优先级：valueSource 引用解析成功 → field.value → field.defaultValue；
// 匹配键为 field（前端约定以 field 为主，不兼容 path）。
func injectSubFlowFields(global *flowDto.FlowExecutionInput, subFlowContent *entity.FlowInfo, fields []map[string]any) {
	if global == nil || subFlowContent == nil || len(fields) == 0 {
		return
	}
	startNode := subFlowStartNode(subFlowContent)
	if startNode == nil {
		return
	}
	byField := make(map[string]map[string]any, len(startNode.OutputConfig))
	for _, output := range startNode.OutputConfig {
		byField[gconv.String(output["field"])] = output
	}
	for _, field := range fields {
		entry := byField[gconv.String(field["field"])]
		if entry == nil {
			continue
		}
		value := field["value"]
		if vs, has := field["valueSource"]; has && vs != nil {
			if vsNodeId, vsField := firstValueSource(vs); vsNodeId != "" && vsField != "" {
				if v, _, ok := values.ResolveValueSource(global, vsNodeId, vsField); ok {
					value = v
				}
			}
		}
		if value == nil {
			value = field["defaultValue"]
		}
		if value != nil {
			entry["value"] = value
		}
	}
}

// firstValueSource 从 valueSource 提取第一个引用源 (nodeId, field)。
// 前端契约统一数组 [{nodeId, field}]，旧 DSL 可能是单对象 {nodeId, field}，两种形态都兼容；
// 子流程字段与引用源一一对应，只取第一个。
func firstValueSource(vs any) (nodeId, field string) {
	if vs == nil {
		return
	}
	switch v := vs.(type) {
	case []any:
		if len(v) > 0 {
			return firstValueSource(v[0])
		}
		return
	case []map[string]any:
		if len(v) > 0 {
			return firstValueSource(v[0])
		}
		return
	}
	m := gconv.Map(vs)
	nodeId = gconv.String(m["nodeId"])
	field = gconv.String(m["field"])
	return
}

// subFlowStartNode 返回工作流开始节点
func subFlowStartNode(content *entity.FlowInfo) *entity.FlowNode {
	if content == nil {
		return nil
	}
	for i := range content.Nodes {
		if content.Nodes[i].Id == content.StartNodeId {
			return &content.Nodes[i]
		}
	}
	return nil
}

// mainFlowMaxConcurrency 取子流程批量执行并发数：从主流程开始节点 outputConfig
// 的 maxConcurrency 字段读取（前端把子流程节点的生成次数表单字段聚合到主流程开始节点），
// 读不到再用子流程节点配置的兜底值。
func mainFlowMaxConcurrency(global *flowDto.FlowExecutionInput, fallback int) int {
	if global == nil {
		return fallback
	}
	for _, n := range global.ConfigMap {
		if n == nil || n.NodeCode != node.NodeTypeStart {
			continue
		}
		for _, output := range n.OutputConfig {
			if gconv.String(output["field"]) != "maxConcurrency" {
				continue
			}
			if v := gconv.Int(output["value"]); v > 0 {
				return v
			}
		}
		return fallback
	}
	return fallback
}

// cloneConfigMap 深拷贝 ConfigMap，保证各批次子流程并发执行时节点输出互不串扰。
// 浅拷贝会共享 *entity.FlowNode，并发写 OutputResult 产生竞态。
func cloneConfigMap(src map[string]*entity.FlowNode) map[string]*entity.FlowNode {
	dst := make(map[string]*entity.FlowNode, len(src))
	for k, v := range src {
		data, err := json.Marshal(v)
		if err != nil {
			dst[k] = v
			continue
		}
		n := new(entity.FlowNode)
		if err = json.Unmarshal(data, n); err != nil {
			dst[k] = v
			continue
		}
		dst[k] = n
	}
	return dst
}

// collectFlowNodeResults 收集一次子流程执行中所有已执行节点的输出，展平成 {字段:值} 列表
func collectFlowNodeResults(execInput *flowDto.FlowExecutionInput) []map[string]any {
	var res []map[string]any
	for _, executed := range execInput.ExecutedNodes {
		if nodeConfig := execInput.ConfigMap[executed.NodeId]; nodeConfig != nil {
			res = append(res, nodeConfig.OutputResult...)
		}
	}
	return res
}