// Copyright (c) 1999-2004 Brian Wellington (bwelling@xbill.org)

package org.xbill.DNS2.clients;

import java.io.IOException;
import java.io.Serializable;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.TreeMap;

import org.xbill.DNS2.resolver.SetResponse;

import dorkbox.network.dns.Name;
import dorkbox.network.dns.constants.DnsClass;
import dorkbox.network.dns.constants.DnsRecordType;
import dorkbox.network.dns.exceptions.ZoneTransferException;
import dorkbox.network.dns.records.DnsRecord;
import dorkbox.network.dns.records.RRset;
import dorkbox.network.dns.records.SOARecord;
import dorkbox.util.OS;

/**
 * A DNS Zone.  This encapsulates all data related to a Zone, and provides
 * convenient lookup methods.
 *
 * @author Brian Wellington
 */

public
class Zone implements Serializable {

    private static final long serialVersionUID = -9220510891189510942L;

    /**
     * A primary zone
     */
    public static final int PRIMARY = 1;

    /**
     * A secondary zone
     */
    public static final int SECONDARY = 2;

    private Map data;
    private Name origin;
    private Object originNode;
    private int dclass = DnsClass.IN;
    private RRset NS;
    private SOARecord SOA;
    private boolean hasWild;


    class ZoneIterator implements Iterator {
        private Iterator zentries;
        private RRset[] current;
        private int count;
        private boolean wantLastSOA;

        ZoneIterator(boolean axfr) {
            synchronized (Zone.this) {
                zentries = data.entrySet()
                               .iterator();
            }
            wantLastSOA = axfr;
            RRset[] sets = allRRsets(originNode);
            current = new RRset[sets.length];
            for (int i = 0, j = 2; i < sets.length; i++) {
                int type = sets[i].getType();
                if (type == DnsRecordType.SOA) {
                    current[0] = sets[i];
                }
                else if (type == DnsRecordType.NS) {
                    current[1] = sets[i];
                }
                else {
                    current[j++] = sets[i];
                }
            }
        }

        @Override
        public
        boolean hasNext() {
            return (current != null || wantLastSOA);
        }

        @Override
        public
        Object next() {
            if (!hasNext()) {
                throw new NoSuchElementException();
            }
            if (current == null) {
                wantLastSOA = false;
                return oneRRset(originNode, DnsRecordType.SOA);
            }
            Object set = current[count++];
            if (count == current.length) {
                current = null;
                while (zentries.hasNext()) {
                    Map.Entry entry = (Map.Entry) zentries.next();
                    if (entry.getKey()
                             .equals(origin)) {
                        continue;
                    }
                    RRset[] sets = allRRsets(entry.getValue());
                    if (sets.length == 0) {
                        continue;
                    }
                    current = sets;
                    count = 0;
                    break;
                }
            }
            return set;
        }

        @Override
        public
        void remove() {
            throw new UnsupportedOperationException();
        }
    }

    /**
     * Creates a Zone from the records in the specified master file.
     *
     * @param zone The name of the zone.
     * @param file The master file to read from.
     *
     * @see Master
     */
    public
    Zone(Name zone, String file) throws IOException {
        data = new TreeMap();

        if (zone == null) {
            throw new IllegalArgumentException("no zone name specified");
        }
        Master m = new Master(file, zone);
        DnsRecord record;

        origin = zone;
        while ((record = m.nextRecord()) != null) {
            maybeAddRecord(record);
        }
        validate();
    }

    private
    void validate() throws IOException {
        originNode = exactName(origin);
        if (originNode == null) {
            throw new IOException(origin + ": no data specified");
        }

        RRset rrset = oneRRset(originNode, DnsRecordType.SOA);
        if (rrset == null || rrset.size() != 1) {
            throw new IOException(origin + ": exactly 1 SOA must be specified");
        }
        Iterator it = rrset.rrs();
        SOA = (SOARecord) it.next();

        NS = oneRRset(originNode, DnsRecordType.NS);
        if (NS == null) {
            throw new IOException(origin + ": no NS set specified");
        }
    }

    private
    void maybeAddRecord(DnsRecord record) throws IOException {
        int rtype = record.getType();
        Name name = record.getName();

        if (rtype == DnsRecordType.SOA && !name.equals(origin)) {
            throw new IOException("SOA owner " + name + " does not match zone origin " + origin);
        }
        if (name.subdomain(origin)) {
            addRecord(record);
        }
    }

    /**
     * Adds a Record to the Zone
     *
     * @param r The record to be added
     *
     * @see DnsRecord
     */
    public
    void addRecord(DnsRecord r) {
        Name name = r.getName();
        int rtype = r.getRRsetType();
        synchronized (this) {
            RRset rrset = findRRset(name, rtype);
            if (rrset == null) {
                rrset = new RRset(r);
                addRRset(name, rrset);
            }
            else {
                rrset.addRR(r);
            }
        }
    }

    private synchronized
    RRset findRRset(Name name, int type) {
        Object types = exactName(name);
        if (types == null) {
            return null;
        }
        return oneRRset(types, type);
    }

    private synchronized
    Object exactName(Name name) {
        return data.get(name);
    }

