/* This file is part of VoltDB.
 * Copyright (C) 2008-2013 VoltDB Inc.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with VoltDB.  If not, see <http://www.gnu.org/licenses/>.
 */

package org.voltdb.sysprocs.saverestore;

import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMultimap;
import com.google.common.collect.Multimap;
import org.json_voltpatches.JSONArray;
import org.json_voltpatches.JSONException;
import org.json_voltpatches.JSONObject;
import org.json_voltpatches.JSONStringer;
import org.voltdb.catalog.Database;
import org.voltdb.catalog.Table;

import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

public class StreamSnapshotRequestConfig extends SnapshotRequestConfig {

    /**
     * Represents a stream in the snapshot. A stream has its own stream pairs,
     * may or may not specify a partition to add and the hash ranges that partition owns.
     *
     * A single snapshot can contain multiple streams.
     */
    public static class Stream {
        // src -> (dest1, dest2,...)
        public final Multimap<Long, Long> streamPairs;
        // the partition the ranges associate to
        public final Integer newPartition;

        /**
         * @param streamPairs     src - > (dest1, dest2,...)
         * @param newPartition    New partition for this stream, if not null, will create a
         *                        post-snapshot task to increment the partition count
         */
        public Stream(Multimap<Long, Long> streamPairs, Integer newPartition)
        {
            this.streamPairs = ImmutableMultimap.copyOf(streamPairs);
            this.newPartition = newPartition;
        }
    }

    // stream configs
    public final List<Stream> streams;
    // true to also do a truncation snapshot
    public final boolean shouldTruncate;

    /**
     * @param tables             See {@link #SnapshotRequestConfig(java.util.List)} for more
     * @param streams            Stream configurations
     * @param shouldTruncate     true to also generate a truncation snapshot
     */
    public StreamSnapshotRequestConfig(List<Table> tables,
                                       List<Stream> streams,
                                       boolean shouldTruncate)
    {
        super(tables);

        this.streams = ImmutableList.copyOf(streams);
        this.shouldTruncate = shouldTruncate;
    }

    public StreamSnapshotRequestConfig(JSONObject jsData,
                                       Database catalogDatabase)
    {
        super(jsData, catalogDatabase);

        this.streams = parseStreams(jsData);
        this.shouldTruncate = jsData.optBoolean("shouldTruncate", false);
    }

    private ImmutableList<Stream> parseStreams(JSONObject jsData)
    {
        ImmutableList.Builder<Stream> builder = ImmutableList.builder();

        try {
            JSONArray streamArray = jsData.getJSONArray("streams");

            for (int i = 0; i < streamArray.length(); i++) {
                JSONObject streamObj = streamArray.getJSONObject(i);

                Stream config = new Stream(parseStreamPairs(streamObj),
                                           streamObj.optInt("newPartition"));

                builder.add(config);
            }
        } catch (JSONException e) {
            SNAP_LOG.warn("Failed to parse stream snapshot request config", e);
        }

        return builder.build();
    }

    private static Multimap<Long, Long> parseStreamPairs(JSONObject jsData)
    {
        ArrayListMultimap<Long, Long> streamPairs = ArrayListMultimap.create();

        if (jsData != null) {
            try {
                JSONObject sp = jsData.getJSONObject("streamPairs");
                @SuppressWarnings("unchecked")
                Iterator<String> it = sp.keys();
                while (it.hasNext()) {
                    String key = it.next();
                    long sourceHSId = Long.valueOf(key);
                    JSONArray destJSONArray = sp.getJSONArray(key);
                    for (int i = 0; i < destJSONArray.length(); i++) {
                        long destHSId = destJSONArray.getLong(i);
                        streamPairs.put(sourceHSId, destHSId);
                    }
                }
            } catch (JSONException e) {
                SNAP_LOG.warn("Failed to parse stream pair information", e);
            }
        }

        return streamPairs;
    }

    @Override
    public void toJSONString(JSONStringer stringer) throws JSONException
    {
        super.toJSONString(stringer);

        stringer.key("shouldTruncate").value(shouldTruncate);
        stringer.key("streams").array();

        for (Stream stream : streams) {
            stringer.object();

            stringer.key("newPartition").value(stream.newPartition);

            stringer.key("streamPairs").object();
            for (Map.Entry<Long, Collection<Long>> entry : stream.streamPairs.asMap().entrySet()) {
                stringer.key(Long.toString(entry.getKey())).array();
                for (long destHSId : entry.getValue()) {
                    stringer.value(destHSId);
                }
                stringer.endArray();
            }
            stringer.endObject();

            stringer.endObject();
        }

        stringer.endArray();
    }
}

@@ -3714,63 +3714,183 @@
     
-Stream config = new Stream(parseStreamPairs(streamObj),
+Integer newPartition = null;%0A                if (!streamObj.isNull(%22newPartition%22)) %7B%0A                    newPartition = Integer.parseInt(streamObj.getString(%22newPartition%22));
 %0A   
@@ -3894,32 +3894,34 @@
                 
+%7D%0A
                 
@@ -3924,37 +3924,64 @@
     
-           streamObj.optInt(%22
+Stream config = new Stream(parseStreamPairs(streamObj), 
 newP
@@ -3988,18 +3988,16 @@
 artition
-%22)
 );%0A%0A    
@@ -5576,32 +5576,134 @@
 eam.newPartition
+ == null ?%0A                                               null : Integer.toString(stream.newPartition)
 );%0A%0A            

Fix StreamSnapshotRequestConfig serialization bug.