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

import java.io.IOException;
import java.io.InputStream;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.stream.Collector;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import com.google.common.base.Stopwatch;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import io.airlift.compress.lz4.Lz4Compressor;
import lombok.Getter;
import lombok.Setter;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.mutable.MutableInt;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.http.HeaderElement;
import org.apache.http.HeaderElementIterator;
import org.apache.http.HttpEntity;
import org.apache.http.HttpHost;
import org.apache.http.HttpRequest;
import org.apache.http.HttpResponse;
import org.apache.http.HttpStatus;
import org.apache.http.StatusLine;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.client.CredentialsProvider;
import org.apache.http.client.config.RequestConfig;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.client.methods.HttpUriRequest;
import org.apache.http.client.protocol.HttpClientContext;
import org.apache.http.client.utils.URIBuilder;
import org.apache.http.conn.ConnectionKeepAliveStrategy;
import org.apache.http.entity.InputStreamEntity;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.NoConnectionReuseStrategy;
import org.apache.http.impl.auth.BasicScheme;
import org.apache.http.impl.client.BasicAuthCache;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.message.BasicHeaderElementIterator;
import org.apache.http.protocol.HTTP;
import org.apache.http.protocol.HttpContext;
import org.apache.http.util.Args;
import org.apache.http.util.EntityUtils;
import org.joda.time.Duration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.util.UriComponentsBuilder;

import ru.yandex.webmaster3.core.WebmasterException;
import ru.yandex.webmaster3.core.data.WebmasterHostId;
import ru.yandex.webmaster3.core.http.HttpConstants;
import ru.yandex.webmaster3.core.http.RequestTrace;
import ru.yandex.webmaster3.core.http.RequestTracer;
import ru.yandex.webmaster3.core.http.WebmasterErrorResponse;
import ru.yandex.webmaster3.core.metrics.externals.AbstractExternalAPIService;
import ru.yandex.webmaster3.core.metrics.externals.ExternalAPICallTracker;
import ru.yandex.webmaster3.core.metrics.externals.ExternalDependencyMethod;
import ru.yandex.webmaster3.core.util.ByteStreamUtil;
import ru.yandex.webmaster3.core.util.CityHash102;
import ru.yandex.webmaster3.core.util.JavaMethodWitness;
import ru.yandex.webmaster3.core.util.W3CollectionUtils;
import ru.yandex.webmaster3.core.util.W3Collectors;
import ru.yandex.webmaster3.storage.host.CommonDataState;
import ru.yandex.webmaster3.storage.host.CommonDataType;
import ru.yandex.webmaster3.storage.settings.dao.CommonDataStateYDao;
import ru.yandex.webmaster3.storage.util.clickhouse2.query.Statement;
import ru.yandex.webmaster3.storage.util.ydb.exception.WebmasterYdbException;

/**
 * @author aherman
 */
public class ClickhouseServer extends AbstractExternalAPIService {
    private static final Logger log = LoggerFactory.getLogger(ClickhouseServer.class);

    private static final Duration DEFAULT_SELECT_TIMEOUT = Duration.standardSeconds(10);
    private static final Duration DEFAULT_INSERT_TIMEOUT = Duration.standardSeconds(60);

    protected static final String MDB_SHARD_PREFIX = "shard";
    protected static final int MDB_CLICKHOUSE_PORT = 8443; //constant?

    private static final String FORMAT_TAB_SEPARATED_WITH_NAMES_AND_TYPES = "TabSeparatedWithNamesAndTypes";
    private static final String FORMAT_TAB_SEPARATED = "TabSeparated";

    private static final String PARAM_DEFAULT_FORMAT = "default_format";
    private static final String PARAM_QUERY = "query";
    private static final String PARAM_READONLY = "readonly";
    private static final String PARAM_HTTP_COMPRESSION = "enable_http_compression";
    private static final String VALUE_TRUE = "1";

    private static final String VALIDATION_QUERY = "SELECT 1";

    private static final Function<Exception, ExternalAPICallTracker.Status> EXCEPTION_MAPPER = (Exception e) -> {
        if (e instanceof ClickhouseUserException) {
            return ExternalAPICallTracker.Status.USER_ERROR;
        } else if (e instanceof IOException) {
            return ExternalAPICallTracker.Status.CONNECTION_ERROR;
        } else {
            return ExternalAPICallTracker.Status.INTERNAL_ERROR;
        }
    };

    @Getter @Setter
    private String clusterId;
    @Getter @Setter
    protected String clickhouseHosts;
    @Getter @Setter
    private String localDCName;
    @Getter @Setter
    private UsernamePasswordCredentials normalUserCredentials;
    @Getter @Setter
    private UsernamePasswordCredentials serviceUserCredentials;
    // TODO временное решение на переходный период, не будем читать/писать на новые шарды
    private int shardsCount;
    @Getter @Setter
    private DefaultUser defaultUser;

