package com.agent.management.engine;

import lombok.extern.slf4j.Slf4j;

import java.nio.file.Path;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 工作流执行上下文
 * 携带节点间传递的变量、工作目录和累积节点输出
 */
@Slf4j
public class WorkflowContext {

    private final Long workflowId;
    private final String runId;
    /** 扁平变量表，所有节点的输出变量合并于此（向后兼容 {{varName}} 模板渲染） */
    private final Map<String, Object> variables = new ConcurrentHashMap<>();
    /**
     * 节点级隔离输出表：nodeId -> {varName -> value}。
     * 用于实现变量作用域隔离，避免不同分支的同名输出相互覆盖。
     */
    private final Map<String, Map<String, Object>> nodeScopedOutputs = new ConcurrentHashMap<>();
    /** 初始运行输入，与节点输出隔离保存 */
    private final Map<String, Object> initialInputs;
    /** 本次运行的工作目录，所有节点共享 */
    private final Path workingDir;
    /** 累积记录每个节点的输出（按执行顺序） */
    private final List<NodeOutput> nodeOutputs = Collections.synchronizedList(new ArrayList<>());
    /**
     * 命名空间扩展槽：承载未来扩展的上下文 section（如 _git / _memory / _runtime）。
     * key 以 "_" 开头表示系统级；用户节点不直接读写，仅 executor 内部使用。
     */
    private final Map<String, Object> sections = new ConcurrentHashMap<>();
    /**
     * 节点流式事件回调；由引擎在节点执行前注入，executor 通过它实时推送 thinking / 工具增量。
     * 为 null 表示当前节点不需要流式推送（向后兼容）。
     */
    private NodeStreamSink streamSink;

    public WorkflowContext(Long workflowId, String runId, Path workingDir, Map<String, Object> initialInputs) {
        this.workflowId = workflowId;
        this.runId = runId;
        this.workingDir = workingDir;
        this.initialInputs = initialInputs == null ? Map.of() : new LinkedHashMap<>(initialInputs);
        if (initialInputs != null) {
            this.variables.putAll(initialInputs);
        }
    }

    public Long getWorkflowId() { return workflowId; }
    public String getRunId() { return runId; }
    public Map<String, Object> getVariables() { return variables; }
    public Path getWorkingDir() { return workingDir; }
    public List<NodeOutput> getNodeOutputs() { return Collections.unmodifiableList(nodeOutputs); }
    public Map<String, Object> getSections() { return sections; }
    public Object getSection(String name) { return sections.get(name); }
    public void putSection(String name, Object value) { sections.put(name, value); }
    public NodeStreamSink getStreamSink() { return streamSink; }
    public void setStreamSink(NodeStreamSink streamSink) { this.streamSink = streamSink; }

    /**
     * 获取初始运行输入（与节点输出隔离）。
     */
    public Map<String, Object> getInitialInputs() {
        return Collections.unmodifiableMap(initialInputs);
    }

    public void setVariable(String key, Object value) {
        variables.put(key, value);
    }

    public Object getVariable(String key) {
        return variables.get(key);
    }

    /**
     * 记录节点输出：同时更新扁平变量、节点级隔离输出和累积输出列表。
     * 变量覆盖检测：若 key 已存在（被上游节点写过），记录 WARN 日志，便于 debug。
     */
    public void setNodeOutput(String nodeId, Map<String, Object> output) {
        if (output != null) {
            setNodeScopedOutput(nodeId, output);
            for (Map.Entry<String, Object> e : output.entrySet()) {
                Object existing = variables.put(e.getKey(), e.getValue());
                if (existing != null) {
                    log.warn("[WorkflowContext] 变量 {} 被节点 {} 覆盖（旧值类型={}）",
                            e.getKey(), nodeId, existing.getClass().getSimpleName());
                }
            }
            nodeOutputs.add(new NodeOutput(nodeId, output));
        }
    }

    /**
     * 以节点级命名空间保存节点输出。
     */
    public void setNodeScopedOutput(String nodeId, Map<String, Object> output) {
        if (output == null || output.isEmpty()) {
            return;
        }
        nodeScopedOutputs.compute(nodeId, (k, existing) -> {
            Map<String, Object> map = existing == null ? new LinkedHashMap<>() : new LinkedHashMap<>(existing);
            map.putAll(output);
            return map;
        });
    }

    /**
     * 获取指定节点的命名空间输出。
     */
    public Map<String, Object> getNodeScopedOutput(String nodeId) {
        Map<String, Object> map = nodeScopedOutputs.get(nodeId);
        return map == null ? Map.of() : Collections.unmodifiableMap(map);
    }

    /**
     * 获取全部节点命名空间输出（只读）。
     */
    public Map<String, Map<String, Object>> getAllNodeScopedOutputs() {
        return Collections.unmodifiableMap(nodeScopedOutputs);
    }

    /**
     * 计算当前 variables 的全量快照（深拷贝），用于 debug 视图。
     */
    public Map<String, Object> snapshotVariables() {
        return new LinkedHashMap<>(variables);
    }
}