package flow

import (
	"ai-agent/gateway"
	flowDao "ai-agent/workflow/dao/flow"
	nodeDao "ai-agent/workflow/dao/node"
	flowDto "ai-agent/workflow/model/dto/flow"
	nodeDto "ai-agent/workflow/model/dto/node"
	"ai-agent/workflow/model/entity"
	"ai-agent/workflow/service/flow/processor/builtin/media"
	"ai-agent/workflow/service/flow/values"
	"context"
	"fmt"
	"sync"

	"github.com/gogf/gf/v2/util/gconv"
)

// StartLambda 启动节点
func StartLambda(ctx context.Context, input any) (any, error) {
	return input, nil
}

// FormLambda 表单调用节点
func FormLambda(ctx context.Context, input any) (any, error) {
	nodeInput, ok := input.(*flowDto.NodeExecutionInput)
	if !ok {
		return nil, fmt.Errorf("入参类型错误")
	}
	// 解析 valueSource 引用，填充表单节点输出配置（供下游引用）
	for _, output := range nodeInput.Config.OutputConfig {
		values.ProcessValueSourceRecursive(output, nodeInput.Global)
	}
	return nodeInput, nil
}

// cacheNodeId 缓存键节点 id：scope 非空（子流程批量子执行）时拼上作用域后缀，使各份子执行的
// async/segment 缓存行互相隔离（同一 exec 下多份内层节点 id 相同，不隔离会互相命中/覆盖 done 结果，
// 见《工作流子流程批量缓存隔离设计.md》）；顶层 scope 为空 → 原样返回，行为不变。
// 仅用于缓存读写键，节点的 Config.Id / node_execution 记录 / 输出引用一律不受影响。
func cacheNodeId(global *flowDto.FlowExecutionInput, nodeId string) string {
	if global == nil || global.SubFlowScope == "" {
		return nodeId
	}
	return nodeId + global.SubFlowScope
}

