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

import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * @author aherman
 */
public abstract class QueryPolicyBuilder {
    protected Function<List<ClickhouseHost>, List<ClickhouseHost>> hostSelector;

    public Function<List<ClickhouseHost>, List<ClickhouseHost>> build() {
        return hostSelector;
    }

    public static QueryPolicyBuilder alive() {
        return new QueryPolicyBuilder() {
            {
                hostSelector = hosts -> hosts.stream().filter(h -> !h.isDown()).collect(Collectors.toList());
            }
        };
    }

    public static QueryPolicyBuilder all() {
        return new QueryPolicyBuilder() {
            {
                hostSelector = hosts -> hosts;
            }
        };
    }

    public QueryPolicyBuilder atRandom() {
        return new QueryPolicyBuilder() {
            {
                hostSelector = QueryPolicyBuilder.this.hostSelector.andThen(
                        hosts -> { Collections.shuffle(hosts); return hosts; }
                );
            }
        };
    }

    public QueryPolicyBuilder inDcOnly(String dcName) {
        return new QueryPolicyBuilder() {
            {
                hostSelector = QueryPolicyBuilder.this.hostSelector.andThen(
                        hosts -> hosts.stream().filter(h -> dcName.equals(h.getDcName())).collect(Collectors.toList())
                );
            }
        };
    }

    public QueryPolicyBuilder inDcFirst(String dcName) {
        return new QueryPolicyBuilder() {
            {
                hostSelector = QueryPolicyBuilder.this.hostSelector.andThen(
                        hosts -> hosts.stream()
                                .sorted((h1, h2) -> {
                                    if (dcName.equals(h1.getDcName())) {
                                        if (dcName.equals(h2.getDcName())) {
                                            return 0;
                                        }
                                        return -1;
                                    } else if (dcName.equals(h2.getDcName())) {
                                        return 1;
                                    }
                                    return 0;
                                })
                                .collect(Collectors.toList())
                );
            }
        };
    }

    public QueryPolicyBuilder byNumber(int number) {
        return new QueryPolicyBuilder() {
            {
                hostSelector = QueryPolicyBuilder.this.hostSelector.andThen(
                        hosts -> {
                            int i = number % hosts.size();
                            return Collections.singletonList(hosts.get(i));
                        }
                );
            }
        };
    }

    public QueryPolicyBuilder byHost(String hostname) {
        return new QueryPolicyBuilder() {
            {
                hostSelector = QueryPolicyBuilder.this.hostSelector.andThen(
                        hosts -> {
                            Optional<ClickhouseHost> host = hosts.stream()
                                    .filter(h -> h.getHostURI().getHost().equals(hostname))
                                    .findFirst();
                            if (host.isPresent()) {
                                return Collections.singletonList(host.get());
                            } else {
                                return Collections.emptyList();
                            }
                        }
                );
            }
        };
    }
}
