package io.delphi.qa.auto.awsmfrw.managers.helpers.db.bigTable;

import com.google.api.gax.rpc.ServerStream;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.cloud.bigtable.data.v2.BigtableDataClient;
import com.google.cloud.bigtable.data.v2.BigtableDataSettings;
import com.google.cloud.bigtable.data.v2.models.Query;
import com.google.cloud.bigtable.data.v2.models.Row;
import com.google.cloud.bigtable.data.v2.models.RowCell;
import com.google.gson.Gson;
import io.qameta.allure.Step;
import lombok.SneakyThrows;
import org.apache.commons.lang3.tuple.Pair;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import static io.delphi.qa.auto.awsmfrw.common_baby.IsaConstants.envConfig;

public class HlBigTable {

    private final BigtableDataClient connection;
    private final Map<String, String> secretMap;
    //at this moment we have are using one instance for dev and qa envs
    String instanceId = (envConfig.env().equals("stage")) ?
            String.format(envConfig.instanceIdPattern(),"stg","hdd") :
            String.format(envConfig.instanceIdPattern(),"qa","hdd");

    public HlBigTable(Map<String, String> secretMap) {
        this.secretMap = secretMap;
        connection = getConnection();
    }

    @SneakyThrows
    protected BigtableDataClient getConnection(){
        BigtableDataClient connection = null;
        String json = new Gson().toJson(secretMap);
        GoogleCredentials credentials = GoogleCredentials.fromStream(new ByteArrayInputStream(json.getBytes()));

        // Set the BigTable data client settings
        BigtableDataSettings settings = BigtableDataSettings.newBuilder()
                .setProjectId(secretMap.get("project_id"))
                .setInstanceId(instanceId)
                .setCredentialsProvider(() -> credentials)
                .build();
        // Create the BigTable data client
        try {
            connection = BigtableDataClient.create(settings);
        } catch (IOException e) {
            throw new RuntimeException("\n-- BigTable connection exception:" + e);
        }
        return connection;
    }

    @Step("Db::BigTable: execute query")
    public Row fetchQueriedEntities(String tableId, String rowKey) {
        Row result = null;
        try {
            result = connection.readRow(tableId, rowKey);

        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return result;
    }

    @Step("Db::BigTable: execute range query")
    public List<List<RowCell>> fetchRangeQueriedEntities(String tableId, Pair<String, String> ranges) {
        Query query = Query.create(tableId);
             query.range(ranges.getLeft(), ranges.getRight());
             query.limit(15000);
        List<List<RowCell>> rows = new ArrayList<>();
        try {
            ServerStream<Row> result = connection.readRows(query);
            for (Row row : result){
                    rows.add(row.getCells());
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return rows;
    }
}