package org.smpte._2071._2012.mdcd.naming.dnsjava;

import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.Set;
import java.util.Stack;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.logging.Level;

import javax.naming.NamingException;

import org.smpte._2071._2012.mdcd.ServiceInstance;
import org.smpte._2071._2012.mdcd.ServiceName;
import org.smpte._2071._2012.mdcd.impl.Utils;
import org.smpte._2071._2012.mdcd.naming.AbstractNamingServiceImpl;
import org.smpte._2071._2012.mdcd.naming.CompoundKeyMap;
import org.smpte._2071._2012.mdcd.naming.NameRecord;
import org.smpte._2071._2012.mdcd.net.InetAddressUtils;
import org.xbill.DNS.AAAARecord;
import org.xbill.DNS.ARecord;
import org.xbill.DNS.CNAMERecord;
import org.xbill.DNS.DClass;
import org.xbill.DNS.ExtendedResolver;
import org.xbill.DNS.Lookup;
import org.xbill.DNS.MXRecord;
import org.xbill.DNS.Message;
import org.xbill.DNS.NAPTRRecord;
import org.xbill.DNS.NSRecord;
import org.xbill.DNS.Name;
import org.xbill.DNS.PTRRecord;
import org.xbill.DNS.Record;
import org.xbill.DNS.Resolver;
import org.xbill.DNS.SOARecord;
import org.xbill.DNS.SRVRecord;
import org.xbill.DNS.Section;
import org.xbill.DNS.SimpleResolver;
import org.xbill.DNS.TXTRecord;
import org.xbill.DNS.TextParseException;
import org.xbill.DNS.Type;

public class NamingServiceImpl extends AbstractNamingServiceImpl
{
    private static final Comparator<? super Record> DNSComparator = new Comparator<Record>()
    {
        @Override
        public int compare(Record o1, Record o2)
        {
            if (o1 == o2 || o1.equals(o2))
            {
                return 0;
            }
            
            int n = o1.getType() - o2.getType();
            if (n != 0)
            {
                return n;
            }
            
            n = o1.getDClass() - o2.getDClass();
            if (n != 0)
            {
                return n;
            }
            
            n = o1.getName().compareTo(o2.getName());
            if (n != 0)
            {
                return n;
            }
            
            return 0;
        }
    };

    private static final long ERROR_THRESHOLD = 60000;
    
    private ExtendedResolver defaultResolver = null;
    
    private Queue<SimpleResolver> specifiedResolvers = null;
    