    @Autowired
    protected CommonDataStateYDao commonDataStateYDao;


    private UsernamePasswordCredentials defaultCredentials;
    private List<ClickhouseHost> hosts = Collections.emptyList();
    private CloseableHttpClient selectHttpClient;
    private CloseableHttpClient insertHttpClient;

    private ScheduledExecutorService updateNodeStateService;

    private Duration selectTimeout = DEFAULT_SELECT_TIMEOUT;
    private Duration insertTimeout = DEFAULT_INSERT_TIMEOUT;
    private ThreadLocal<ClickhouseHost> tx = new ThreadLocal<>();
    private ThreadLocal<UsernamePasswordCredentials> currentCredentials = new ThreadLocal<>();

    public void init() throws Exception {
        Stopwatch stopwatch = Stopwatch.createStarted();

        log.info("Default user set to {}", defaultUser);
        switch (defaultUser) {
            case OFFLINE:
                defaultCredentials = serviceUserCredentials;
                break;
            case USER:
                defaultCredentials = normalUserCredentials;
                break;
            default:
                throw new RuntimeException("Unknown user " + defaultUser);
        }

        Pair<List<ClickhouseHost>, List<ClickhouseHost>> allhosts = computeClickhouseHosts();
        hosts = allhosts.getLeft();

        if (hosts.isEmpty()) {
            throw new RuntimeException("Clickhouse hosts is empty");
        }

        RequestConfig requestConfig = RequestConfig
                .custom()
                .setConnectTimeout(HttpConstants.DEFAULT_CONNECT_TIMEOUT)
                .setConnectionRequestTimeout(200)
                .setContentCompressionEnabled(true)
                .setSocketTimeout((int) selectTimeout.getMillis())
                .build();

        insertHttpClient = HttpClients
                .custom()
                .setDefaultRequestConfig(requestConfig)
                .setMaxConnTotal(hosts.size() * 20)
                .setMaxConnPerRoute(20)
                .setConnectionTimeToLive(2, TimeUnit.SECONDS)
                .setConnectionReuseStrategy(NoConnectionReuseStrategy.INSTANCE)
                .disableAutomaticRetries()
                .evictExpiredConnections()
                .build();

        selectHttpClient = HttpClients
                .custom()
                .setDefaultRequestConfig(requestConfig)
                .setMaxConnTotal(hosts.size() * 20)
                .setMaxConnPerRoute(20)
                .setConnectionTimeToLive(2, TimeUnit.SECONDS)
                .setKeepAliveStrategy(new OverrideKeepAliveStrategy(2, TimeUnit.SECONDS))
                .disableAutomaticRetries()
                .evictExpiredConnections()
                .build();

        checkHosts(true);

        Map<String, Pair<MutableInt, MutableInt>> hostInfo = new HashMap<>();
        for (ClickhouseHost host : hosts) {
            Pair<MutableInt, MutableInt> upDownHostCount =
                    hostInfo.computeIfAbsent(host.getDcName(), dc -> Pair.of(new MutableInt(), new MutableInt()));
            if (host.isDown()) {
                upDownHostCount.getRight().increment();
            } else {
                upDownHostCount.getLeft().increment();
            }
        }
        shardsCount = (int) hosts.stream().map(ClickhouseHost::getShard).distinct().count();

        updateNodeStateService = Executors.newSingleThreadScheduledExecutor(
                new ThreadFactoryBuilder()
                        .setDaemon(true)
                        .setNameFormat("clickhouse-state-%d")
                        .build()
        );

        ScheduledFuture<?> future =
                updateNodeStateService.scheduleAtFixedRate(new NodeUpdater(), 5, 5, TimeUnit.SECONDS);
        log.info("Node updater started: {} {}", future.isCancelled(), future.isDone());

        log.info("Clickhouse initialized: {} in {} ms",
                hostInfo.entrySet()
                        .stream()
                        .map(e -> e.getKey() + "(" + e.getValue().getLeft() + "/" + e.getValue().getRight() + ")")
                        .collect(Collectors.joining(" ")),
                stopwatch.elapsed(TimeUnit.MILLISECONDS)
        );
    }

    protected Pair<List<ClickhouseHost>, List<ClickhouseHost>> computeClickhouseHosts() throws Exception {
        Pair<List<ClickhouseHost>, List<ClickhouseHost>> allhosts;
        if (!StringUtils.isEmpty(clickhouseHosts)) {
            return computeShards(
                    Arrays.stream(clickhouseHosts.split(","))
                            .map(ClickhouseServer::parseClickhouseHost)
                            .filter(x -> x != null)
            );
        } else {
            throw new RuntimeException("Clickhouse hosts not set");
        }
    }

