package com.agent.management.engine;

import lombok.extern.slf4j.Slf4j;

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

/**
 * 工作流执行上下文
 *
 * <p>变量作用域采用节点命名空间隔离：每个节点的输出存于 {@code nodeScopedOutputs[nodeId]}，
 * 同名变量互不覆盖。后继节点通过 {@link NodeInputResolver} 的显式 mapping 或四级隐式回退
 * 从前驱命名空间取值，参见 docs/design/workflow-variable-scope-design.md。</p>
 */
@Slf4j
public class WorkflowContext {

    private final Long workflowId;
    private final String runId;
    /**
     * 节点级隔离输出表：nodeId -> {varName -> value}。
     * 同名变量按 nodeId 隔离，互不覆盖。
     */
    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);
    }

    public Long getWorkflowId() { return workflowId; }
    public String getRunId() { return runId; }
    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);
    }

    /**
     * 记录节点输出：写入节点级隔离输出表与累积输出列表。
     * 不再维护扁平变量表，同名输出按 nodeId 隔离，互不覆盖。
     */
    public void setNodeOutput(String nodeId, Map<String, Object> output) {
        if (output != null) {
            setNodeScopedOutput(nodeId, output);
            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);
    }

    /**
     * 全量快照（debug 视图与运行历史回放）。
     *
     * <p>格式：初始输入（扁平形式） + 各节点命名空间输出（{nodeId}__{varName} 形式）。
     * 自动穿透 envelope：从节点 envelope 的 {@code data} 字段内部扁平化，
     * 不暴露 status/message 壳层。</p>
     */
    public Map<String, Object> snapshotAllOutputs() {
        Map<String, Object> snapshot = new LinkedHashMap<>(initialInputs);
        for (Map.Entry<String, Map<String, Object>> nodeEntry : nodeScopedOutputs.entrySet()) {
            String nodeId = nodeEntry.getKey();
            Map<String, Object> dataView = unwrapDataView(nodeEntry.getValue());
            for (Map.Entry<String, Object> varEntry : dataView.entrySet()) {
                snapshot.put(nodeId + NodeWorkspaceBuilder.NODE_FIELD_SEPARATOR + varEntry.getKey(),
                        varEntry.getValue());
            }
        }
        return snapshot;
    }

    /**
     * 解包节点输出的 data 视图（envelope.data 或原始 Map）。
     */
    private static Map<String, Object> unwrapDataView(Map<String, Object> scoped) {
        if (scoped == null || scoped.isEmpty()) return Collections.emptyMap();
        if (NodeOutputEnvelope.isEnvelope(scoped)) {
            NodeOutputEnvelope env = NodeOutputEnvelope.fromObject(scoped);
            return env == null ? Collections.emptyMap() : env.getData();
        }
        return scoped;
    }
}
