package ru.yandex.webmaster3.storage.util.ydb;

import java.util.Map;

import lombok.Builder;
import lombok.Value;
import org.apache.commons.lang3.text.StrSubstitutor;

/**
 * @author leonidrom
 */
public class YdbUtils {
    private static final String COMMON_YQL_HEADER =
            "$ydb_endpoint = \"${YDB_ENDPOINT}\";\n" +
                    "$ydb_database = \"${YDB_DATABASE}\";\n" +
                    "$ydb_table = \"${YDB_TABLE}\";\n" +
                    "$ydb_batch_size_bytes = ${YDB_BATCH_SIZE};\n" +
                    "$ydb_oauth_token = AsTuple(\"token\", SecureParam(\"token:default_ydb\"));\n" +
                    "$ydb_batch_size_rows = ${YDB_BATCH_SIZE_ROWS};\n" +
                    "$ydb_max_retries = ${YDB_MAX_RETRIES};\n" +
                    "\n" +
                    "PRAGMA yt.QueryCacheMode = \"disable\";\n" +
                    "PRAGMA yt.DataSizePerJob = \"500485760\";\n" +
                    "PRAGMA yt.MaxJobCount = \"${MAX_JOBS_COUNT}\";\n" +
                    "PRAGMA yt.UserSlots = \"${USER_SLOTS}\";\n" +
                    "PRAGMA yt.DefaultMaxJobFails = \"1\";\n" +
                    "\n";

    private static final String PUSH_DATA_STRING =
            "USING YDB::${YDB_CLONE_OPERATION}(\n" +
                    "    TableRows(),\n" +
                    "    $ydb_endpoint,\n" +
                    "    $ydb_database,\n" +
                    "    $ydb_table,\n" +
                    "    $ydb_oauth_token,\n" +
                    "    $ydb_batch_size_bytes,\n" +
                    "    $ydb_batch_size_rows,\n" +
                    "    $ydb_max_retries\n" +
                    ");\n";

    public static String getImportQueryString(String ydbEndpoint, String ydbDatabase, String tablePath,
                                              String dataSelectStatement, QueryProperty queryProperty) {
        String importQuery = YdbUtils.COMMON_YQL_HEADER +
                "PROCESS (\n" +
                dataSelectStatement +
                ")\n" +
                YdbUtils.PUSH_DATA_STRING +
                "COMMIT;";

        StrSubstitutor substitutor = new StrSubstitutor(Map.of(
                "YDB_ENDPOINT", ydbEndpoint,
                "YDB_DATABASE", ydbDatabase,
                "YDB_TABLE", tablePath,
                "YDB_BATCH_SIZE", queryProperty.batchSize,
                "YDB_BATCH_SIZE_ROWS", queryProperty.batchSizeRows,
                "YDB_MAX_RETRIES", queryProperty.maxRetries,
                "MAX_JOBS_COUNT", queryProperty.maxJobsCount,
                "USER_SLOTS", queryProperty.userSlots,
                "YDB_CLONE_OPERATION", queryProperty.getOperationPrefix().toString()
        ));

        return substitutor.replace(importQuery);
    }

    public static String getImportQueryString(String ydbEndpoint, String ydbDatabase, String tablePath,
                                              String dataSelectStatement) {
        return getImportQueryString(ydbEndpoint, ydbDatabase, tablePath,
                dataSelectStatement, QueryProperty.builder().build());
    }

    public enum YdbCopyOperationType {
        BulkPushData,
        PushData
    }

    @Value
    @Builder
    public static class QueryProperty {
        Integer batchSize = 1048576;
        Integer batchSizeRows = 1048576;
        Integer maxRetries = 10;
        Integer maxJobsCount = 16384;
        Integer userSlots = 25;
        YdbCopyOperationType OperationPrefix = YdbCopyOperationType.PushData;
    }
}