    public void destroy() {
        updateNodeStateService.shutdownNow();
        try {
            updateNodeStateService.awaitTermination(5, TimeUnit.SECONDS);
        } catch (InterruptedException e) {
            // Ignore
        }
        IOUtils.closeQuietly(selectHttpClient);
        IOUtils.closeQuietly(insertHttpClient);
    }

    private void checkHosts(boolean verbose) {
        for (ClickhouseHost host : hosts) {
            if (verbose) {
                log.debug("Check host: {} {}", host.getDcName(), host.getHostURI());
            }
            long start = System.nanoTime();
            try {
                ClickhouseQueryContext.Builder chContextB = ClickhouseQueryContext.useDefaults()
                        .setHost(host)
                        .setTimeout(Duration.standardSeconds(1))
                        .setLogQuery(verbose);

                execute(chContextB, QueryType.SELECT, VALIDATION_QUERY, Optional.empty(), Optional.of(r -> {
                    if (verbose || host.isDown()) {
                        long delta = System.nanoTime() - start;
                        log.debug("Host: {} {} is up {}",
                                host.getDcName(), host.getHostURI(),
                                TimeUnit.NANOSECONDS.toMillis(delta));
                    }
                    host.markAsUp();
                    return null;
                }));
            } catch (ClickhouseException e) {
                if (verbose || !host.isDown()) {
                    String causeMessage = null;
                    Throwable cause = e.getCause();
                    if (cause != null) {
                        causeMessage = cause.getMessage();
                    }
                    long delta = System.nanoTime() - start;
                    log.debug("Host: {} {} is down {} {} {}",
                            host.getDcName(), host.getHostURI(),
                            e.getMessage(), causeMessage,
                            TimeUnit.NANOSECONDS.toMillis(delta));
                }
                host.markAsDown();
            }
        }
    }

    protected Pair<List<ClickhouseHost>, List<ClickhouseHost>> computeShards(Stream<ClickhouseHost> hosts) {
        String disableDcName = null;
        CommonDataState dataState = null;
        try {
            dataState = commonDataStateYDao.getValue(CommonDataType.DISABLED_CLICKHOUSE_DC_NAME);
        } catch (WebmasterYdbException e) {
            // ignore exception
            log.error("Unable to get disable DC name from cassandra", e);
        }
        if (dataState != null) {
            disableDcName = dataState.getValue();
        }

        List<ClickhouseHost> enableHosts = new ArrayList<>();
        List<ClickhouseHost> disableHosts = new ArrayList<>();
        Map<String, List<ClickhouseHost>> dc2Hosts = hosts.collect(Collectors.groupingBy(ClickhouseHostLocation::getDcName));
        for (String dcName : dc2Hosts.keySet()) {
            List<ClickhouseHost> hostsInDC = dc2Hosts.get(dcName);
            hostsInDC.sort(Comparator.comparing(ClickhouseHostLocation::getShard));
            for (int shard = 0; shard < hostsInDC.size(); shard++) {
                if (disableDcName != null && disableDcName.equals(dcName)) {
                    disableHosts.add(hostsInDC.get(shard));
                } else {
                    enableHosts.add(hostsInDC.get(shard));
                }
            }
        }
        return Pair.of(enableHosts, disableHosts);
    }

    protected static ClickhouseHost parseClickhouseHost(String uriSpec) {
        String[] parts = uriSpec.split(":", 3);
        String dc = parts[0];
        String host = parts[1];
        String shardStr = parts[2];
        int shard;
        if (shardStr.startsWith(MDB_SHARD_PREFIX)) {
            shard = Integer.parseInt(shardStr.substring(MDB_SHARD_PREFIX.length())) - 1;
        } else {
            shard = Integer.parseInt(shardStr) - 1;
        }
        URIBuilder uriBuilder = new URIBuilder();
        uriBuilder.setScheme(WebmasterHostId.Schema.HTTPS.getSchemaName());
        uriBuilder.setHost(host);
        uriBuilder.setPort(MDB_CLICKHOUSE_PORT);
        return new ClickhouseHost(URI.create(uriBuilder.toString()), dc.toUpperCase(), shard);
    }

    public <T> Optional<T> selectOne(String query, Function<CHRow, T> mapper) {
        return queryOne(query, mapper);
    }

    public <T> Optional<T> queryOne(String query, Function<CHRow, T> mapper) {
        return queryAll(query, mapper).stream().findFirst();
    }

    public <T> List<T> selectAll(String query, Function<CHRow, T> mapper) {
        return queryAll(query, mapper);
    }

