package ru.yandex.webmaster3.storage.clickhouse.system.dao;

import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;

import com.google.common.base.Strings;
import com.google.common.collect.Lists;
import org.apache.commons.collections.CollectionUtils;

import ru.yandex.webmaster3.storage.clickhouse.system.data.ClickhouseSystemReplicaInfo;
import ru.yandex.webmaster3.storage.util.clickhouse2.AbstractClickhouseDao;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseException;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseHost;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseQueryContext;
import ru.yandex.webmaster3.storage.util.clickhouse2.query.QueryBuilder;
import ru.yandex.webmaster3.storage.util.clickhouse2.query.Where;

/**
 * @author avhaliullin
 */
public class ClickhouseSystemReplicasCHDao extends AbstractClickhouseDao {
    private static final String TABLE_NAME = "system.replicas";

    public Optional<Integer> getInsertsInQueue(ClickhouseHost host, String db, String table) throws ClickhouseException {
        return getClickhouseServer().executeWithFixedHost(host,
                () -> queryOne(
                        QueryBuilder.select(F.INSERTS_IN_QUEUE)
                                .from(TABLE_NAME)
                                .where(QueryBuilder.eq(F.DATABASE, db))
                                .and(QueryBuilder.eq(F.TABLE, table)),
                        chRow -> chRow.getIntUnsafe(F.INSERTS_IN_QUEUE)
                )
        );
    }

    public Map<String, Integer> getInsertsInQueue(ClickhouseHost host, Collection<String> tables) throws ClickhouseException {
        if (CollectionUtils.isEmpty(tables)) {
            return Collections.emptyMap();
        }
        ClickhouseQueryContext.Builder builder = ClickhouseQueryContext.useDefaults().setHost(host);
        return getClickhouseServer()
                .collectAll(builder,
                        QueryBuilder.select(F.DATABASE, F.TABLE, F.INSERTS_IN_QUEUE).from(TABLE_NAME)
                                .where(QueryBuilder.in(F.DATABASE + " || '.' || " + F.TABLE, tables)).toString(),
                        Collectors.toMap(
                                row -> row.getString(F.DATABASE) + "." + row.getString(F.TABLE),
                                row -> row.getIntUnsafe(F.INSERTS_IN_QUEUE)
                        )
                );
    }

    public List<ClickhouseSystemReplicaInfo> getReplicas(
            ClickhouseHost host, String database, boolean retrieveZKFields) throws ClickhouseException {
        return getReplicas(host, database, null, retrieveZKFields);
    }

    public List<ClickhouseSystemReplicaInfo> getReplicas(
            ClickhouseHost host, String database, String tablePrefix, boolean retrieveZKFields) throws ClickhouseException {

        List<String> fields = Lists.newArrayList(F.DATABASE, F.TABLE, F.ENGINE, F.IS_LEADER, F.ZOOKEEPER_PATH,
                F.REPLICA_NAME, F.REPLICA_PATH, F.COLUMNS_VERSION, F.INSERTS_IN_QUEUE, F.MERGES_IN_QUEUE);
        if (retrieveZKFields) {
            fields.add(F.TOTAL_REPLICAS);
            fields.add(F.ACTIVE_REPLICAS);
        }

        Where st = QueryBuilder.select(fields).from(TABLE_NAME).where(QueryBuilder.eq(F.DATABASE, database));
        if (!Strings.isNullOrEmpty(tablePrefix)) {
            st = st.and(QueryBuilder.startsWith(F.TABLE, tablePrefix));
        }

        return getClickhouseServer().queryAll(ClickhouseQueryContext.useDefaults().setHost(host), st.toString(),
                row -> new ClickhouseSystemReplicaInfo(
                        row.getString(F.DATABASE),
                        row.getString(F.TABLE),
                        row.getString(F.ENGINE),
                        row.getIntUnsafe(F.IS_LEADER) == 1,
                        row.getString(F.ZOOKEEPER_PATH),
                        row.getString(F.REPLICA_NAME),
                        row.getString(F.REPLICA_PATH),
                        row.getIntUnsafe(F.COLUMNS_VERSION),
                        row.getIntUnsafe(F.INSERTS_IN_QUEUE),
                        row.getIntUnsafe(F.MERGES_IN_QUEUE),
                        retrieveZKFields ? row.getIntUnsafe(F.TOTAL_REPLICAS) : 0,
                        retrieveZKFields ? row.getIntUnsafe(F.ACTIVE_REPLICAS) : 0
                )
        );
    }

    private interface F {
        String DATABASE = "database";
        String TABLE = "table";
        String ENGINE = "engine";
        String IS_LEADER = "is_leader";
        String ZOOKEEPER_PATH = "zookeeper_path";
        String REPLICA_NAME = "replica_name";
        String REPLICA_PATH = "replica_path";
        String COLUMNS_VERSION = "columns_version";
        String INSERTS_IN_QUEUE = "inserts_in_queue";
        String MERGES_IN_QUEUE = "merges_in_queue";
        String TOTAL_REPLICAS = "total_replicas";
        String ACTIVE_REPLICAS = "active_replicas";

    }
}
