/**
 * 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.storm.daemon;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.Set;
import org.apache.storm.Config;
import org.apache.storm.Thrift;
import org.apache.storm.generated.GlobalStreamId;
import org.apache.storm.generated.Grouping;
import org.apache.storm.grouping.CustomStreamGrouping;
import org.apache.storm.grouping.LoadAwareCustomStreamGrouping;
import org.apache.storm.grouping.LoadAwareShuffleGrouping;
import org.apache.storm.grouping.LoadMapping;
import org.apache.storm.grouping.ShuffleGrouping;
import org.apache.storm.shade.com.google.common.collect.Ordering;
import org.apache.storm.shade.com.google.common.collect.Sets;
import org.apache.storm.task.WorkerTopologyContext;
import org.apache.storm.tuple.Fields;
import org.apache.storm.utils.TupleUtils;
import org.apache.storm.utils.Utils;

public class GrouperFactory {

    // A no-op grouper
    public static final LoadAwareCustomStreamGrouping DIRECT = new LoadAwareCustomStreamGrouping() {
        @Override
        public void refreshLoad(LoadMapping loadMapping) {

        }

        @Override
        public void prepare(WorkerTopologyContext context, GlobalStreamId stream, List<Integer> targetTasks) {

        }

        @Override
        public List<Integer> chooseTasks(int taskId, List<Object> values) {
            return null;
        }

    };

    public static LoadAwareCustomStreamGrouping mkGrouper(WorkerTopologyContext context, String componentId, String streamId,
                                                          Fields outFields,
                                                          Grouping thriftGrouping,
                                                          List<Integer> unsortedTargetTasks,
                                                          Map<String, Object> topoConf) {
        List<Integer> targetTasks = Ordering.natural().sortedCopy(unsortedTargetTasks);
        final boolean isNotLoadAware = (null != topoConf.get(Config.TOPOLOGY_DISABLE_LOADAWARE_MESSAGING) && (boolean) topoConf
            .get(Config.TOPOLOGY_DISABLE_LOADAWARE_MESSAGING));
        CustomStreamGrouping result = null;
        switch (Thrift.groupingType(thriftGrouping)) {
            case FIELDS:
                if (Thrift.isGlobalGrouping(thriftGrouping)) {
                    result = new GlobalGrouper();
                } else {
                    result = new FieldsGrouper(outFields, thriftGrouping);
                }
                break;
            case SHUFFLE:
                if (isNotLoadAware) {
                    result = new ShuffleGrouping();
                } else {
                    result = new LoadAwareShuffleGrouping();
                }
                break;
            case ALL:
                result = new AllGrouper();
                break;
            case LOCAL_OR_SHUFFLE:
                // Prefer local tasks as target tasks if possible
                Set<Integer> sameTasks = Sets.intersection(Sets.newHashSet(targetTasks), Sets.newHashSet(context.getThisWorkerTasks()));
                targetTasks = (sameTasks.isEmpty()) ? targetTasks : new ArrayList<>(sameTasks);
                if (isNotLoadAware) {
                    result = new ShuffleGrouping();
                } else {
                    result = new LoadAwareShuffleGrouping();
                }
                break;
            case NONE:
                result = new NoneGrouper();
                break;
            case CUSTOM_OBJECT:
                result = (CustomStreamGrouping) Thrift.instantiateJavaObject(thriftGrouping.get_custom_object());
                break;
            case CUSTOM_SERIALIZED:
                result = Utils.javaDeserialize(thriftGrouping.get_custom_serialized(), CustomStreamGrouping.class);
                break;
            case DIRECT:
                result = DIRECT;
                break;
            default:
                result = null;
                break;
        }

        if (null != result) {
            result.prepare(context, new GlobalStreamId(componentId, streamId), targetTasks);
        }

        if (result instanceof LoadAwareCustomStreamGrouping) {
            return (LoadAwareCustomStreamGrouping) result;
        } else {
            return new BasicLoadAwareCustomStreamGrouping(result);
        }
    }

    /**
     * A bridge between CustomStreamGrouping and LoadAwareCustomStreamGrouping.
     */
    public static class BasicLoadAwareCustomStreamGrouping implements LoadAwareCustomStreamGrouping {

        private final CustomStreamGrouping customStreamGrouping;

        public BasicLoadAwareCustomStreamGrouping(CustomStreamGrouping customStreamGrouping) {
            this.customStreamGrouping = customStreamGrouping;
        }

        @Override
        public void refreshLoad(LoadMapping loadMapping) {

        }

        @Override
        public void prepare(WorkerTopologyContext context, GlobalStreamId stream, List<Integer> targetTasks) {
            customStreamGrouping.prepare(context, stream, targetTasks);
        }

        @Override
        public List<Integer> chooseTasks(int taskId, List<Object> values) {
            return customStreamGrouping.chooseTasks(taskId, values);
        }
    }

    public static class FieldsGrouper implements CustomStreamGrouping {

        private Fields outFields;
        private List<List<Integer>> targetTasks;
        private Fields groupFields;
        private int numTasks;

        public FieldsGrouper(Fields outFields, Grouping thriftGrouping) {
            this.outFields = outFields;
            this.groupFields = new Fields(Thrift.fieldGrouping(thriftGrouping));

        }

        @Override
        public void prepare(WorkerTopologyContext context, GlobalStreamId stream, List<Integer> targetTasks) {
            this.targetTasks = new ArrayList<List<Integer>>();
            for (Integer targetTask : targetTasks) {
                this.targetTasks.add(Collections.singletonList(targetTask));
            }
            this.numTasks = targetTasks.size();
        }

        @Override
        public List<Integer> chooseTasks(int taskId, List<Object> values) {
            int targetTaskIndex = TupleUtils.chooseTaskIndex(outFields.select(groupFields, values), numTasks);
            return targetTasks.get(targetTaskIndex);
        }

    }

    public static class GlobalGrouper implements CustomStreamGrouping {

        private List<Integer> targetTasks;

        public GlobalGrouper() {
        }

        @Override
        public void prepare(WorkerTopologyContext context, GlobalStreamId stream, List<Integer> targetTasks) {
            this.targetTasks = targetTasks;
        }

        @Override
        public List<Integer> chooseTasks(int taskId, List<Object> values) {
            if (targetTasks.isEmpty()) {
                return null;
            }
            // It's possible for target to have multiple tasks if it reads multiple sources
            return Collections.singletonList(targetTasks.get(0));
        }
    }

    public static class NoneGrouper implements CustomStreamGrouping {

        private final Random random;
        private List<Integer> targetTasks;
        private int numTasks;

        public NoneGrouper() {
            random = new Random();
        }

        @Override
        public void prepare(WorkerTopologyContext context, GlobalStreamId stream, List<Integer> targetTasks) {
            this.targetTasks = targetTasks;
            this.numTasks = targetTasks.size();
        }

        @Override
        public List<Integer> chooseTasks(int taskId, List<Object> values) {
            int index = random.nextInt(numTasks);
            return Collections.singletonList(targetTasks.get(index));
        }
    }

    public static class AllGrouper implements CustomStreamGrouping {

        private List<Integer> targetTasks;

        @Override
        public void prepare(WorkerTopologyContext context, GlobalStreamId stream, List<Integer> targetTasks) {
            this.targetTasks = targetTasks;
        }

        @Override
        public List<Integer> chooseTasks(int taskId, List<Object> values) {
            return targetTasks;
        }
    }

}