package io.delphiplatform.api.v3.bigtable.client;

import com.google.api.gax.rpc.ServerStream;
import com.google.api.gax.rpc.WatchdogTimeoutException;
import com.google.cloud.bigtable.data.v2.BigtableDataClient;
import com.google.cloud.bigtable.data.v2.models.Query;
import com.google.cloud.bigtable.data.v2.models.Row;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.retry.annotation.Backoff;
import org.springframework.retry.annotation.Retryable;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

@Component
public class DelphiBigtableClient {

    private final BigtableDataClient dataClient;

    @Autowired
    public DelphiBigtableClient(BigtableDataClient dataClient) {
        this.dataClient = dataClient;
    }

    /**
     * This call is not leveraging ServerStream's "stream" nature and just loads whole rows list in place before returning.
     */
    @Retryable(value = WatchdogTimeoutException.class, backoff = @Backoff(delay = 200))
    public List<Row> getRows(Query query) {
        ServerStream<Row> rowServerStream = dataClient.readRows(query);
        return StreamSupport.stream(rowServerStream.spliterator(), false)
            .collect(Collectors.toList());
    }
}