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

import com.google.common.collect.ImmutableMap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.*;
import java.net.Socket;
import java.nio.charset.Charset;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;


/**
 * @author avhaliullin
 */
public class WhoisRecordReader implements Closeable {
    private static final Logger log = LoggerFactory.getLogger(WhoisRecordReader.class);

    public static final long WAITING_TIME = 2100;

    private static final Pattern REDIRECT = Pattern.compile("\\s*whois server:\\s+([\\w.-]+)\\s*", Pattern.CASE_INSENSITIVE);
    private static final Pattern REMOVE_DOMAIN_WORD =
            Pattern.compile("(.*no match for .*domain.*|.*no entries found.*|.*domain.*not (found|valid).*|.*no information.*for.*domain.*)|" +
                    "(unsupported character encoding)|invalid query|.*name invalid.*|.*queried object does not exist.*|" +
                    ".*no object found|(.*Syntax error.*)", Pattern.CASE_INSENSITIVE);
    private static final char[] END_OF_STREAM = {13, 10};
    private static final Map<String, String> CUSTOM_WHOIS_SERVERS = ImmutableMap.<String, String>builder()
            .put("xn--p1acf", "whois.nic.xn--p1acf") // .рус
            .put("group", "whois.nic.ru")
            .put("top", "whois.nic.ru")
            .put("website", "whois.nic.ru")
            .put("info", "whois.nic.ru")
            .put("space", "whois.nic.ru")
            .put("site", "whois.nic.ru")
            .put("online", "whois.nic.ru")
            .build();

    private final Queue<String> serversQueue;
    private final Set<String> serversSet;
    private final String hostName;
    private BufferedReader reader;
    private Socket socket;
    private String servName;
    private boolean usingDomainWord = false;

    public WhoisRecordReader(String hostName) {
        this.hostName = hostName;

        String topLevelDomain = hostName.substring(hostName.lastIndexOf('.') + 1);
        String servName;
        if (CUSTOM_WHOIS_SERVERS.containsKey(topLevelDomain)) {
            servName = CUSTOM_WHOIS_SERVERS.get(topLevelDomain);
        } else {
            servName = topLevelDomain + ".whois-servers.net";
        }

        serversQueue = new LinkedList<>();
        serversSet = new HashSet<>();

        serversQueue.add(servName);
        serversSet.add(servName);
    }

    private void connect(String servName, boolean useDomainWord, boolean waiting) throws IOException {
        close();
        if (waiting) {
            log.debug("Waiting for " + WAITING_TIME + "ms.");
            try {
                Thread.sleep(WAITING_TIME);
            } catch (InterruptedException e) {
                log.debug("Interrupted!");
                return;
            }
        }
        log.info("Connecting to \"" + servName + "\", useDomainWord=" + useDomainWord);
        socket = new Socket(servName, 43);
        Writer out = new OutputStreamWriter(socket.getOutputStream(), Charset.forName("US-ASCII"));

        log.info("Asking about \"" + hostName + "\"");
        out.write((useDomainWord ? "domain " : "") + hostName);
        out.write(END_OF_STREAM);
        out.flush();

        reader = new BufferedReader(new InputStreamReader(socket.getInputStream()));
        usingDomainWord = useDomainWord;
    }

    public String readLine() throws IOException {
        String line;
        if (reader == null || (line = reader.readLine()) == null) {
            if (serversQueue.isEmpty()) {
                close();
                return null;
            }

            servName = serversQueue.poll();

            connect(servName, true, false);
            return readLine();
        }

        Matcher redirectMatcher = REDIRECT.matcher(line);
        if (redirectMatcher.find()) {
            String redirectedTo = redirectMatcher.group(1);
            if (!serversSet.contains(redirectedTo)) {
                log.debug("Redirect to \"" + redirectedTo + "\" found. Line: \"" + line + "\"");
                serversSet.add(redirectedTo);
                serversQueue.add(redirectedTo);
            }
        }

        if (usingDomainWord && REMOVE_DOMAIN_WORD.matcher(line).find()) {
            log.debug("Using \"domain\" prefix is unsupported. Line: " + line);
            connect(servName, false, true);
            return readLine();
        }
        return line;
    }

    @Override
    public void close() throws IOException {
        if (socket != null) {
            socket.close();
        }
    }

    public static void main(String[] args) throws Exception {
        WhoisRecordReader whoisRecordReader = new WhoisRecordReader("mmyy.site");
        String line;
        while ((line = whoisRecordReader.readLine()) != null) {
            System.out.println(line);
        }
    }
}