    public <T> List<T> queryAll(String query, Function<CHRow, T> mapper) {
        List<T> result = new ArrayList<>();
        select(query, row -> {
            result.add(mapper.apply(row));
            return null;
        });
        return result;
    }

    public <T> Optional<T> queryOne(ClickhouseQueryContext.Builder chContext, Statement query, Function<CHRow, T> mapper) {
        return queryOne(chContext, query.toString(), mapper);
    }

    public <T> Optional<T> queryOne(ClickhouseQueryContext.Builder chContext, String query, Function<CHRow, T> mapper) {
        List<T> res = collectAll(chContext, query, Collectors.mapping(mapper, Collectors.toList()));
        return res.isEmpty() ? Optional.empty() : Optional.of(res.get(0));
    }

    public <T> List<T> queryAll(ClickhouseQueryContext.Builder chContext, Statement query, Function<CHRow, T> mapper) {
        return queryAll(chContext, query.toString(), mapper);
    }

    public <T> List<T> queryAll(ClickhouseQueryContext.Builder chContext, String query, Function<CHRow, T> mapper) {
        return collectAll(chContext, query, Collectors.mapping(mapper, Collectors.toList()));
    }

    public <T, A> T collectAll(ClickhouseQueryContext.Builder chContext, String query, Collector<CHRow, A, T> collector) {
        A acc = collector.supplier().get();
        execute(chContext, QueryType.SELECT, query, Optional.empty(), Optional.of(row -> {
            collector.accumulator().accept(acc, row);
            return null;
        }));
        return collector.finisher().apply(acc);
    }

    public <T, A> T collectAll(String query, InputStream data, Collector<CHRow, A, T> collector) {
        A acc = collector.supplier().get();
        select(query, data, row -> {
            collector.accumulator().accept(acc, row);
            return null;
        });
        return collector.finisher().apply(acc);
    }

    public <T, A> T collectAll(String query, Collector<CHRow, A, T> collector) {
        A acc = collector.supplier().get();
        select(query, row -> {
            collector.accumulator().accept(acc, row);
            return null;
        });
        return collector.finisher().apply(acc);
    }

    public void forEach(String query, InputStream data, Consumer<CHRow> consumer) {
        select(query, data, row -> {
            consumer.accept(row);
            return null;
        });
    }

    public void forEach(String query, Consumer<CHRow> consumer) {
        select(query, row -> {
            consumer.accept(row);
            return null;
        });
    }

    public void insert(String query) {
        insert(ClickhouseQueryContext.useDefaults(), query);
    }

    public void insert(ClickhouseQueryContext.Builder chContext, String query) {
        execute(chContext, QueryType.INSERT, query, Optional.empty(), Optional.empty());
    }

    public void insert(String query, InputStream data) {
        insert(ClickhouseQueryContext.useDefaults(), query, data);
    }

    public void insert(ClickhouseQueryContext.Builder chContext, String query, InputStream data) {
        execute(chContext, QueryType.INSERT, query, Optional.of(data), Optional.empty());
    }

    private static class CompressedHeader {
        private static final int CHECKSUM_SIZE = 16;
        private static final int METHOD_SIZE = 1;
        private static final int UNCOMPRESSED_SIZE = 4;
        private static final int COMPRESSED_SIZE = 4;

        public static void writeChecksumm(byte[] buffer, long cs0, long cs1) {
            ByteStreamUtil.writeLongLE(buffer, 0, cs0);
            ByteStreamUtil.writeLongLE(buffer, 8, cs1);
        }

        public static void writeCompressionMethod(byte[] buffer) {
            buffer[CHECKSUM_SIZE] = (byte) 0x82;
        }

        public static void writeCompressedSize(byte[] buffer, int size) {
            ByteStreamUtil.writeIntLE(buffer, CHECKSUM_SIZE + METHOD_SIZE, size);
        }

        public static void writeUncompressedSize(byte[] buffer, int size) {
            ByteStreamUtil.writeIntLE(buffer, CHECKSUM_SIZE + METHOD_SIZE + COMPRESSED_SIZE, size);
        }

        public static int headerSize() {
            return METHOD_SIZE + UNCOMPRESSED_SIZE + COMPRESSED_SIZE;
        }

        public static int fullHeaderSize() {
            return CHECKSUM_SIZE + headerSize();
        }
    }

    public <T> T executeWithFixedHost(TxCallback<T> callback) {
        ClickhouseHost host = tx.get();
        if (host != null) {
            return callback.execute();
        } else {
            host = pickRandomAliveHost();
            return executeWithFixedHost(host, callback);
        }
    }

    public <T> T executeWithFixedHost(ClickhouseHost host, TxCallback<T> callback) {
        tx.set(host);
        try {
            return callback.execute();
        } finally {
            tx.remove();
        }
    }

