content: /*
   * Copyright 2024-2025 the original author or authors.
   *
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *      https://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package com.alibaba.cloud.ai.graph.state;
  
  import com.alibaba.cloud.ai.graph.KeyStrategy;
  import com.alibaba.cloud.ai.graph.NodeOutput;
  import com.alibaba.cloud.ai.graph.OverAllState;
  import com.alibaba.cloud.ai.graph.RunnableConfig;
  import com.alibaba.cloud.ai.graph.checkpoint.Checkpoint;
  
  import java.util.Map;
  
  import static java.lang.String.format;
  
  public final class StateSnapshot extends NodeOutput {
  
  	private final RunnableConfig config;
  
  	public String next() {
  		return config.nextNode().orElse(null);
  	}
  
  	public RunnableConfig config() {
  		return config;
  	}
  
  	private StateSnapshot(String node, OverAllState state, RunnableConfig config) {
  		super(node, state);
  		this.config = config;
  	}
  
  	@Override
  	public String toString() {
  
  		return format("StateSnapshot{node=%s, state=%s, config=%s}", node(), state(), config());
  	}
  
  	public static StateSnapshot of(Map<String, KeyStrategy> keyStrategyMap, Checkpoint checkpoint,
  			RunnableConfig config, AgentStateFactory<OverAllState> factory) {
  
  		RunnableConfig newConfig = RunnableConfig.builder(config)
  			.checkPointId(checkpoint.getId())
  			.nextNode(checkpoint.getNextNodeId())
  			.build();
  		return new StateSnapshot(checkpoint.getNodeId(),
  				factory.apply(checkpoint.getState()).registerKeyAndStrategy(keyStrategyMap), newConfig);
  	}
  
  	public static StateSnapshot of(OverAllState overAllState, Checkpoint checkpoint, RunnableConfig config,
  			AgentStateFactory<OverAllState> factory) {
  
  		RunnableConfig newConfig = RunnableConfig.builder(config)
  			.checkPointId(checkpoint.getId())
  			.nextNode(checkpoint.getNextNodeId())
  			.build();
  		return new StateSnapshot(checkpoint.getNodeId(),
  				factory.apply(checkpoint.getState()).registerKeyAndStrategy(overAllState.keyStrategies()), newConfig);
  	}
  
  }
  