package io.delphiplatform.api.v3.bigtable;

import com.google.cloud.bigtable.data.v2.models.Filters.ChainFilter;
import com.google.cloud.bigtable.data.v2.models.Filters.Filter;
import com.google.cloud.bigtable.data.v2.models.Filters.InterleaveFilter;
import com.google.cloud.bigtable.data.v2.models.Query;
import com.google.cloud.bigtable.data.v2.models.Row;
import com.google.common.collect.Lists;

import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;

import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import io.delphiplatform.api.util.CollectionUtils;
import io.delphiplatform.api.v3.bigtable.client.DelphiBigtableClient;
import io.delphiplatform.api.v3.bigtable.config.RequestRowKeyHandle;
import io.delphiplatform.api.v3.bigtable.config.RowKeyRequestType;
import io.delphiplatform.api.v3.bigtable.reader.BigtableColumns;
import io.delphiplatform.api.v3.constant.ApplicationConstants;
import lombok.extern.slf4j.Slf4j;

import static com.google.cloud.bigtable.data.v2.models.Filters.FILTERS;

@Slf4j
@Service
public class BigtableQueryExecutor {

    private final DelphiBigtableClient bigtableClient;
    private final Executor btQueryTaskExecutor;
    private final Integer btQueryBatchSize;

    public BigtableQueryExecutor(
        DelphiBigtableClient bigtableClient,
        @Qualifier("btQueriesExecutor") Executor btQueryTaskExecutor,
        @Value("${delphi.bigtable.query-batch-size}") Integer btQueryBatchSize
    ) {
        this.bigtableClient = bigtableClient;
        this.btQueryTaskExecutor = btQueryTaskExecutor;
        this.btQueryBatchSize = btQueryBatchSize;
    }

    @Async
    public CompletableFuture<List<Row>> executeQuery(Query query, Collection<String> ignoredColumns) {
        return executeQuerySync(query, ignoredColumns);
    }

    public CompletableFuture<List<Row>> executeQuerySync(Query query, Collection<String> ignoredColumns) {
        if (!CollectionUtils.isEmpty(ignoredColumns)) {
            query.filter(getIgnoreColumnsFilter(ignoredColumns));
        }

        return CompletableFuture.completedFuture(bigtableClient.getRows(query));
    }

    @Async
    public CompletableFuture<Stream<Row>> executePartitioned(String tableName, List<RequestRowKeyHandle> rowKeys) {
        return executePartitioned(tableName, rowKeys, Collections.emptyList());
    }

    @Async
    public CompletableFuture<Stream<Row>> executePartitioned(
        String tableName, List<RequestRowKeyHandle> rowKeys, Collection<String> ignoredColumns
    ) {
        if (CollectionUtils.isEmpty(rowKeys)) {
            return CompletableFuture.completedFuture(Stream.empty());
        }

        log.debug("Requested rowkey handles to table {} are {}", tableName, rowKeys);
        log.debug("Ignored columns are {}", ignoredColumns);

        List<List<RequestRowKeyHandle>> partitions = Lists.partition(rowKeys, btQueryBatchSize);
        Filter ignoredColumnsFilter = getIgnoreColumnsFilter(ignoredColumns);

        List<CompletableFuture<List<Row>>> futures = partitions.stream().map(partition -> {
            Query query = buildBTQuery(partition, tableName, ignoredColumns, ignoredColumnsFilter);

            return CompletableFuture.supplyAsync(() -> bigtableClient.getRows(query), btQueryTaskExecutor);
        }).collect(Collectors.toList());

        return CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
            .thenApply(ignored -> futures.stream().flatMap(f -> {
                try {
                    return f.get().stream();
                } catch (InterruptedException | ExecutionException e) {
                    throw new RuntimeException(e);
                }
            }));
    }

    private Query buildBTQuery(
        List<RequestRowKeyHandle> rowKeys, String tableName, Collection<String> ignoredColumns, Filter ignoredColumnsFilter
    ) {
        if (CollectionUtils.isEmpty(rowKeys)) {
            throw new IllegalArgumentException("rowkeys list is empty");
        }

        Query query = Query.create(tableName);
        ChainFilter filters = FILTERS.chain();

        RowKeyRequestType rowKeysRequestType = rowKeys.iterator().next().getRowKeyRequestType();

        switch (rowKeysRequestType) {
            case RANGE:
                rowKeys.forEach(rowKey -> query.range(rowKey.getRowKeyRange().getLeft(), rowKey.getRowKeyRange().getRight()));
                break;
            case LEADING_SEGMENTS:
                rowKeys.forEach(rowKey -> query.prefix(rowKey.getRowKeyLeadingSegments() + ApplicationConstants.ROW_KEY_SEPARATOR));
                break;
            case EXACT:
                rowKeys.forEach(rowKey -> query.rowKey(rowKey.getRowKeyExact()));
                break;
        }

        if (!CollectionUtils.isEmpty(ignoredColumns)) {
            filters.filter(ignoredColumnsFilter);
        }

        query.filter(filters);

        return query;
    }

    private Filter getIgnoreColumnsFilter(Collection<String> ignoredColumns) {
        InterleaveFilter requestedColumns = FILTERS.interleave();
        BigtableColumns.ALL_COLUMNS.stream()
            .filter(column -> !ignoredColumns.contains(column))
            .forEach(column -> requestedColumns.filter(FILTERS.qualifier().exactMatch(column)));

        return requestedColumns;
    }
}