package ru.yandex.webmaster3.storage.verification.dns;

import java.net.UnknownHostException;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.stream.Collectors;

import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.xbill.DNS.Cache;
import org.xbill.DNS.DClass;
import org.xbill.DNS.ExtendedResolver;
import org.xbill.DNS.Lookup;
import org.xbill.DNS.Record;
import org.xbill.DNS.TextParseException;
import org.xbill.DNS.Type;

import ru.yandex.webmaster3.core.WebmasterException;
import ru.yandex.webmaster3.core.host.verification.fail.DNSRecord;
import ru.yandex.webmaster3.core.http.WebmasterErrorResponse;

/**
 * @author akhazhoyan 06/2018
 */
public class DnsLookupService {

    private static final Logger log = LoggerFactory.getLogger(DnsLookupService.class);
    private static final int TRY_AGAIN_MAX_ATTEMPTS = 5;

    public enum LookupStatus {
        SUCCESSFUL,
        HOST_NOT_FOUND,
        TYPE_NOT_FOUND,
        TRY_AGAIN,
        UNRECOVERABLE;

        private static LookupStatus fromIntegerLookupResult(int lookupResult) {
            switch (lookupResult) {
                case Lookup.HOST_NOT_FOUND:
                    return HOST_NOT_FOUND;
                case Lookup.SUCCESSFUL:
                    return SUCCESSFUL;
                case Lookup.TRY_AGAIN:
                    return TRY_AGAIN;
                case Lookup.TYPE_NOT_FOUND:
                    return TYPE_NOT_FOUND;
                case Lookup.UNRECOVERABLE:
                    return UNRECOVERABLE;
                default:
                    throw new IllegalArgumentException("Unexpected lookup result value: " + lookupResult);
            }
        }
    }

    private Pair<LookupStatus, List<DNSRecord>> doResolve(Collection<String> nameServers, String hostName, int recordType) {
        int attemptsCount = 0;
        Pair<Lookup, List<DNSRecord>> pair;
        do {
            pair = doResolveNoRetries(nameServers, hostName, recordType);
            Lookup lookup = pair.getLeft();
            if (lookup.getResult() != Lookup.TRY_AGAIN) {
                break;
            }

            attemptsCount++;
            log.info("Retrying DNS request, attempt {}", attemptsCount);
        } while (attemptsCount < TRY_AGAIN_MAX_ATTEMPTS);

        Lookup lookup = pair.getLeft();
        List<DNSRecord> recordList = pair.getRight();
        var lookupStatus = LookupStatus.fromIntegerLookupResult(lookup.getResult());
        switch (lookupStatus) {
            case HOST_NOT_FOUND:
            case SUCCESSFUL:
            case TYPE_NOT_FOUND:
                return Pair.of(lookupStatus, recordList);

            default:
                String error = String.valueOf(lookup.getErrorString());
                throw new WebmasterException("DNS lookup failed with status " + lookupStatus + ": " + error,
                        new WebmasterErrorResponse.InternalUnknownErrorResponse(DnsLookupService.class, null)
                );
        }
    }

    private Pair<Lookup, List<DNSRecord>> doResolveNoRetries(Collection<String> nameServers, String hostName, int recordType) {
        log.info("Running DNS lookup against {} with type={}", hostName, recordType);
        try {
            Lookup lookup = new Lookup(hostName, recordType);
            Cache dnsCache = new Cache();
            lookup.setCache(dnsCache);
            lookup.setSearchPath(new String[]{});
            if (!CollectionUtils.isEmpty(nameServers)) {
                lookup.setResolver(new ExtendedResolver(nameServers.toArray(new String[0])));
            }

            Record[] records = lookup.run();
            if (records == null) {
                records = new Record[0];
            }

            log.info("Got DNS records: {}", Arrays.toString(records));

            List<DNSRecord> recordList = Arrays.stream(records)
                .filter(r -> r.getType() == recordType)
                .map(r -> new DNSRecord(
                        r.getName().toString(),
                        Type.string(r.getType()),
                        DClass.string(r.getDClass()),
                        r.getTTL(),
                        r.rdataToString()
                ))
                .collect(Collectors.toList());

            log.info("DNS lookup status: {}", LookupStatus.fromIntegerLookupResult(lookup.getResult()));
            return Pair.of(lookup, recordList);
        } catch (TextParseException | UnknownHostException e) {
            throw new WebmasterException("Failed to get dns records or host " + hostName,
                    new WebmasterErrorResponse.InternalUnknownErrorResponse(DnsLookupService.class, null), e
            );
        }
    }

    /**
     * @param hostName   хост в punycode
     * @param recordType тип DNS-записи из класса {@link org.xbill.DNS.Type}
     * @return пару из статуса обхода и списка DNS-записей данного типа ({@code record.getType() == recordType})
     * @throws WebmasterException если библиотека org.xbill.DNS бросает исключение
     */
    public Pair<LookupStatus, List<DNSRecord>> lookup(Collection<String> nameServers, String hostName, int recordType) {
        return doResolve(nameServers, hostName, recordType);
    }

    /**
     * @param hostName   хост в punycode
     * @param recordType тип DNS-записи из класса {@link org.xbill.DNS.Type}
     * @return пару из статуса обхода и списка DNS-записей данного типа ({@code record.getType() == recordType})
     * @throws WebmasterException если библиотека org.xbill.DNS бросает исключение
     */
    public Pair<LookupStatus, List<DNSRecord>> lookup(String hostName, int recordType) {
        var pair = doResolve(null, hostName, Type.NS);
        if (recordType == Type.NS) {
            return pair;
        }

        LookupStatus lookupStatus = pair.getLeft();
        if (lookupStatus == LookupStatus.SUCCESSFUL) {
            List<DNSRecord> recordList = pair.getRight();
            List<String> nameServers = recordList.stream().map(DNSRecord::getData).collect(Collectors.toList());
            log.info("Found name servers: {}", String.join(", ", nameServers));
            try {
                pair = doResolve(nameServers, hostName, recordType);
                lookupStatus = pair.getLeft();

                if (lookupStatus == LookupStatus.SUCCESSFUL) {
                    return pair;
                }
            } catch (Exception e) {
                log.error("Failed to make the optimized DNS lookup", e);

            }
        }

        return doResolve(null, hostName, recordType);
    }

}
