/*
    * Licensed to the Apache Software Foundation (ASF) under one
    * or more contributor license agreements.  See the NOTICE file
    * distributed with this work for additional information
    * regarding copyright ownership.  The ASF licenses this file
    * to you 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
    *
   *     http://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 org.apache.giraph.graph;
  
  import org.apache.giraph.partition.PartitionOwner;
  import org.apache.giraph.master.MasterInfo;
  import org.apache.giraph.utils.ReflectionUtils;
  import org.apache.giraph.utils.WritableUtils;
  import org.apache.giraph.worker.WorkerInfo;
  import org.apache.hadoop.io.Writable;
  
  import com.google.common.collect.Iterables;
  import com.google.common.collect.Lists;
  import com.google.common.collect.Maps;
  
  import java.io.DataInput;
  import java.io.DataOutput;
  import java.io.IOException;
  import java.util.Collection;
  import java.util.List;
  import java.util.Map;
  
  /**
   * Helper class to write descriptions of master, workers and partition owners
   */
  public class AddressesAndPartitionsWritable implements Writable {
    /** Master information */
    private MasterInfo masterInfo;
    /** List of all workers */
    private List<WorkerInfo> workerInfos;
    /** Collection of partitions */
    private Collection<PartitionOwner> partitionOwners;
  
    /**
     * Constructor when we want to serialize object
     *
     * @param masterInfo Master information
     * @param workerInfos List of all workers
     * @param partitionOwners Collection of partitions
     */
    public AddressesAndPartitionsWritable(MasterInfo masterInfo,
        List<WorkerInfo> workerInfos,
        Collection<PartitionOwner> partitionOwners) {
      this.masterInfo = masterInfo;
      this.workerInfos = workerInfos;
      this.partitionOwners = partitionOwners;
    }
  
    /** Constructor for reflection */
    public AddressesAndPartitionsWritable() {
    }
  
    /**
     * Get master information
     *
     * @return Master information
     */
    public MasterInfo getMasterInfo() {
      return masterInfo;
    }
  
    /**
     * Get all workers
     *
     * @return List of all workers
     */
    public List<WorkerInfo> getWorkerInfos() {
      return workerInfos;
    }
  
    /**
     * Get partition owners
     *
     * @return Collection of partition owners
     */
    public Collection<PartitionOwner> getPartitionOwners() {
      return partitionOwners;
    }
  
    @Override
    public void write(DataOutput output) throws IOException {
      masterInfo.write(output);
  
     output.writeInt(workerInfos.size());
     for (WorkerInfo workerInfo : workerInfos) {
       workerInfo.write(output);
     }
 
     Map<Integer, WorkerInfo> workerInfoMap = getAsWorkerInfoMap(workerInfos);
     // Also write out the previous worker information that are used
     // in the partition owners
     List<WorkerInfo> previousWorkerInfos = Lists.newArrayList();
     for (PartitionOwner partitionOwner : partitionOwners) {
       if (partitionOwner.getPreviousWorkerInfo() != null) {
         if (!workerInfoMap.containsKey(
             partitionOwner.getPreviousWorkerInfo().getTaskId())) {
           previousWorkerInfos.add(partitionOwner.getPreviousWorkerInfo());
         }
       }
     }
     output.writeInt(previousWorkerInfos.size());
     for (WorkerInfo workerInfo : previousWorkerInfos) {
       workerInfo.write(output);
     }
 
     output.writeInt(partitionOwners.size());
     if (partitionOwners.size() > 0) {
       WritableUtils.writeClass(
           partitionOwners.iterator().next().getClass(), output);
     }
     for (PartitionOwner partitionOwner : partitionOwners) {
       partitionOwner.writeWithWorkerIds(output);
     }
   }
 
   @Override
   public void readFields(DataInput input) throws IOException {
     masterInfo = new MasterInfo();
     masterInfo.readFields(input);
 
     int workerInfosSize = input.readInt();
     workerInfos = Lists.newArrayListWithCapacity(workerInfosSize);
     for (int i = 0; i < workerInfosSize; i++) {
       WorkerInfo workerInfo = new WorkerInfo();
       workerInfo.readFields(input);
       workerInfos.add(workerInfo);
     }
 
     Map<Integer, WorkerInfo> workerInfoMap = getAsWorkerInfoMap(workerInfos);
     int additionalWorkerInfos = input.readInt();
     for (int i = 0; i < additionalWorkerInfos; i++) {
       WorkerInfo workerInfo = new WorkerInfo();
       workerInfo.readFields(input);
       workerInfoMap.put(workerInfo.getTaskId(), workerInfo);
     }
 
     int partitionOwnersSize = input.readInt();
     Class<PartitionOwner> partitionOwnerClass = null;
     if (partitionOwnersSize > 0) {
       partitionOwnerClass = WritableUtils.readClass(input);
     }
     partitionOwners = Lists.newArrayListWithCapacity(partitionOwnersSize);
     for (int i = 0; i < partitionOwnersSize; i++) {
       PartitionOwner partitionOwner =
           ReflectionUtils.newInstance(partitionOwnerClass);
       partitionOwner.readFieldsWithWorkerIds(input, workerInfoMap);
       partitionOwners.add(partitionOwner);
     }
   }
 
   /**
    * Convert Iterable of WorkerInfos to the map from task id to WorkerInfo.
    *
    * @param workerInfos Iterable of WorkerInfos
    * @return The map from task id to WorkerInfo
    */
   private static Map<Integer, WorkerInfo> getAsWorkerInfoMap(
       Iterable<WorkerInfo> workerInfos) {
     Map<Integer, WorkerInfo> workerInfoMap =
         Maps.newHashMapWithExpectedSize(Iterables.size(workerInfos));
     for (WorkerInfo workerInfo : workerInfos) {
       workerInfoMap.put(workerInfo.getTaskId(), workerInfo);
     }
     return workerInfoMap;
   }
 }