    private Map<Resolver, Long> resolverErrors = new HashMap<Resolver, Long>();

    
    public NamingServiceImpl()
    throws NamingException
    {
        super();
        
        try
        {
            defaultResolver = new ExtendedResolver();
        } catch (Exception e)
        {
            Utils.log(log, Level.WARNING, "Could not create default resolver " + (e.getMessage() != null ? " - " + e.getMessage() : "."), Level.FINE, e);
        }
    }
    
    
    @Override
    public void setLocalAddress(final InetAddress localAddress)
    throws NamingException
    {
        super.setLocalAddress(localAddress);
        
        if (specifiedResolvers != null)
        {
            for (SimpleResolver resolver : specifiedResolvers)
            {
                resolver.setLocalAddress(localAddress);
            }
        } else
        {
            Stack<Resolver> stack = new Stack<Resolver>();
            stack.addAll(Arrays.asList(defaultResolver.getResolvers()));
            while (!stack.isEmpty())
            {
                Resolver resolver = stack.pop();
                if (resolver instanceof SimpleResolver)
                {
                    ((SimpleResolver) resolver).setLocalAddress(localAddress);
                } else if (resolver instanceof ExtendedResolver)
                {
                    stack.addAll(Arrays.asList(((ExtendedResolver) resolver).getResolvers()));
                }
            }
        }
    }
    
    
    @Override
    public void setServerAddresses(Set<InetSocketAddress> serverAddresses)
    throws NamingException
    {
        if (this.specifiedResolvers != null)
        {
            this.specifiedResolvers.clear();
        }
        
        if (serverAddresses != null)
        {
            this.serverAddresses.clear();
        }
        
        addServerAddresses(serverAddresses);
    }
    
    
    @Override
    public boolean addServerAddresses(Set<InetSocketAddress> serverAddresses)
    throws NamingException
    {
        if (this.specifiedResolvers == null)
        {
            this.specifiedResolvers = new ConcurrentLinkedQueue<SimpleResolver>();
        }
        
        if (serverAddresses == null)
        {
            this.serverAddresses = new LinkedHashSet<InetSocketAddress>();
        }
        
        HashSet<InetSocketAddress> newServerAddresses = new HashSet<InetSocketAddress>();
        for (InetSocketAddress serverAddress : serverAddresses)
        {
            if (!this.serverAddresses.contains(serverAddress))
            {
                try
                {
                    SimpleResolver newResolver = new SimpleResolver(serverAddress.getAddress().getHostAddress());
                    newResolver.setAddress(serverAddress);
                    newResolver.setLocalAddress(localAddress);
                    
                    Lookup lookup = new Lookup("@", Type.SOA);
                    lookup.setResolver(newResolver);
                    lookup.run();
                    if (lookup.getResult() == Lookup.SUCCESSFUL)
                    {
                        this.specifiedResolvers.add(newResolver);
                        newServerAddresses.add(serverAddress);
                    } else
                    {
                        log.warning("Could not connect to Naming Service \"" + this + " at server address \"" + serverAddress + "\", address not added.");
                    }
                } catch (TextParseException e)
                {
                    NamingException ne = new NamingException(e.getMessage());
                    ne.setStackTrace(e.getStackTrace());
                    throw ne;
                } catch (UnknownHostException e)
                {
                    NamingException ne = new NamingException(e.getMessage());
                    ne.setStackTrace(e.getStackTrace());
                    throw ne;
                }
            }
        }
        
        if (newServerAddresses.size() > 0)
        {
            return super.addServerAddresses(newServerAddresses);
        } else
        {
            return false;
        }
    }
    
    
    @Override
    public List<NameRecord> reverseQuery(InetAddress address)
    throws NamingException
    {
        List<NameRecord> results = new ArrayList<NameRecord>();
        
        List<ResourceRecord> records = new ArrayList<ResourceRecord>(); 
        List<ResourceRecord> tempRecords;
        
        int[] queryTypes = new int[] {Type.PTR};
        
        for (int type : queryTypes)
        {
            tempRecords = listRecords(new ServiceName(InetAddressUtils.reverseMapAddress(address)), type);
            if (tempRecords != null && tempRecords.size() > 0)
            {
                records.addAll(tempRecords);
            }
        }
        
        if (records.size() > 0)
        {
            for (ResourceRecord record : records)
            {
                Name host = record.getHost();
                Name target = record.getTarget();
                Name value = host != null ? host : target;
                results.add(new NameRecord(value.toString(), address));
            }
        }
        
        return results;
    }
    
    
    @Override
    public List<NameRecord> query(ServiceName serviceName)
    throws NamingException
    {
        List<NameRecord> results = new ArrayList<NameRecord>();
        
        List<ResourceRecord> records = new ArrayList<ResourceRecord>(); 
        List<ResourceRecord> tempRecords;
        
        int[] queryTypes = new int[] {Type.A, Type.PTR, Type.SRV};
        
        for (int type : queryTypes)
        {
            tempRecords = listRecords(serviceName, type);
            if (tempRecords != null && tempRecords.size() > 0)
            {
                records.addAll(tempRecords);
            }
        }
        
        if (records.size() > 0)
        {
            for (ResourceRecord record : records)
            {
                Name host = record.getHost();
                Name target = record.getTarget();
                Name value = host != null ? host : target;
                results.add(new NameRecord(value.toString(), record.getAddresses()));
            }
        }
        
        return results;
    }


