/*
 * 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
 *
 *     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 com.facebook.presto.sql.planner;

import com.facebook.presto.spi.predicate.NullableValue;
import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;

import javax.annotation.concurrent.Immutable;

import java.util.Collection;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.function.Function;

import static com.google.common.base.MoreObjects.toStringHelper;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static java.util.Objects.requireNonNull;

@Immutable
public final class Partitioning
{
    private final PartitioningHandle handle;
    private final List<ArgumentBinding> arguments;

    private Partitioning(PartitioningHandle handle, List<ArgumentBinding> arguments)
    {
        this.handle = requireNonNull(handle, "handle is null");
        this.arguments = ImmutableList.copyOf(requireNonNull(arguments, "arguments is null"));
    }

    public static Partitioning create(PartitioningHandle handle, List<Symbol> columns)
    {
        return new Partitioning(handle, columns.stream()
                .map(ArgumentBinding::columnBinding)
                .collect(toImmutableList()));
    }

    // Factory method for JSON serde only!
    @JsonCreator
    public static Partitioning jsonCreate(
            @JsonProperty("handle") PartitioningHandle handle,
            @JsonProperty("arguments") List<ArgumentBinding> arguments)
    {
        return new Partitioning(handle, arguments);
    }

    @JsonProperty
    public PartitioningHandle getHandle()
    {
        return handle;
    }

    @JsonProperty
    public List<ArgumentBinding> getArguments()
    {
        return arguments;
    }

    public Set<Symbol> getColumns()
    {
        return arguments.stream()
                .filter(ArgumentBinding::isVariable)
                .map(ArgumentBinding::getColumn)
                .collect(toImmutableSet());
    }

    public boolean isPartitionedWith(Partitioning right,
            Function<Symbol, Set<Symbol>> leftToRightMappings,
            Function<Symbol, Optional<NullableValue>> leftConstantMapping,
            Function<Symbol, Optional<NullableValue>> rightConstantMapping)
    {
        if (!handle.equals(right.handle)) {
            return false;
        }

        if (arguments.size() != right.arguments.size()) {
            return false;
        }

        for (int i = 0; i < arguments.size(); i++) {
            ArgumentBinding leftArgument = arguments.get(i);
            ArgumentBinding rightArgument = right.arguments.get(i);

            if (!isPartitionedWith(leftArgument, leftConstantMapping, rightArgument, rightConstantMapping, leftToRightMappings)) {
                return false;
            }
        }
        return true;
    }

    private static boolean isPartitionedWith(
            ArgumentBinding leftArgument,
            Function<Symbol, Optional<NullableValue>> leftConstantMapping,
            ArgumentBinding rightArgument,
            Function<Symbol, Optional<NullableValue>> rightConstantMapping,
            Function<Symbol, Set<Symbol>> leftToRightMappings)
    {
        if (leftArgument.isVariable()) {
            if (rightArgument.isVariable()) {
                // variable == variable
                Set<Symbol> mappedColumns = leftToRightMappings.apply(leftArgument.getColumn());
                return mappedColumns.contains(rightArgument.getColumn());
            }
            else {
                // variable == constant
                // Normally, this would be a false condition, but if we happen to have an external
                // mapping from the symbol to a constant value and that constant value matches the
                // right value, then we are co-partitioned.
                Optional<NullableValue> leftConstant = leftConstantMapping.apply(leftArgument.getColumn());
                return leftConstant.isPresent() && leftConstant.get().equals(rightArgument.getConstant());
            }
        }
        else {
            if (rightArgument.isConstant()) {
                // constant == constant
                return leftArgument.getConstant().equals(rightArgument.getConstant());
            }
            else {
                // constant == variable
                Optional<NullableValue> rightConstant = rightConstantMapping.apply(rightArgument.getColumn());
                return rightConstant.isPresent() && rightConstant.get().equals(leftArgument.getConstant());
            }
        }
    }

    public boolean isPartitionedOn(Collection<Symbol> columns, Set<Symbol> knownConstants)
    {
        // partitioned on (k_1, k_2, ..., k_n) => partitioned on (k_1, k_2, ..., k_n, k_n+1, ...)
        // can safely ignore all constant columns when comparing partition properties
        return arguments.stream()
                .filter(ArgumentBinding::isVariable)
                .map(ArgumentBinding::getColumn)
                .filter(symbol -> !knownConstants.contains(symbol))
                .allMatch(columns::contains);
    }

    public boolean isEffectivelySinglePartition(Set<Symbol> knownConstants)
    {
        return isPartitionedOn(ImmutableSet.of(), knownConstants);
    }

    public boolean isRepartitionEffective(Collection<Symbol> keys, Set<Symbol> knownConstants)
    {
        Set<Symbol> keysWithoutConstants = keys.stream()
                .filter(symbol -> !knownConstants.contains(symbol))
                .collect(toImmutableSet());
        Set<Symbol> nonConstantArgs = arguments.stream()
                .filter(ArgumentBinding::isVariable)
                .map(ArgumentBinding::getColumn)
                .filter(symbol -> !knownConstants.contains(symbol))
                .collect(toImmutableSet());
        return !nonConstantArgs.equals(keysWithoutConstants);
    }

    public Partitioning translate(Function<Symbol, Symbol> translator)
    {
        return new Partitioning(handle, arguments.stream()
                .map(argument -> argument.translate(translator))
                .collect(toImmutableList()));
    }

    public Optional<Partitioning> translate(Function<Symbol, Optional<Symbol>> translator, Function<Symbol, Optional<NullableValue>> constants)
    {
        ImmutableList.Builder<ArgumentBinding> newArguments = ImmutableList.builder();
        for (ArgumentBinding argument : arguments) {
            Optional<ArgumentBinding> newArgument = argument.translate(translator, constants);
            if (!newArgument.isPresent()) {
                return Optional.empty();
            }
            newArguments.add(newArgument.get());
        }

        return Optional.of(new Partitioning(handle, newArguments.build()));
    }

    @Override
    public int hashCode()
    {
        return Objects.hash(handle, arguments);
    }

    @Override
    public boolean equals(Object obj)
    {
        if (this == obj) {
            return true;
        }
        if (obj == null || getClass() != obj.getClass()) {
            return false;
        }
        final Partitioning other = (Partitioning) obj;
        return Objects.equals(this.handle, other.handle) &&
                Objects.equals(this.arguments, other.arguments);
    }

    @Override
    public String toString()
    {
        return toStringHelper(this)
                .add("handle", handle)
                .add("arguments", arguments)
                .toString();
    }

    @Immutable
    public static final class ArgumentBinding
    {
        private final Symbol column;
        private final NullableValue constant;

        @JsonCreator
        public ArgumentBinding(
                @JsonProperty("column") Symbol column,
                @JsonProperty("constant") NullableValue constant)
        {
            this.column = column;
            this.constant = constant;
            checkArgument((column == null) != (constant == null), "Either column or constant must be set");
        }

        public static ArgumentBinding columnBinding(Symbol column)
        {
            return new ArgumentBinding(requireNonNull(column, "column is null"), null);
        }

        public static ArgumentBinding constantBinding(NullableValue constant)
        {
            return new ArgumentBinding(null, requireNonNull(constant, "constant is null"));
        }

        public boolean isConstant()
        {
            return constant != null;
        }

        public boolean isVariable()
        {
            return column != null;
        }

        @JsonProperty
        public Symbol getColumn()
        {
            return column;
        }

        @JsonProperty
        public NullableValue getConstant()
        {
            return constant;
        }

        public ArgumentBinding translate(Function<Symbol, Symbol> translator)
        {
            if (isConstant()) {
                return this;
            }
            return columnBinding(translator.apply(column));
        }

        public Optional<ArgumentBinding> translate(Function<Symbol, Optional<Symbol>> translator, Function<Symbol, Optional<NullableValue>> constants)
        {
            if (isConstant()) {
                return Optional.of(this);
            }

            Optional<ArgumentBinding> newColumn = translator.apply(column)
                    .map(ArgumentBinding::columnBinding);
            if (newColumn.isPresent()) {
                return newColumn;
            }

            // As a last resort, check for a constant mapping for the symbol
            // Note: this MUST be last because we want to favor the symbol representation
            // as it makes further optimizations possible.
            return constants.apply(column)
                    .map(ArgumentBinding::constantBinding);
        }

        @Override
        public String toString()
        {
            if (constant != null) {
                return constant.toString();
            }
            return "\"" + column + "\"";
        }

        @Override
        public boolean equals(Object o)
        {
            if (this == o) {
                return true;
            }
            if (o == null || getClass() != o.getClass()) {
                return false;
            }
            ArgumentBinding that = (ArgumentBinding) o;
            return Objects.equals(column, that.column) &&
                    Objects.equals(constant, that.constant);
        }

        @Override
        public int hashCode()
        {
            return Objects.hash(column, constant);
        }
    }
}