    public <T> T executeByServiceUser(TxCallback<T> callback) {
        currentCredentials.set(serviceUserCredentials);
        try {
            return callback.execute();
        } finally {
            currentCredentials.remove();
        }
    }

    private void select(String query, Function<CHRow, Void> consumer) {
        execute(ClickhouseQueryContext.useDefaults(), QueryType.SELECT, query, Optional.empty(), Optional.of(consumer));
    }

    private void select(String query, InputStream data, Function<CHRow, Void> consumer) {
        execute(ClickhouseQueryContext.useDefaults(), QueryType.SELECT, query, Optional.of(data), Optional.of(consumer));
    }

    private boolean isLocalHost(ClickhouseHost h) {
        return localDCName.equals(h.getDcName());
    }

    public static Pair<byte[], Integer> compressData(byte[] data) {
        Lz4Compressor lz4Compressor = new Lz4Compressor();
        byte[] compressedData = new byte[CompressedHeader.fullHeaderSize() + lz4Compressor.maxCompressedLength(data.length)];
        int compressedSize =
                lz4Compressor.compress(data, 0, data.length, compressedData, CompressedHeader.fullHeaderSize(), compressedData.length)
                        + CompressedHeader.headerSize();

        CompressedHeader.writeCompressionMethod(compressedData);
        CompressedHeader.writeCompressedSize(compressedData, compressedSize);
        CompressedHeader.writeUncompressedSize(compressedData, data.length);

        long[] hash128 = CityHash102.cityHash128(compressedData, CompressedHeader.CHECKSUM_SIZE, compressedSize);
        CompressedHeader.writeChecksumm(compressedData, hash128[0], hash128[1]);
        return Pair.of(compressedData, compressedSize + CompressedHeader.CHECKSUM_SIZE);
    }

    public void execute(String query) {
        execute(ClickhouseQueryContext.useDefaults(), QueryType.INSERT, query, Optional.empty(), Optional.empty());
    }

    public void execute(ClickhouseHost chHost, String query) {
        execute(ClickhouseQueryContext.useDefaults().setHost(chHost), QueryType.INSERT, query, Optional.empty(), Optional.empty());
    }

    public void execute(ClickhouseQueryContext.Builder chContextB, String query) {
        execute(chContextB, QueryType.INSERT, query, Optional.empty(), Optional.empty());
    }

    public void execute(ClickhouseQueryContext.Builder chContextB, QueryType queryType, String query,
                        Optional<InputStream> insertDataO,
                        Optional<Function<CHRow, Void>> responseConsumerO)
            {
        if (!chContextB.getTimeout().isPresent()) {
            if (queryType == QueryType.SELECT) {
                chContextB = chContextB.setTimeout(selectTimeout);
            } else if (queryType == QueryType.INSERT) {
                chContextB = chContextB.setTimeout(insertTimeout);
            }
        }

        if (!chContextB.getHost().isPresent()) {
            ClickhouseHost host = tx.get();
            if (host == null) {
                host = pickRandomAliveHost();
            }
            chContextB = chContextB.setHost(host);
        }

        if (!chContextB.getCredentials().isPresent()) {
            UsernamePasswordCredentials credentials = currentCredentials.get();
            if (credentials == null) {
                credentials = defaultCredentials;
            }
            chContextB = chContextB.setCredentials(credentials);
        }

        ClickhouseQueryContext chContext = chContextB.build();
        Stopwatch qWatch = Stopwatch.createStarted();
        try {
            execute(chContext, queryType, query, insertDataO, responseConsumerO);
        } finally {
            if (chContext.isLogQuery()) {
                log.debug("Q: {} {}-{} {}", qWatch.elapsed(TimeUnit.MILLISECONDS), chContext.getHost().getDcName(), chContext.getHost().getShard(), query);

                if (queryType == QueryType.SELECT) {
                    RequestTracer.getCurrentTrace()
                            .stopRead(RequestTrace.Db.CLICKHOUSE, query, qWatch.elapsed(TimeUnit.NANOSECONDS));
                } else if (queryType == QueryType.INSERT) {
                    RequestTracer.getCurrentTrace()
                            .stopWrite(RequestTrace.Db.CLICKHOUSE, query, qWatch.elapsed(TimeUnit.NANOSECONDS));
                }
            }
        }
    }

