package ru.yandex.webmaster3.storage.util.yt;


import java.util.Arrays;
import java.util.List;
import java.util.Map;

/**
 * @author avhaliullin
 */
public class YtMapReduceCommand extends AbstractYtJob {
    private final YtJobSpec mapperSpec;
    private final YtJobSpec reducerSpec;
    private final YtJobSpec combinerSpec;
    private final List<String> sortBy;
    private final List<String> reduceBy;

    YtMapReduceCommand(String ytCluster, List<YtPath> inputTables, List<YtPath> outputTables,
                       Map<String, Object> secureVault, YtJobSpec mapperSpec, YtJobSpec reducerSpec, YtJobSpec combinerSpec,
                       List<String> sortBy, List<String> reduceBy, Map<String, Object> spec) {
        super("map_reduce", ytCluster, inputTables, outputTables, secureVault, spec);
        this.mapperSpec = mapperSpec;
        this.reducerSpec = reducerSpec;
        this.combinerSpec = combinerSpec;
        this.sortBy = sortBy;
        this.reduceBy = reduceBy;
    }

    @Override
    protected void fillSpec(Map<String, Object> spec) {
        super.fillSpec(spec);
        spec.put("mapper", mapperSpec);
        spec.put("reducer", reducerSpec);
        spec.put("reduce_combiner", combinerSpec);
        spec.put("sort_by", sortBy);
        spec.put("reduce_by", reduceBy);
    }

    public static Builder newBuilder() {
        return new Builder();
    }

    public static class Builder extends AbstractYtJobBuilder<Builder> {
        private YtJobSpec mapperSpec;
        private YtJobSpec reducerSpec;
        private YtJobSpec combinerSpec;
        private List<String> reduceBy;
        private List<String> sortBy;
        private YtReduceIOSpec reduceJobIo;
        private YtMapIOSpec mapJobIo;

        @Override
        protected Builder getSelf() {
            return this;
        }

        public Builder setMapper(YtJobSpec mapperSpec) {
            this.mapperSpec = mapperSpec;
            return this;
        }

        public Builder setReducer(YtJobSpec reducerSpec) {
            this.reducerSpec = reducerSpec;
            return this;
        }

        public Builder setCombiner(YtJobSpec reducerSpec) {
            this.combinerSpec = reducerSpec;
            return this;
        }

        public Builder reduceBy(String... columns) {
            reduceBy = Arrays.asList(columns);
            return this;
        }

        public Builder sortBy(String... columns) {
            sortBy = Arrays.asList(columns);
            return this;
        }

        public Builder setReduceJobIo(YtReduceIOSpec ioSpec) {
            this.reduceJobIo = ioSpec;
            return this;
        }

        public Builder setMapperJobIo(YtMapIOSpec ioSpec) {
            this.mapJobIo = ioSpec;
            return this;
        }

        public YtMapReduceCommand build() {
            preBuild();
            if (reduceJobIo != null) {
                spec.put("reduce_job_io", reduceJobIo);
            }
            if (mapJobIo != null) {
                spec.put("map_job_io", mapJobIo);
            }
            return new YtMapReduceCommand(ytCluster, inputTables, outputTables, secureVault, mapperSpec, reducerSpec, combinerSpec, sortBy, reduceBy, spec);
        }
    }
}