    @Override
    public List<ServiceInstance> serviceDiscovery(ServiceName serviceName)
    throws NamingException
    {
        ArrayList<ServiceInstance> services = new ArrayList<ServiceInstance>();
        List<ResourceRecord> records = new ArrayList<ResourceRecord>(); 
        List<ResourceRecord> tempRecords;
        
        int[] queryTypes = new int[] {Type.PTR, Type.SRV, Type.TXT};
        
        for (int type : queryTypes)
        {
            tempRecords = listRecords(serviceName, type);
            if (tempRecords != null && tempRecords.size() > 0)
            {
                records.addAll(tempRecords);
            }
        }
        
        if (records != null && records.size() > 0)
        {
            CompoundKeyMap<Name, Integer, ResourceRecord> cache = new CompoundKeyMap<Name, Integer, ResourceRecord>();
            
            Collections.sort(records, new Comparator<ResourceRecord>()
            {
                @Override
                public int compare(ResourceRecord nr1, ResourceRecord nr2)
                {
                    int value = nr1.getName().compareTo(nr2.getName());
                    if (value == 0)
                    {
                        int type1 = nr1.getType();
                        int type2 = nr2.getType();
                        
                        if (type1 == type2)
                        {
                            return 0;
                        }
                        
                        switch (type1)
                        {
                            case Type.SRV :
                                return -1;
                            case Type.TXT :
                                return (type2 == Type.A || type2 == Type.CNAME ? -1 : 1);
                            case Type.CNAME :
                                return (type2 == Type.A ? -1 : 1);
                            case Type.A :
                                return 1;
                            default :
                                if (type2 == Type.SRV)
                                {
                                    return 1;
                                } else
                                {
                                    return (type1 < type2 ? -1 : (type1 > type2 ? 1 : 0));
                                }
                        }
                    }
                    
                    return value;
                }
            });
            
            // Cache Complete List of results, but service target and type.
            for (ResourceRecord record : records)
            {
                int type = record.getType();
                Name target = record.getTarget();
                String targetStr = target != null ? target.toString() : null;
                cache.put(target, record.getType(), record);
                
                if (type == Type.PTR)
                {
                    if (target != null)
                    {
                        services.addAll(serviceDiscovery(new ServiceName(targetStr, null)));
                    }
                } else if (type == Type.SRV)
                {
                    ServiceInstance service = new ServiceInstance(serviceName);
                    services.add(service);
                    
                    Name host = record.getHost();
                    service.setHost(host != null ? host.toString() : null);
                    service.setPort(record.getPort());
                    service.setPriority(record.getPriority());
                    service.setTarget(targetStr);
                    service.setWeight(record.getWeight());
                    if (targetStr != null)
                    {
                        List<InetAddress> serviceAddresses = listAddresses(targetStr);
                        if (serviceAddresses != null)
                        {
                            service.setAddresses(serviceAddresses);
                        }
                    }
                } else if (type == Type.TXT)
                {
                    if (services.size() > 0)
                    {
                        ServiceInstance service = services.get(services.size() - 1);
                        if (record.getTextRecords() != null)
                        {
                            service.addTextRecords(record.getTextRecords());
                        } else
                        {
                            Exception e = new Exception("We Recieved a TXT record with no Text!  This should not have happened!");
                            Utils.log(log, Level.WARNING, e.getMessage(), Level.FINE, e);
                        }
                    }
                }
            }
        }
            
        return services;
    }


    private List<ResourceRecord> listRecords(ServiceName serviceName, int type)
    throws NamingException
    {
        List<ResourceRecord> records = new ArrayList<ResourceRecord>();
        List<String> domains = new ArrayList<String>();
        
        final String domain = serviceName.getDomain();
        
        if (domain == null || domain.length() == 0)
        {
            synchronized (this.domains)
            {
                domains.addAll(this.domains);
            }
        } else
        {
            domains.add(domain);
        }
        
        if (this.specifiedResolvers != null)
        {
            for (String tempDomain : domains)
            {
                for (SimpleResolver resolver : this.specifiedResolvers)
                {
                    try
                    {
                        Long lastError = resolverErrors.get(resolver);
                        if (lastError == null || lastError + ERROR_THRESHOLD < System.currentTimeMillis())
                        {
                            records.addAll(lookup(resolver, type, serviceName.getNameWithoutDomain(), tempDomain));
                        }
                    } catch (NamingException e)
                    {
                        Utils.log(log, Level.WARNING, "Error querying \"" + serviceName.getNameWithoutDomain() + "." + tempDomain + "\" of type \"" + typeToString(type) + "\"" + (e.getMessage() != null ? " - " + e.getMessage() : "."), Level.FINE, e);
                    } catch (IOException e)
                    {
                        resolverErrors.put(resolver, System.currentTimeMillis());
                        Utils.log(log, Level.WARNING, "Error querying \"" + serviceName.getNameWithoutDomain() + "." + tempDomain + "\" of type \"" + typeToString(type) + "\"" + (e.getMessage() != null ? " - " + e.getMessage() : "."), Level.FINE, e);
                    }
                }
            }
        } else
        {
            for (String tempDomain : domains)
            {
                try
                {
                    records.addAll(lookup(this.defaultResolver, type, serviceName.getNameWithoutDomain(), tempDomain));
                } catch (NamingException e)
                {
                    Utils.log(log, Level.WARNING, "Error querying \"" + serviceName.getNameWithoutDomain() + "." + tempDomain + "\" of type \"" + typeToString(type) + "\"" + (e.getMessage() != null ? " - " + e.getMessage() : "."), Level.FINE, e);
                } catch (IOException e)
                {
                    Utils.log(log, Level.WARNING, "Error querying \"" + serviceName.getNameWithoutDomain() + "." + tempDomain + "\" of type \"" + typeToString(type) + "\"" + (e.getMessage() != null ? " - " + e.getMessage() : "."), Level.FINE, e);
                }
            }
        }
        
        return records;
    }
    
    
    public List<InetAddress> listAddresses(String name)
    throws NamingException
    {
        List<InetAddress> addresses = new ArrayList<InetAddress>();
        List<ResourceRecord> records = new ArrayList<ResourceRecord>(); 
        List<ResourceRecord> tempRecords;
        
        int[] queryTypes = new int[] {Type.A, Type.AAAA, Type.CNAME, Type.PTR, Type.SRV};
        
        for (int type : queryTypes)
        {
            tempRecords = listRecords(new ServiceName(name), type);
            if (tempRecords != null && tempRecords.size() > 0)
            {
                records.addAll(tempRecords);
            }
        }
        
        if (records != null)
        {
            for (ResourceRecord record : records)
            {
                Name target = record.getTarget();
                switch (record.getType())
                {
                    case Type.A :
                    case Type.AAAA :
                        addresses.addAll(record.getAddresses());
                        break;
                    case Type.CNAME :
                    case Type.PTR :
                    case Type.SRV :
                        try
                        {
                            if (target != null)
                            {
                                addresses.addAll(listAddresses(target.toString()));
                            }
                        } catch (Exception e)
                        {
                            Utils.log(log, Level.WARNING, "Error resolving addresses for \"" + name + "\"" + (e.getMessage() != null ? " - " + e.getMessage() : "."), Level.FINE, e);
                        }
                        break;
                }
            }
        }
        
        return addresses;
    }