    @ExternalDependencyMethod("execute")
    private void execute(ClickhouseQueryContext chContext,
                         QueryType queryType, String query, Optional<InputStream> insertDataO,
                         Optional<Function<CHRow, Void>> responseConsumerO) {

        ClickhouseHost host = chContext.getHost();
        URI serverURI = host.getHostURI();
        UriComponentsBuilder uriComponentsBuilder = UriComponentsBuilder
                .fromUri(serverURI).queryParam(PARAM_HTTP_COMPRESSION, VALUE_TRUE);

        if (queryType == QueryType.SELECT) {
            uriComponentsBuilder = uriComponentsBuilder
                    .queryParam(PARAM_DEFAULT_FORMAT, FORMAT_TAB_SEPARATED_WITH_NAMES_AND_TYPES);
        } else if (insertDataO.isPresent()) {
            uriComponentsBuilder = uriComponentsBuilder
                    .queryParam(PARAM_DEFAULT_FORMAT, FORMAT_TAB_SEPARATED);
        }
        for (Map.Entry<String, String> entry : chContext.getSettings().entrySet()) {
            uriComponentsBuilder = uriComponentsBuilder.queryParam(entry.getKey(), entry.getValue());
        }

        URI uri;

        HttpUriRequest request;
        if (queryType == QueryType.SELECT) {
            uriComponentsBuilder = uriComponentsBuilder.queryParam(PARAM_READONLY, VALUE_TRUE);
        }
        if (insertDataO.isPresent()) {
            uri = uriComponentsBuilder
                    .queryParam(PARAM_QUERY, query)
                    .build().toUri();
            HttpPost post = new HttpPost(uri);
            post.setEntity(new InputStreamEntity(insertDataO.get()));
            request = post;
        } else {
            uri = uriComponentsBuilder.build().toUri();
            HttpPost post = new HttpPost(uri);
            post.setEntity(new StringEntity(query, StandardCharsets.UTF_8));
            request = post;
        }

        HttpClientContext context = createContext(request, chContext, uri);
        RequestConfig.Builder rcBuilder = RequestConfig.copy(context.getRequestConfig())
                .setSocketTimeout((int) chContext.getTimeout().getMillis());

        context.setRequestConfig(rcBuilder.build());

        final CloseableHttpClient httpClient = queryType == QueryType.SELECT? selectHttpClient : insertHttpClient;
        trackExecution(new JavaMethodWitness() {}, EXCEPTION_MAPPER, () -> {
            try (CloseableHttpResponse httpResponse = httpClient.execute(request, context)) {
                StatusLine statusLine = httpResponse.getStatusLine();
                HttpEntity entity = httpResponse.getEntity();
                int statusCode = statusLine.getStatusCode();
                if (statusCode != HttpStatus.SC_OK) {
                    String error = EntityUtils.toString(entity);
                    log.error("Error executing query {}: {}", query, error);
                    String message = "Error from " + host.getHostURI() + ": " + statusLine;
                    if (statusCode == HttpStatus.SC_NOT_FOUND) {
                        throw new ClickhouseUserException(message, query, error);
                    } else {
                        throw new ClickhouseException(message, query, error);
                    }
                }

                if (queryType == QueryType.INSERT) {
                    EntityUtils.consume(entity);
                } else if (!responseConsumerO.isPresent()) {
                    EntityUtils.consume(entity);
                } else {
                    Function<CHRow, Void> consumer = responseConsumerO.get();
                    InputStream content = entity.getContent();
                    LineInputStream lineInputStream = new LineInputStream(content, 10);
                    Iterator<LineInputStream.Line> lineIterator = lineInputStream.iterator();
                    String[] columnNames = extractColumnNames(lineIterator, query);
                    CHType[] columnTypes = extractTypes(lineIterator, query);
                    if (columnNames.length != columnTypes.length) {
                        throw new ClickhouseException("Illegal response format, column and type count mismatch", query, null);
                    }

                    // Если в query содержится WITH TOTALS, то при парсинге response нужно пропустить пустую строку
                    if (query.contains("WITH TOTALS")) {
                        if (lineIterator.hasNext()) {
                            lineIterator.next();
                        }
                    }
                    while (lineIterator.hasNext()) {
                        LineInputStream.Line line = lineIterator.next();
                        consumer.apply(CHRow.createFromLine(line, columnNames, columnTypes));
                    }
                }
            } catch (IOException e) {
                //host.markAsDown();
                log.debug("Host: {} {} is down {} {}", host.getDcName(), host.getHostURI(), e.getMessage(), e.getMessage());
                throw new ClickhouseException("Unable to execute query on " + serverURI.getHost(), query, null, e);
            }
        });
    }

    private List<String> findAliveDCs() {
        return hosts.stream()
                .collect(Collectors.groupingBy(ClickhouseHost::getDcName))
                .entrySet().stream()
                .filter(entry -> entry.getValue().stream().noneMatch(ClickhouseHost::isDown))
                .map(Map.Entry::getKey)
                .collect(Collectors.toList());
    }

