/*
 * 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 io.prestosql.operator.index;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import io.airlift.slice.Slice;
import io.prestosql.metadata.Metadata;
import io.prestosql.spi.connector.RecordCursor;
import io.prestosql.spi.connector.RecordSet;
import io.prestosql.spi.type.Type;

import java.lang.invoke.MethodHandle;
import java.util.Iterator;
import java.util.List;
import java.util.Set;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Throwables.throwIfUnchecked;
import static io.prestosql.spi.function.OperatorType.EQUAL;
import static java.lang.Boolean.TRUE;
import static java.util.Objects.requireNonNull;

/**
 * Only retains rows that have identical values for each respective fieldSet.
 */
public class FieldSetFilteringRecordSet
        implements RecordSet
{
    private final RecordSet delegate;
    private final List<Set<Field>> fieldSets;

    public FieldSetFilteringRecordSet(Metadata metadata, RecordSet delegate, List<Set<Integer>> fieldSets)
    {
        requireNonNull(metadata, "metadata is null");
        this.delegate = requireNonNull(delegate, "delegate is null");

        ImmutableList.Builder<Set<Field>> fieldSetsBuilder = ImmutableList.builder();
        List<Type> columnTypes = delegate.getColumnTypes();
        for (Set<Integer> fieldSet : requireNonNull(fieldSets, "fieldSets is null")) {
            ImmutableSet.Builder<Field> fieldSetBuilder = ImmutableSet.builder();
            for (int field : fieldSet) {
                fieldSetBuilder.add(new Field(
                        field,
                        metadata.getScalarFunctionImplementation(metadata.resolveOperator(EQUAL, ImmutableList.of(columnTypes.get(field), columnTypes.get(field)))).getMethodHandle()));
            }
            fieldSetsBuilder.add(fieldSetBuilder.build());
        }
        this.fieldSets = fieldSetsBuilder.build();
    }

    @Override
    public List<Type> getColumnTypes()
    {
        return delegate.getColumnTypes();
    }

    @Override
    public RecordCursor cursor()
    {
        return new FieldSetFilteringRecordCursor(delegate.cursor(), fieldSets);
    }

    private static class Field
    {
        private final int field;
        private final MethodHandle equalsMethodHandle;

        public Field(int field, MethodHandle equalsMethodHandle)
        {
            this.field = field;
            this.equalsMethodHandle = requireNonNull(equalsMethodHandle, "equalsMethodHandle is null");
        }

        public int getField()
        {
            return field;
        }

        public MethodHandle getEqualsMethodHandle()
        {
            return equalsMethodHandle;
        }
    }

    private static class FieldSetFilteringRecordCursor
            implements RecordCursor
    {
        private final RecordCursor delegate;
        private final List<Set<Field>> fieldSets;

        private FieldSetFilteringRecordCursor(RecordCursor delegate, List<Set<Field>> fieldSets)
        {
            this.delegate = delegate;
            this.fieldSets = fieldSets;
        }

        @Override
        public long getCompletedBytes()
        {
            return delegate.getCompletedBytes();
        }

        @Override
        public long getReadTimeNanos()
        {
            return delegate.getReadTimeNanos();
        }

        @Override
        public Type getType(int field)
        {
            return delegate.getType(field);
        }

        @Override
        public boolean advanceNextPosition()
        {
            while (delegate.advanceNextPosition()) {
                if (fieldSetsEqual(delegate, fieldSets)) {
                    return true;
                }
            }
            return false;
        }

        private static boolean fieldSetsEqual(RecordCursor cursor, List<Set<Field>> fieldSets)
        {
            for (Set<Field> fieldSet : fieldSets) {
                if (!fieldsEquals(cursor, fieldSet)) {
                    return false;
                }
            }
            return true;
        }

        private static boolean fieldsEquals(RecordCursor cursor, Set<Field> fields)
        {
            if (fields.size() < 2) {
                return true; // Nothing to compare
            }
            Iterator<Field> fieldIterator = fields.iterator();
            Field firstField = fieldIterator.next();
            while (fieldIterator.hasNext()) {
                if (!fieldEquals(cursor, firstField, fieldIterator.next())) {
                    return false;
                }
            }
            return true;
        }

        private static boolean fieldEquals(RecordCursor cursor, Field field1, Field field2)
        {
            checkArgument(cursor.getType(field1.getField()).equals(cursor.getType(field2.getField())), "Should only be comparing fields of the same type");

            if (cursor.isNull(field1.getField()) || cursor.isNull(field2.getField())) {
                return false;
            }

            Class<?> javaType = cursor.getType(field1.getField()).getJavaType();
            try {
                if (javaType == long.class) {
                    return TRUE.equals((Boolean) field1.getEqualsMethodHandle().invokeExact(cursor.getLong(field1.getField()), cursor.getLong(field2.getField())));
                }
                if (javaType == double.class) {
                    return TRUE.equals((Boolean) field1.getEqualsMethodHandle().invokeExact(cursor.getDouble(field1.getField()), cursor.getDouble(field2.getField())));
                }
                if (javaType == boolean.class) {
                    return TRUE.equals((Boolean) field1.getEqualsMethodHandle().invokeExact(cursor.getBoolean(field1.getField()), cursor.getBoolean(field2.getField())));
                }
                if (javaType == Slice.class) {
                    return TRUE.equals((Boolean) field1.getEqualsMethodHandle().invokeExact(cursor.getSlice(field1.getField()), cursor.getSlice(field2.getField())));
                }
                return TRUE.equals((Boolean) field1.getEqualsMethodHandle().invoke(cursor.getObject(field1.getField()), cursor.getObject(field2.getField())));
            }
            catch (Throwable t) {
                throwIfUnchecked(t);
                throw new RuntimeException(t);
            }
        }

        @Override
        public boolean getBoolean(int field)
        {
            return delegate.getBoolean(field);
        }

        @Override
        public long getLong(int field)
        {
            return delegate.getLong(field);
        }

        @Override
        public double getDouble(int field)
        {
            return delegate.getDouble(field);
        }

        @Override
        public Slice getSlice(int field)
        {
            return delegate.getSlice(field);
        }

        @Override
        public Object getObject(int field)
        {
            // equals is super expensive for the currently supported types: Block
            throw new UnsupportedOperationException();
        }

        @Override
        public boolean isNull(int field)
        {
            return delegate.isNull(field);
        }

        @Override
        public void close()
        {
            delegate.close();
        }
    }
}
