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.serializer.std;
  
  import com.alibaba.cloud.ai.graph.OverAllState;
  import com.alibaba.cloud.ai.graph.serializer.Serializer;
  import com.alibaba.cloud.ai.graph.serializer.StateSerializer;
  import com.alibaba.cloud.ai.graph.state.AgentStateFactory;
  
  import java.io.IOException;
  import java.io.ObjectInput;
  import java.io.ObjectOutput;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  
  public class ObjectStreamStateSerializer extends StateSerializer {
  
  	private static final org.slf4j.Logger log = org.slf4j.LoggerFactory.getLogger(ObjectStreamStateSerializer.class);
  
  	private final SerializerMapper mapper = new SerializerMapper();
  
  	private final MapSerializer mapSerializer = new MapSerializer();
  
  	public ObjectStreamStateSerializer(AgentStateFactory<OverAllState> stateFactory) {
  		super(stateFactory);
  		mapper.register(Collection.class, new ListSerializer());
  		mapper.register(Map.class, new MapSerializer());
  	}
  
  	public SerializerMapper mapper() {
  		return mapper;
  	}
  
  	@Override
  	public final void writeData(Map<String, Object> data, ObjectOutput out) throws IOException {
  		mapSerializer.write(data, mapper.objectOutputWithMapper(out));
  	}
  
  	@Override
  	public final Map<String, Object> readData(ObjectInput in) throws IOException, ClassNotFoundException {
  		return mapSerializer.read(mapper.objectInputWithMapper(in));
  	}
  
  	static class ListSerializer implements NullableObjectSerializer<List<Object>> {
  
  		@Override
  		public void write(List<Object> object, ObjectOutput out) throws IOException {
  			out.writeInt(object.size());
  
  			for (Object value : object) {
  				try {
  					writeNullableObject(value, out);
  				}
  				catch (IOException ex) {
  					log.error("Error writing collection value", ex);
  					throw ex;
  				}
  			}
  
  			out.flush();
  
  		}
  
  		@Override
  		public List<Object> read(ObjectInput in) throws IOException, ClassNotFoundException {
  			List<Object> result = new ArrayList<>();
  
  			int size = in.readInt();
  
  			for (int i = 0; i < size; i++) {
  
  				Object value = readNullableObject(in).orElse(null);
  
  				result.add(value);
  
  			}
  
  			return result;
  		}
  
  	}
  
  	static class MapSerializer implements NullableObjectSerializer<Map<String, Object>> {
  
  		@Override
  		public void write(Map<String, Object> object, ObjectOutput out) throws IOException {
  			out.writeInt(object.size());
  
  			for (Map.Entry<String, Object> e : object.entrySet()) {
  				try {
  					Serializer.writeUTF(e.getKey(), out);
  
  					writeNullableObject(e.getValue(), out);
  
  				}
  				catch (IOException ex) {
  					log.error("Error writing map key '{}'", e.getKey(), ex);
  					throw ex;
  				}
  			}
  
  			out.flush();
  
  		}
  
  		@Override
  		public Map<String, Object> read(ObjectInput in) throws IOException, ClassNotFoundException {
  			Map<String, Object> result = new HashMap<>();
  
  			int size = in.readInt();
  
  			for (int i = 0; i < size; i++) {
  				String key = Serializer.readUTF(in);
  
  				Object value = readNullableObject(in).orElse(null);
  
  				result.put(key, value);
  
  			}
  			return result;
  		}
  
  	}
  
  }
  