    private ClickhouseHost pickRandomHost(List<ClickhouseHost> aliveHosts) {
        // prefer localDc
        List<ClickhouseHost> aliveLocalHosts = aliveHosts.stream().filter(this::isLocalHost).collect(Collectors.toList());
        if (!aliveLocalHosts.isEmpty()) {
            return W3CollectionUtils.getRandomItem(aliveLocalHosts);
        }
        return W3CollectionUtils.getRandomItem(aliveHosts);
    }

    public ClickhouseHost pickRandomAliveHost() {
        List<ClickhouseHost> aliveHosts = getHosts().stream().filter(ClickhouseHost::isUp).collect(Collectors.toList());
        if (aliveHosts.isEmpty()) {
            throw new ClickhouseException("Not found alive host ", null, null);
        }
        return pickRandomHost(aliveHosts);
    }

    public ClickhouseHost pickAliveHostOrFail(int shard) {
        List<ClickhouseHost> aliveHosts = hosts.stream().filter(h -> !h.isDown()).filter(h -> h.getShard() == shard).collect(Collectors.toList());
        if (aliveHosts.isEmpty()) {
            throw new ClickhouseException("Not found alive host for shard " + shard, null, null);
        }
        return pickRandomHost(aliveHosts);
    }

    public ClickhouseHost pickAliveHostWithMaxShardOrFail(int shard) {
        List<ClickhouseHost> aliveHosts = hosts.stream().filter(h -> !h.isDown()).filter(h -> h.getShard() < shard).collect(Collectors.toList());
        if (aliveHosts.isEmpty()) {
            throw new ClickhouseException("Not found alive host for max shard " + shard, null, null);
        }
        return pickRandomHost(aliveHosts);
    }

    public Optional<String> pickAliveDC() {
        List<String> aliveDCs = findAliveDCs();
        if (aliveDCs.isEmpty()) {
            return Optional.empty();
        }
        return Optional.of(W3CollectionUtils.getRandomItem(aliveDCs));
    }

    public Collection<ClickhouseHost> getAliveHostsForEachShard() {
        List<ClickhouseHost> hosts = getHosts().stream().filter(ClickhouseHost::isUp).collect(Collectors.toList());
        Collections.shuffle(hosts);
        // оставим по одному шарду
        Collection<ClickhouseHost> shards = hosts.stream()
                .collect(Collectors.toMap(ClickhouseHost::getShard, Function.identity(), W3Collectors.acceptAnyMerger())).values();
        if (shards.size() < getShardsCount()) {
            throw new WebmasterException("Not all shards are available",
                    new WebmasterErrorResponse.ClickhouseErrorResponse(getClass(), "Not all shards are available"));
        }
        return shards;
    }

    public void executeOnAllHosts(String query) {
        executeOnAllHosts(ClickhouseQueryContext.useDefaults(), query, Integer.MAX_VALUE);
    }

    public void executeOnAllHosts(ClickhouseQueryContext.Builder chContext, String query) {
        executeOnAllHosts(chContext, query, Integer.MAX_VALUE);
    }

    public void executeOnAllHosts(ClickhouseQueryContext.Builder chContext, String query, int maxShard) {
        if (hosts.stream().anyMatch(ClickhouseHost::isDown)) {
            throw new ClickhouseException("Some hosts are down", query, "Some hosts are down, cannot perform query on all hosts");
        }
        for (ClickhouseHost host : hosts) {
            if (host.getShard() < maxShard) {
                execute(chContext.setHost(host), QueryType.INSERT, query, Optional.empty(), Optional.empty());
            }
        }
    }

    public List<ClickhouseHost> getHosts() {
        return Collections.unmodifiableList(hosts);
    }

    public int getShardsCount() {
        return shardsCount;
    }

    private String[] extractColumnNames(Iterator<LineInputStream.Line> lineIterator, String query)
            {
        if (!lineIterator.hasNext()) {
            throw new ClickhouseException("Illegal response format, missing names", query, null);
        }
        LineInputStream.Line columnsLine = lineIterator.next();
        int columnsCount = columnsLine.count((byte) '\t') + 1;

        if (columnsCount == 1) {
            return new String[]{columnsLine.getString(0, StandardCharsets.UTF_8)};
        }
        String[] columnNames = new String[columnsCount];
        int start = 0;
        int end;
        for (int i = 0; i < columnsCount - 1; i++) {
            end = columnsLine.indexOf((byte) '\t', start);
            columnNames[i] = columnsLine.getString(start, end, StandardCharsets.UTF_8);
            start = end + 1;
        }
        columnNames[columnsCount - 1] = columnsLine.getString(start, StandardCharsets.UTF_8);
        return columnNames;
    }