// ModelLambda 模型调用节点
func ModelLambda(ctx context.Context, input any) (any, error) {
	nodeInput, ok := input.(*flowDto.NodeExecutionInput)
	if !ok {
		return nil, fmt.Errorf("入参类型错误")
	}

	modelParams, err := values.BuildModelRequestBody(nodeInput.Config.ModelConfig.ModelRequestParamsPath, nodeInput.Global)
	if err != nil {
		return nil, err
	}
	// async/segment 缓存键节点 id（子流程批量子执行带 scope，顶层即 Config.Id）
	cNodeId := cacheNodeId(nodeInput.Global, nodeInput.Config.Id)

	// 2. 前置工具：决定模型调用入参（单次/多次）
	// 入参统一为扁平模型请求体（BuildModelRequestBody 输出，key 为点分路径）。
	// 分批处理器按默认上限拆分集合字段，其余前置工具（如 split_shots_pipeline）读取扁平参数。
	preToolParams := modelParams
	paramsList, err := invokePreTool(ctx, nodeInput.Config.PreTool, preToolParams)
	if err != nil {
		return nil, err
	}

	// 3. 逐批调用模型，汇总输出（保持请求顺序），累计 token/费用供节点记录落库
	var outputRes []map[string]any
	var totalTokens int64
	var totalPrompt int64
	var totalCompletion int64
	var totalCost float64
	var totalDuration int64
	// 计价用生效模型 id：引用行由 model-gateway 解析为系统模型 id（ModelCallRes.ModelId，计价按系统模型）；
	// 未返回（model-gateway 旧版本）时回落节点配置的模型 id。media_type 供 per_token 命中媒体价。
	effModelID := nodeInput.Config.ModelConfig.ModelId
	var effMediaType string
	if len(paramsList) > 1 {
		// 段级续跑仅在"多段 + 视频模型"启用；非视频分段（批量文本等）走原逻辑零影响。
		// 段身份 = 列表位置（0-based）：paramsList 顺序即段序，concat 按列表顺序拼接；
		// 位置互不重复且跨 reExecute 稳定（参数一致 → 段数/顺序不变）。不依赖 params 里的
		// segment_index——真实链路（上游 split_shots_pipeline 转写 → 下游 split_segment 按
		// __segment_fields 拆分，invokePreTool 剥离 __ 内部键）下 paramsList 只有模型参数。
		// segVideo=false 时走既有非段级合并路径（全量生成、不落库、不复用），全新执行行为不变，
		// 后续自动 concat 判断（独立的 isVideoModel 调用）仍正常执行。
		segVideo := isVideoModel(ctx, nodeInput.Config.ModelConfig.ModelId)

		// 续跑（!ForceNewRun）时读取该节点已成功段；全新执行不查（BuildExecution 已清旧段），saved 为 nil → 全量重生成
		var saved map[int]entity.SegmentRef
		if !nodeInput.Global.ForceNewRun && segVideo {
			saved, err = flowDao.FlowSegmentResultDao.ListByNode(ctx, nodeInput.Global.ExecutionId, cNodeId)
			if err != nil {
				return nil, err
			}
		}

		idxList, needGen := planSegmentResume(paramsList, saved)

		results := make([][]map[string]any, len(paramsList))
		tokenRes := make([]*gateway.ModelCallRes, len(paramsList))
		errs := make([]error, len(paramsList))
		saveErrs := make([]error, len(paramsList))
		isInference := make([]bool, len(paramsList))
		var wg sync.WaitGroup
		for i, params := range paramsList {
			if !needGen[i] {
				continue
			}
			wg.Add(1)
			go func(i int, params map[string]any) {
				defer wg.Done()
				// 每段单次调用，不原地重试：段失败即走节点失败收口（HandleFailedNodeExecution → Interrupt），
				// 下次 reExecute 由 planSegmentResume 复用已成功段、仅重生成失败段
				results[i], tokenRes[i], isInference[i], errs[i] = ModelCallResultLambda(ctx, nodeInput.Config.ModelConfig.ModelId, nodeInput.Global.SessionId, params, nodeInput.Config.Prompt, nodeInput.Global.ExecutionId, cNodeId, idxList[i])
				// 每段成功立即落库：该段刚成功即持久化，其他段仍在跑时已成功段也不丢；
				// 后续段失败或进程崩溃（panic/OOM/kill）时，已完成段已在库中，reExecute 可直接复用
				if segVideo && errs[i] == nil {
					for _, rec := range results[i] {
						key := media.FindVideoKey(rec)
						url := media.FindVideoURL(ctx, rec)
						if key == "" || url == "" {
							continue
						}
						if err := flowDao.FlowSegmentResultDao.Save(ctx, nodeInput.Global.ExecutionId, cNodeId, idxList[i], key, url); err != nil {
							saveErrs[i] = err
						}
					}
				}
			}(i, params)
		}
		wg.Wait()

		// 仍有失败段或落库失败 → 节点失败（成功段已立即落库，供下次 reExecute 复用）
		for i := range results {
			if saveErrs[i] != nil {
				return nil, saveErrs[i]
			}
			if needGen[i] && errs[i] != nil {
				return nil, errs[i]
			}
			if needGen[i] && tokenRes[i] != nil {
				totalTokens += tokenRes[i].TotalTokens
				totalPrompt += tokenRes[i].PromptTokens
				totalCompletion += tokenRes[i].CompletionTokens
				totalCost += tokenRes[i].Cost
				if isVideoModel(ctx, nodeInput.Config.ModelConfig.ModelId) {
					totalDuration += tokenRes[i].Duration
				}
				if tokenRes[i].ModelId > 0 {
					effModelID = tokenRes[i].ModelId
				}
				if effMediaType == "" {
					effMediaType = tokenRes[i].MediaType
				}
			}
		}

		if segVideo {
			// 复用段 + 新生段按段序号升序合并，concat 按列表顺序拼接 → 顺序保证
			outputRes = mergeSegmentOutputs(idxList, needGen, results, saved)
		} else {
			if isInference[0] {
				outputRes = mergeInferenceBatchResults(results)
			} else {
				for _, res := range results {
					outputRes = append(outputRes, res...)
				}
			}
		}
	} else {
		for _, params := range paramsList {
			res, modelRes, _, err := ModelCallResultLambda(ctx, nodeInput.Config.ModelConfig.ModelId, nodeInput.Global.SessionId, params, nodeInput.Config.Prompt, nodeInput.Global.ExecutionId, cNodeId, flowDao.FlowAsyncSegSentinel)
			if err != nil {
				return nil, err
			}
			if modelRes != nil {
				totalTokens += modelRes.TotalTokens
				totalPrompt += modelRes.PromptTokens
				totalCompletion += modelRes.CompletionTokens
				totalCost += modelRes.Cost
				if isVideoModel(ctx, nodeInput.Config.ModelConfig.ModelId) {
					totalDuration += modelRes.Duration
				}
				if modelRes.ModelId > 0 {
					effModelID = modelRes.ModelId
				}
				if effMediaType == "" {
					effMediaType = modelRes.MediaType
				}
			}
			outputRes = append(outputRes, res...)
		}
	}

	// 3.5 把本次节点消耗的 token/费用/生成视频时长写入节点执行记录，供汇总节点聚合到 exec_workflow
	// model_id 供 per_token 结算按模型聚合 token；total_duration 供 per_item/per_second 按生成视频总时长计费；
	// per_char 模型把输出字数映射到 completion_tokens 传输，随 token 拆分一并落库
	if nodeInput.NodeExecutionId > 0 && (totalTokens > 0 || totalCost > 0 || totalDuration > 0) {
		if _, err = nodeDao.NodeExecutionDao.Update(ctx, &nodeDto.UpdateNodeExecutionReq{
			Id: nodeInput.NodeExecutionId,
			TokenInfo: []map[string]any{{
				// model_id 写字符串：token_info 为 JSONB，int64 落库成 JSON 数字，读回是 float64，
				// billing 侧按 model 聚合时 (string) 断言会失败导致 per_token 永远记 0。
				// 与 ModelItem.ModelId 的 json:"modelId,string" 约定一致，字符串精确回环（雪花id>2^53 无损）。
				// effModelID 为解析后的系统模型 id（引用行），per_token 结算按此查价。
				"model_id":          gconv.String(effModelID),
				"prompt_tokens":     totalPrompt,
				"completion_tokens": totalCompletion,
				"media_type":        effMediaType,
				"total_tokens":      totalTokens,
				"total_fee":         totalCost,
				"total_duration":    totalDuration,
			}},
		}); err != nil {
			return nil, fmt.Errorf("节点：%v 写入token信息失败: %v", nodeInput.Config.Name, err)
		}
	}

	// 4.5 视频模型节点返回多个视频时，自动调用视频合成工具（concat_videos）合并为单条；
	// 已显式配置 concat_videos 后置工具时跳过，避免重复合并
	if nodeInput.Config.PostTool != media.ProcessorName && len(outputRes) > 1 && isVideoModel(ctx, nodeInput.Config.ModelConfig.ModelId) {
		outputRes, err = invokePostTool(ctx, media.ProcessorName, outputRes, map[string]any{"callback_url": "callback_url", "upload": true})
		if err != nil {
			return nil, err
		}
	} else {
		// 4. 后置工具：加工模型输出（透传原始请求参数，供后置工具读取合并配置等）
		outputRes, err = invokePostTool(ctx, nodeInput.Config.PostTool, outputRes, modelParams)
		if err != nil {
			return nil, err
		}
	}
	nodeInput.Config.OutputResult = outputRes
	return nodeInput, nil
}

// mergeInferenceBatchResults 推理模型分批结果拼接为单条输出记录：
// 各批结果按批序对同名 key 的值做字符串拼接（"拼到一个字段"），最终返回单条 {key:值} 记录。
// 非字符串值（如结构/数组字段）取最后一份，避免误拼接。
func mergeInferenceBatchResults(results [][]map[string]any) []map[string]any {
	merged := make(map[string]any)
	for _, res := range results {
		for _, record := range res {
			for key, val := range record {
				prev, has := merged[key]
				if !has {
					merged[key] = val
					continue
				}
				sPrev, pOK := prev.(string)
				sVal, vOK := val.(string)
				if pOK && vOK {
					merged[key] = sPrev + "\n" + sVal
					continue
				}
				merged[key] = val
			}
		}
	}
	return []map[string]any{merged}
}