/* DELETE ME */
import org.xbill.DNS.Cache;
import org.xbill.DNS.Credibility;
import org.xbill.DNS.DClass;
import org.xbill.DNS.ExtendedResolver;
import org.xbill.DNS.FindServer;
import org.xbill.DNS.Message;
import org.xbill.DNS.MXRecord;
import org.xbill.DNS.Name;
import org.xbill.DNS.Rcode;
import org.xbill.DNS.Record;
import org.xbill.DNS.Resolver;
import org.xbill.DNS.RRset;
import org.xbill.DNS.SetResponse;
import org.xbill.DNS.TextParseException;
import org.xbill.DNS.Type;
/* DELETE ME */

import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.*;

public class MXTest
{
    private Collection serverNames = Collections.synchronizedSet(new HashSet());
    private Collection foreignSite = Collections.synchronizedSet(new WeakSet());
    private DNSServer dns = null;
    private String mxServer = null;
    private long successTime, failureTime;
    private long success, failure;

    boolean isLocalDomain(String domain)
    {
        boolean isLocal = true;
        long startTime = System.currentTimeMillis();
        if (!serverNames.contains(domain))
        {
            if (foreignSite.contains(domain)) isLocal = false;
            else
            {
                // Collection mxRecords = getMailetContext().getMailServers(host);
                Collection mxRecords = dns.findMXRecords(domain);
                if (mxRecords != null && mxRecords.contains(mxServer))
                {
                    serverNames.add(domain);
                }
                else
                {
                    isLocal = false;
                    foreignSite.add(domain);
                }
            }
        }
        long time = System.currentTimeMillis() - startTime;
        if (isLocal)
        {
            successTime += time;
            success++;
        }
        else
        {
            failureTime += time;
            failure++;
        }
        return isLocal;
    }

    public MXTest(String[] args)
    {
        try {
            dns = new DNSServer();
            mxServer = args[0] + ".";

            long startTime = System.currentTimeMillis();
            for (int iter = 0; iter < 1000000; iter++)
            {
                for (int i = 1; i < args.length; i++)
                {
                    boolean isLocal = isLocalDomain(args[i]);
//                  System.out.println(mxServer + " does" + (isLocal ? "" : " not") + " handle mail for " + args[i]);
                }
            }
            System.out.println("Total time: " + (System.currentTimeMillis() - startTime));
            System.out.println("Successes: " + success + ", time: " + successTime);
            System.out.println("Failures: " + failure + ", time: " + failureTime);
        } catch  (Exception e) {
            e.printStackTrace();
        }
    }

    static public void main(String[] args)
    {
        new MXTest(args);
    }
}

/* DELETE ME */
/**
 * Provides DNS client functionality to services running
 * inside James
 */
class DNSServer {

    /**
     * A resolver instance used to retrieve DNS records.  This
     * is a reference to a third party library object.
     */
    private Resolver resolver;

    /**
     * A TTL cache of results received from the DNS server.  This
     * is a reference to a third party library object.
     */
    private Cache cache;

    /**
     * Whether the DNS response is required to be authoritative
     */
    private byte dnsCredibility = Credibility.NONAUTH_ANSWER;

    /**
     * The DNS servers to be used by this service
     */
    private List dnsServers = new ArrayList();

    /**
     * The MX Comparator used in the MX sort.
     */
    private Comparator mxComparator = new MXRecordComparator();

    public DNSServer() throws Exception {
        String[] serversArray = FindServer.servers();
        if (serversArray != null) {
            for ( int i = 0; i < serversArray.length; i++ ) {
                dnsServers.add(serversArray[ i ]);
            }
        }

        //Create the extended resolver...
        serversArray = (String[])dnsServers.toArray(new String[0]);

        for(int c = 0; c < serversArray.length; c++) {
            // System.out.println("DNS Server is: " + serversArray[c]);
        }

        try {
            resolver = new ExtendedResolver( serversArray );
        } catch (UnknownHostException uhe) {
            throw uhe;
        }

        cache = new Cache (DClass.IN);
    }