    /**
     * Lookup Naming Records from the resolver.
     * 
     * @param resolver the DNS Resolver.
     * @param name The host name
     * @param domain The domain name.  If null, name is expected to be fully qualified.
     * @return
     * @throws NamingException
     * @throws IOException 
     */
    @SuppressWarnings("unchecked")
    private List<ResourceRecord> lookup(Resolver resolver, int type, String name, String domain)
    throws NamingException, IOException
    {
        List<ResourceRecord> records = new ArrayList<ResourceRecord>();
        
        Name domainName;
        Name hostName;
        if (domain != null)
        {
            if (!domain.endsWith("."))
            {
                domain += ".";
            }
            domainName = new Name(domain);
            hostName = new Name(name, domainName);
        } else
        {
            if (!name.endsWith("."))
            {
                name += ".";
            }
            domainName = null;
            hostName = new Name(name);
        }
        
        Message query = Message.newQuery(Record.newRecord(hostName, type, DClass.IN));
        Message response = null;
        boolean errorOccurred = false;
        byte retry = 1;
        do
        {
            try
            {
                response = resolver.send(query);
                errorOccurred = false;
                break;
            } catch (IOException e)
            {
                errorOccurred = true;
                try
                {
                    Thread.sleep(500);
                } catch (InterruptedException ie)
                {
                    // ignore
                }
                
                if (retry++ > 3)
                {
                    throw e;
                }
            }
        } while (errorOccurred);
        
        switch (response.getRcode())
        {
            case 1 :
                throw new NamingException("Error querying " + typeToString(type) + " \"" + hostName + "\" - DNS error code 1: Format Error");
            case 2 :
                throw new NamingException("Error querying " + typeToString(type) + " \"" + hostName + "\" - DNS error code 2: Server Failure");
            case 3 :
                Utils.log(log, Level.WARNING, "Issue querying " + typeToString(type) + " \"" + hostName + "\" - DNS error code 3: Name Not Found", Level.FINE, new NamingException("DNS error code 3: Name Not Found"));
                break;
            case 4 :
                throw new NamingException("Error querying " + typeToString(type) + " \"" + hostName + "\" - DNS error code 4: Not Implemented");
            case 5 :
                throw new NamingException("Error querying " + typeToString(type) + " \"" + hostName + "\" - DNS error code 5: Request Refused");
        }

        Record[] answers = response.getSectionArray(Section.ANSWER);
        Arrays.sort(answers, DNSComparator);
        if (answers.length > 0)
        {
            for (Record answer : answers)
            {
                ResourceRecord record = null;
                switch(answer.getType())
                {
                    case Type.A :
                        ARecord a = (ARecord) answer;
                        record = new ResourceRecord(answer.getType(), answer.getName());
                        record.setHost(answer.getName());
                        record.addAddress(a.getAddress());
                        records.add(record);
                        break;
                    case Type.AAAA :
                        AAAARecord aaaa = (AAAARecord) answer;
                        record = new ResourceRecord(answer.getType(), answer.getName());
                        record.setHost(answer.getName());
                        record.addAddress(aaaa.getAddress());
                        records.add(record);
                        break;
                    case Type.CNAME :
                        CNAMERecord cname = (CNAMERecord) answer;
                        record = new ResourceRecord(answer.getType(), answer.getName());
                        record.setHost(cname.getTarget());
                        record.setTarget(cname.getTarget());
                        break;
                    case Type.PTR :
                        PTRRecord ptr = (PTRRecord) answer;
                        record = new ResourceRecord(answer.getType(), answer.getName());
                        record.setHost(ptr.getTarget());
                        record.setTarget(ptr.getTarget());
                        records.add(record);
                        break;
                    case Type.SRV :
                        SRVRecord srv = (SRVRecord) answer;
                        record = new ResourceRecord(answer.getType(), answer.getName());
                        record.setHost(srv.getTarget());
                        record.setPort(srv.getPort());
                        record.setPriority(srv.getPriority());
                        record.setWeight(srv.getWeight());
                        record.setTarget(srv.getTarget());
                        records.add(record);
                        break;
                    case Type.TXT :
                        TXTRecord txt = (TXTRecord) answer;
                        record = new ResourceRecord(answer.getType(), answer.getName());
                        record.setHost(null);
                        record.addTextRecords(txt.getStrings());
                        records.add(record);
                        break;
                    case Type.NAPTR :
                        NAPTRRecord naptr = (NAPTRRecord) answer;
                        record = new ResourceRecord(answer.getType(), answer.getName());
                        record.setFlags(naptr.getFlags());
                        record.setOrder(naptr.getOrder());
                        record.setPreference(naptr.getPreference());
                        record.setRegexp(naptr.getRegexp());
                        record.setReplacement(naptr.getReplacement());
                        record.setService(naptr.getService());
                        records.add(record);
                        break;
                    case Type.NS :
                        NSRecord ns = (NSRecord) answer;
                        record = new ResourceRecord(answer.getType(), answer.getName());
                        record.setHost(ns.getTarget());
                        record.setTarget(ns.getTarget());
                        records.add(record);
                        break;
                    case Type.MX :
                        MXRecord mx = (MXRecord) answer;
                        record = new ResourceRecord(answer.getType(), answer.getName());
                        record.setPriority(mx.getPriority());
                        record.setHost(mx.getTarget());
                        record.setTarget(mx.getTarget());
                        records.add(record);
                        break;
                    case Type.SOA :
                        SOARecord soa = (SOARecord) answer;
                        record = new ResourceRecord(answer.getType(), answer.getName());
                        record.setHost(soa.getHost());
                        record.setTarget(soa.getName());
                        records.add(record);
                        break;
                }
            }
        }
        
        return records;
    }
    
    
    private String typeToString(int type)
    {
        switch (type)
        {
            case Type.A :
                return "A";
            case Type.A6 :
                return "A6";
            case Type.AAAA :
                return "AAAA";
            case Type.ANY :
                return "ANY";
            case Type.AXFR :
                return "AXFR";
            case Type.CERT :
                return "CERT";
            case Type.CNAME :
                return "CNAME";
            case Type.DHCID :
                return "DHCID";
            case Type.DNAME :
                return "DNAME";
            case Type.DNSKEY :
                return "DNSKEY";
            case Type.DS :
                return "DS";
            case Type.EID :
                return "EID";
            case Type.IPSECKEY :
                return "IPSECKEY";
            case Type.IXFR :
                return "IXFR";
            case Type.KEY :
                return "KEY";
            case Type.MX :
                return "MX";
            case Type.NAPTR :
                return "NAPTR";
            case Type.NS :
                return "NS";
            case Type.PTR :
                return "PTR";
            case Type.SOA :
                return "SOA";
            case Type.SIG :
                return "SIG";
            case Type.TKEY :
                return "TKEY";
            case Type.TSIG :
                return "TSIG";
            case Type.TXT :
                return "TXT";
            case Type.SRV :
                return "SRV";
            default :
                return Integer.toBinaryString(type);
        }
    }
}