    private synchronized
    RRset oneRRset(Object types, int type) {
        if (type == DnsRecordType.ANY) {
            throw new IllegalArgumentException("oneRRset(ANY)");
        }
        if (types instanceof List) {
            List list = (List) types;
            for (int i = 0; i < list.size(); i++) {
                RRset set = (RRset) list.get(i);
                if (set.getType() == type) {
                    return set;
                }
            }
        }
        else {
            RRset set = (RRset) types;
            if (set.getType() == type) {
                return set;
            }
        }
        return null;
    }

    private synchronized
    void addRRset(Name name, RRset rrset) {
        if (!hasWild && name.isWild()) {
            hasWild = true;
        }
        Object types = data.get(name);
        if (types == null) {
            data.put(name, rrset);
            return;
        }
        int rtype = rrset.getType();
        if (types instanceof List) {
            List list = (List) types;
            for (int i = 0; i < list.size(); i++) {
                RRset set = (RRset) list.get(i);
                if (set.getType() == rtype) {
                    list.set(i, rrset);
                    return;
                }
            }
            list.add(rrset);
        }
        else {
            RRset set = (RRset) types;
            if (set.getType() == rtype) {
                data.put(name, rrset);
            }
            else {
                LinkedList list = new LinkedList();
                list.add(set);
                list.add(rrset);
                data.put(name, list);
            }
        }
    }

    /**
     * Creates a Zone from an array of records.
     *
     * @param zone The name of the zone.
     * @param records The records to add to the zone.
     *
     * @see Master
     */
    public
    Zone(Name zone, DnsRecord[] records) throws IOException {
        data = new TreeMap();

        if (zone == null) {
            throw new IllegalArgumentException("no zone name specified");
        }
        origin = zone;
        for (int i = 0; i < records.length; i++) {
            maybeAddRecord(records[i]);
        }
        validate();
    }

    /**
     * Creates a Zone by doing the specified zone transfer.
     *
     * @param xfrin The incoming zone transfer to execute.
     *
     * @see ZoneTransferIn
     */
    public
    Zone(ZoneTransferIn xfrin) throws IOException, ZoneTransferException {
        fromXFR(xfrin);
    }

    private
    void fromXFR(ZoneTransferIn xfrin) throws IOException, ZoneTransferException {
        data = new TreeMap();

        origin = xfrin.getName();
        List records = xfrin.run();
        for (Iterator it = records.iterator(); it.hasNext(); ) {
            DnsRecord record = (DnsRecord) it.next();
            maybeAddRecord(record);
        }
        if (!xfrin.isAXFR()) {
            throw new IllegalArgumentException("zones can only be " + "created from AXFRs");
        }
        validate();
    }

    /**
     * Creates a Zone by performing a zone transfer to the specified host.
     *
     * @see ZoneTransferIn
     */
    public
    Zone(Name zone, int dclass, String remote) throws IOException, ZoneTransferException {
        ZoneTransferIn xfrin = ZoneTransferIn.newAXFR(zone, remote, null);
        xfrin.setDClass(dclass);
        fromXFR(xfrin);
    }

    /**
     * Returns the Zone's origin
     */
    public
    Name getOrigin() {
        return origin;
    }

    /**
     * Returns the Zone origin's NS records
     */
    public
    RRset getNS() {
        return NS;
    }

    /**
     * Returns the Zone's SOA record
     */
    public
    SOARecord getSOA() {
        return SOA;
    }

    /**
     * Returns the Zone's class
     */
    public
    int getDClass() {
        return dclass;
    }

    /**
     * Looks up Records in the Zone.  This follows CNAMEs and wildcards.
     *
     * @param name The name to look up
     * @param type The type to look up
     *
     * @return A SetResponse object
     *
     * @see SetResponse
     */
    public
    SetResponse findRecords(Name name, int type) {
        return lookup(name, type);
    }