    /**
     * <p>Return a prioritized unmodifiable list of MX records
     * obtained from the server.</p>
     *
     * @param hostname domain name to look up
     *
     * @return a unmodifiable list of MX records corresponding to
     *         this mail domain name
     */
    public Collection findMXRecords(String hostname) {
        Record answers[] = lookup(hostname, Type.MX);
        List servers = new ArrayList();
        try {
            if (answers == null) {
                return servers;
            }

            MXRecord mxAnswers[] = new MXRecord[answers.length];
            for (int i = 0; i < answers.length; i++) {
                mxAnswers[i] = (MXRecord)answers[i];
            }

            Arrays.sort(mxAnswers, mxComparator);

            for (int i = 0; i < mxAnswers.length; i++) {
                servers.add(mxAnswers[i].getTarget ().toString ());
                // System.out.println(new StringBuffer("Found MX record ").append(mxAnswers[i].getTarget ().toString ()).toString());
            }
            return Collections.unmodifiableCollection(servers);
        } finally {
            //If we found no results, we'll add the original domain name if
            //it's a valid DNS entry
            if (servers.size () == 0) {
                StringBuffer logBuffer =
                    new StringBuffer(128)
                            .append("Couldn't resolve MX records for domain ")
                            .append(hostname)
                            .append(".");
                System.err.println(logBuffer.toString());
                try {
                    InetAddress.getByName(hostname);
                    servers.add(hostname);
                } catch (UnknownHostException uhe) {
                    // The original domain name is not a valid host,
                    // so we can't add it to the server list.  In this
                    // case we return an empty list of servers
                    logBuffer = new StringBuffer(128)
                                .append("Couldn't resolve IP address for host ")
                                .append(hostname)
                                .append(".");
                    System.err.println(logBuffer.toString());
                }
            }
        }
    }

    /**
     * Looks up DNS records of the specified type for the specified name.
     *
     * This method is a public wrapper for the private implementation
     * method
     *
     * @param name the name of the host to be looked up
     * @param type the type of record desired
     */
    public Record[] lookup(String name, short type) {
        return rawDNSLookup(name,false,type);
    }

    /**
     * Looks up DNS records of the specified type for the specified name
     *
     * @param namestr the name of the host to be looked up
     * @param querysent whether the query has already been sent to the DNS servers
     * @param type the type of record desired
     */
    private Record[] rawDNSLookup(String namestr, boolean querysent, short type) {
        Name name = null;
        try {
            name = Name.fromString(namestr, Name.root);
        } catch (TextParseException tpe) {
            // TODO: Figure out how to handle this correctly.
            System.err.println("Couldn't parse name " + namestr);
            tpe.printStackTrace();
            return null;
        }
        short dclass = DClass.IN;

        Record [] answers;
        int answerCount = 0, n = 0;

        SetResponse cached = cache.lookupRecords(name, type, dnsCredibility);
        if (cached.isSuccessful()) {
            // System.out.println(new StringBuffer(256)
            //                 .append("Retrieving MX record for ")
            //                 .append(name).append(" from cache")
            //                 .toString());
            RRset [] rrsets = cached.answers();
            answerCount = 0;
            for (int i = 0; i < rrsets.length; i++) {
                answerCount += rrsets[i].size();
            }

            answers = new Record[answerCount];

            for (int i = 0; i < rrsets.length; i++) {
                Iterator iter = rrsets[i].rrs();
                while (iter.hasNext()) {
                    Record r = (Record)iter.next();
                    answers[n++] = r;
                }
            }
        }
        else if (cached.isNXDOMAIN() || cached.isNXRRSET()) {
            return null;
        }
        else if (querysent) {
            return null;
        }
        else {
            // System.out.println(new StringBuffer(256)
            //                 .append("Looking up MX record for ")
            //                 .append(name)
            //                .toString());
            Record question = Record.newRecord(name, type, dclass);
            Message query = Message.newQuery(question);
            Message response = null;

            try {
                response = resolver.send(query);
            }
            catch (Exception ex) {
                System.err.println("Query error!");
                ex.printStackTrace();
                return null;
            }

            short rcode = response.getHeader().getRcode();
            if (rcode == Rcode.NOERROR || rcode == Rcode.NXDOMAIN) {
                cache.addMessage(response);
            }

            if (rcode != Rcode.NOERROR) {
                return null;
            }

            return rawDNSLookup(namestr, true, type);
        }

        return answers;
    }

    private static class MXRecordComparator
        implements Comparator {

        public int compare (Object a, Object b) {
            MXRecord ma = (MXRecord)a;
            MXRecord mb = (MXRecord)b;
            return ma.getPriority () - mb.getPriority ();
        }
    }
}