    private CHType[] extractTypes(Iterator<LineInputStream.Line> lineIterator, String query)
            {
        if (!lineIterator.hasNext()) {
            throw new ClickhouseException("Illegal response format, missing types", query, null);
        }
        LineInputStream.Line line = lineIterator.next();
        int count = line.count((byte) '\t') + 1;
        if (count == 1) {
            return new CHType[]{CHType.getType(line.getString(0, StandardCharsets.UTF_8))};
        }

        CHType[] types = new CHType[count];
        int start = 0;
        int end;
        for (int i = 0; i < count - 1; i++) {
            end = line.indexOf((byte) '\t', start);
            types[i] = CHType.getType(line.getString(start, end, StandardCharsets.UTF_8));
            start = end + 1;
        }
        types[count - 1] = CHType.getType(line.getString(start, StandardCharsets.UTF_8));
        return types;
    }

    protected HttpClientContext createContext(HttpRequest request, ClickhouseQueryContext chContext, URI uri) {
        CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
        BasicAuthCache authCache = new BasicAuthCache();

        HttpHost httpHost = new HttpHost(uri.getHost(), uri.getPort());
        credentialsProvider.setCredentials(new AuthScope(httpHost.getHostName(), httpHost.getPort()), chContext.getCredentials());
        authCache.put(httpHost, new BasicScheme());

        HttpClientContext clientContext = HttpClientContext.create();
        clientContext.setCredentialsProvider(credentialsProvider);
        clientContext.setAuthCache(authCache);

        return clientContext;
    }

    public static Function<List<ClickhouseHost>, List<ClickhouseHost>> alive() {
        return (hostList) -> hostList.stream().filter(h -> !h.isDown()).collect(Collectors.toList());
    }

    public static Function<List<ClickhouseHost>, List<ClickhouseHost>> all() {
        return (hostList) -> hostList.stream().filter(h -> !h.isDown()).collect(Collectors.toList());
    }

    public Function<List<ClickhouseHost>, List<ClickhouseHost>> onlyLocalDc() {
        return (hostsList) -> {
            return new ArrayList<>(hostsList.stream()
                    .filter(h -> localDCName.equals(h.getDcName()))
                    .collect(Collectors.toList()));
        };
    }

    public Function<List<ClickhouseHost>, List<ClickhouseHost>> localDcFirst() {
        return (hostsList) -> {
            List<ClickhouseHost> result = new ArrayList<>(hostsList);
            result.sort((g1, g2) -> {
                if (localDCName.equals(g1.getDcName())) {
                    if (localDCName.equals(g2.getDcName())) {
                        return 0;
                    }
                    return -1;
                } else if (localDCName.equals(g2.getDcName())) {
                    return 1;
                }
                return 0;
            });
            return result;
        };
    }

    public static Function<List<ClickhouseHost>, List<ClickhouseHost>> randomize() {
        return (hostsList) -> {
            List<ClickhouseHost> result = new ArrayList<>(hostsList);
            Collections.shuffle(result);
            return result;
        };
    }

    public enum QueryType {
        INSERT,
        SELECT
    }

    private static class OverrideKeepAliveStrategy implements ConnectionKeepAliveStrategy {
        private final long keepAliveOverride;

        private OverrideKeepAliveStrategy(long keepAliveOverride, TimeUnit timeUnit) {
            this.keepAliveOverride = timeUnit.toMillis(keepAliveOverride);
        }

        @Override
        public long getKeepAliveDuration(HttpResponse response, HttpContext context) {
            Args.notNull(response, "HTTP response");
            final HeaderElementIterator it = new BasicHeaderElementIterator(
                    response.headerIterator(HTTP.CONN_KEEP_ALIVE));
            while (it.hasNext()) {
                final HeaderElement he = it.nextElement();
                final String param = he.getName();
                final String value = he.getValue();
                if (value != null && param.equalsIgnoreCase("timeout")) {
                    try {
                        return TimeUnit.SECONDS.toMillis(Long.parseLong(value));
                    } catch (final NumberFormatException ignore) {
                    }
                }
            }
            return keepAliveOverride;
        }
    }

    private static class RawClickhouseHost {
        private final URI hostURI;
        private final String dcName;

        public RawClickhouseHost(URI hostURI, String dcName) {
            this.hostURI = hostURI;
            this.dcName = dcName;
        }
    }

    private static class ClickhouseUserException extends ClickhouseException {
        public ClickhouseUserException(String message, String query, String error) {
            super(message, query, error);
        }
    }

    private class NodeUpdater implements Runnable {
        @Override
        public void run() {
            try {
                checkHosts(false);
            } catch (Throwable e) {
                log.error("checkHosts failed", e);
            }
        }
    }

    public interface TxCallback<T> {
        T execute() throws ClickhouseException;
    }

    public enum DefaultUser {
        OFFLINE,
        USER,
    }
}