    private synchronized
    SetResponse lookup(Name name, int type) {
        int labels;
        int olabels;
        int tlabels;
        RRset rrset;
        Name tname;
        Object types;
        SetResponse sr;

        if (!name.subdomain(origin)) {
            return SetResponse.ofType(SetResponse.NXDOMAIN);
        }

        labels = name.labels();
        olabels = origin.labels();

        for (tlabels = olabels; tlabels <= labels; tlabels++) {
            boolean isOrigin = (tlabels == olabels);
            boolean isExact = (tlabels == labels);

            if (isOrigin) {
                tname = origin;
            }
            else if (isExact) {
                tname = name;
            }
            else {
                tname = new Name(name, labels - tlabels);
            }

            types = exactName(tname);
            if (types == null) {
                continue;
            }

		/* If this is a delegation, return that. */
            if (!isOrigin) {
                RRset ns = oneRRset(types, DnsRecordType.NS);
                if (ns != null) {
                    return new SetResponse(SetResponse.DELEGATION, ns);
                }
            }

		/* If this is an ANY lookup, return everything. */
            if (isExact && type == DnsRecordType.ANY) {
                sr = new SetResponse(SetResponse.SUCCESSFUL);
                RRset[] sets = allRRsets(types);
                for (int i = 0; i < sets.length; i++) {
                    sr.addRRset(sets[i]);
                }
                return sr;
            }

		/*
         * If this is the name, look for the actual type or a CNAME.
		 * Otherwise, look for a DNAME.
		 */
            if (isExact) {
                rrset = oneRRset(types, type);
                if (rrset != null) {
                    sr = new SetResponse(SetResponse.SUCCESSFUL);
                    sr.addRRset(rrset);
                    return sr;
                }
                rrset = oneRRset(types, DnsRecordType.CNAME);
                if (rrset != null) {
                    return new SetResponse(SetResponse.CNAME, rrset);
                }
            }
            else {
                rrset = oneRRset(types, DnsRecordType.DNAME);
                if (rrset != null) {
                    return new SetResponse(SetResponse.DNAME, rrset);
                }
            }

		/* We found the name, but not the type. */
            if (isExact) {
                return SetResponse.ofType(SetResponse.NXRRSET);
            }
        }

        if (hasWild) {
            for (int i = 0; i < labels - olabels; i++) {
                tname = name.wild(i + 1);

                types = exactName(tname);
                if (types == null) {
                    continue;
                }

                rrset = oneRRset(types, type);
                if (rrset != null) {
                    sr = new SetResponse(SetResponse.SUCCESSFUL);
                    sr.addRRset(rrset);
                    return sr;
                }
            }
        }

        return SetResponse.ofType(SetResponse.NXDOMAIN);
    }

    private synchronized
    RRset[] allRRsets(Object types) {
        if (types instanceof List) {
            List typelist = (List) types;
            return (RRset[]) typelist.toArray(new RRset[typelist.size()]);
        }
        else {
            RRset set = (RRset) types;
            return new RRset[] {set};
        }
    }

    /**
     * Looks up Records in the zone, finding exact matches only.
     *
     * @param name The name to look up
     * @param type The type to look up
     *
     * @return The matching RRset
     *
     * @see RRset
     */
    public
    RRset findExactMatch(Name name, int type) {
        Object types = exactName(name);
        if (types == null) {
            return null;
        }
        return oneRRset(types, type);
    }

    /**
     * Adds an RRset to the Zone
     *
     * @param rrset The RRset to be added
     *
     * @see RRset
     */
    public
    void addRRset(RRset rrset) {
        Name name = rrset.getName();
        addRRset(name, rrset);
    }

    /**
     * Removes a record from the Zone
     *
     * @param r The record to be removed
     *
     * @see DnsRecord
     */
    public
    void removeRecord(DnsRecord r) {
        Name name = r.getName();
        int rtype = r.getRRsetType();
        synchronized (this) {
            RRset rrset = findRRset(name, rtype);
            if (rrset == null) {
                return;
            }
            if (rrset.size() == 1 && rrset.first()
                                          .equals(r)) {
                removeRRset(name, rtype);
            }
            else {
                rrset.deleteRR(r);
            }
        }
    }

    private synchronized
    void removeRRset(Name name, int type) {
        Object types = data.get(name);
        if (types == null) {
            return;
        }
        if (types instanceof List) {
            List list = (List) types;
            for (int i = 0; i < list.size(); i++) {
                RRset set = (RRset) list.get(i);
                if (set.getType() == type) {
                    list.remove(i);
                    if (list.size() == 0) {
                        data.remove(name);
                    }
                    return;
                }
            }
        }
        else {
            RRset set = (RRset) types;
            if (set.getType() != type) {
                return;
            }
            data.remove(name);
        }
    }

    /**
     * Returns an Iterator over the RRsets in the zone.
     */
    public
    Iterator iterator() {
        return new ZoneIterator(false);
    }

    /**
     * Returns an Iterator over the RRsets in the zone that can be used to
     * construct an AXFR response.  This is identical to {@link #iterator} except
     * that the SOA is returned at the end as well as the beginning.
     */
    public
    Iterator AXFR() {
        return new ZoneIterator(true);
    }

    /**
     * Returns the contents of the Zone as a string (in master file format).
     */
    public
    String toString() {
        return toMasterFile();
    }

    /**
     * Returns the contents of the Zone in master file format.
     */
    public synchronized
    String toMasterFile() {
        Iterator zentries = data.entrySet()
                                .iterator();
        StringBuilder sb = new StringBuilder();
        nodeToString(sb, originNode);
        while (zentries.hasNext()) {
            Map.Entry entry = (Map.Entry) zentries.next();
            if (!origin.equals(entry.getKey())) {
                nodeToString(sb, entry.getValue());
            }
        }
        return sb.toString();
    }

    private
    void nodeToString(StringBuilder sb, Object node) {
        RRset[] sets = allRRsets(node);
        for (int i = 0; i < sets.length; i++) {
            RRset rrset = sets[i];
            Iterator it = rrset.rrs();
            while (it.hasNext()) {
                sb.append(it.next())
                  .append(OS.LINE_SEPARATOR);
            }
            it = rrset.sigs();
            while (it.hasNext()) {
                sb.append(it.next())
                  .append(OS.LINE_SEPARATOR);
            }
        }
    }

}