import java.util.*;
import java.io.*;
import java.net.UnknownHostException;
import edu.mit.media.hive.support.Debug;

public class NTPSpider {
    Hashtable knownHosts = new Hashtable(800011);       // Host -> HostInfo

    // List of jobs we need to do, prioritized by stratum.
    Vector stratumJobs[] = new Vector[17];

    WorkerPool workerPool;

    BufferedWriter toExploreFile, exploredFile;
    
    /** Call this method to add a new host to the databases. **/
    public synchronized void noteHost(String hostName, int stratum) {
        Host host = null;
        try {
            host = new Host(hostName);
        } catch (UnknownHostException ex) {
            Debug.error(ex);
        }
        
        if (knownHosts.get(host) != null)         // we've seen it already.
            return;                               // that's OK.

        if (stratum > 16 || stratum < 0) {
            Debug.warning("Stratum was " + stratum + " setting to 16 " + host);
            stratum = 16;
        }

        // Add it to the host table.
        noteNewHost(host, stratum);
        // Add to the appropriate queue.
        stratumJobs[stratum].addElement(host);
        // Add to the file
        try {
            toExploreFile.write(host + " " + stratum);
            toExploreFile.newLine();
            toExploreFile.flush();
        } catch (IOException ex) {
            Debug.error(ex);
        }
    }

    public synchronized void noteDone(String host) {
        // Store info in the file
        try {
            exploredFile.write(host);
            exploredFile.newLine();
            exploredFile.flush();
        } catch (IOException ex) {
            Debug.error(ex);
        }
    }
    
    /** Get a stream to write output on for data on "host". **/
    public OutputStream getOutputStream(String host) {
        // Create a file for output, handle if it exists already.

        // Construct a filename in a subdir.
        int dot = host.indexOf('.');
        String filename;
        if (dot > 0)
            filename = host.substring(0, dot) + "/" + host;
        else
            filename = host;
        
        File file = new File(filename);
        FileOutputStream outFile = null;
        try {
            if (file.exists()) {
                Debug.warning("File " + filename + " exists. Appending.");
                outFile = new FileOutputStream(filename, true);
                outFile.write(12);                // stick ^L in there
            } else {
                outFile = new FileOutputStream(file);
            }
        } catch (IOException ex) {
            Debug.error(ex);
        }
        return outFile;
    }
    
    public NTPSpider(WorkerPool wp, String hostsToExplore, String hostsExplored) {
        // set up worker pool
        workerPool = wp;
        wp.createWorkers(this);

        // create our job lists
        for (int i = 0; i < stratumJobs.length; i++)
            stratumJobs[i] = new Vector();

        // initialize data structs for our state
        Reader r = null;
        try {
            r = new FileReader(hostsToExplore);
        } catch (Exception ex) {
            Debug.error(ex);
            System.exit(1);
        }
        readHostsToExplore(r);

        try {
            r = new FileReader(hostsExplored);
        } catch (Exception ex) {
            Debug.error(ex);
        }
        readHostsExplored(r);

        // build a list of jobs to do
        buildJobsList();

        // open up the files for keeping state
        FileWriter fw;
        try {
            fw = new FileWriter(hostsToExplore, true);
            toExploreFile = new BufferedWriter(fw);
            fw = new FileWriter(hostsExplored, true);
            exploredFile = new BufferedWriter(fw);
        } catch (IOException ex) {
            Debug.error(ex);
            System.exit(1);
        }
    }

    /** Read in a database of all the hosts we know about.
     ** Format is: "ip.ad.dr.ess stratum"
     **/
    private void readHostsToExplore(Reader r) {
        BufferedReader br = new BufferedReader(r);
        String s;
        long i = 0;
        try {
            while ((s = br.readLine()) != null) {
                int space = s.indexOf(' ');
                String hostName = s.substring(0, space);
                Host host = null;
                try {
                    host = new Host(hostName);
                } catch (UnknownHostException ex) {
                    Debug.error(ex);
                }
                
                String stratum = s.substring(space + 1);
                int sv = Integer.parseInt(stratum);
                noteNewHost(host, sv);
                if (++i % 10000 == 0)
                    Debug.notice("Read "+ i);
            }
        } catch (Exception ex) {
            Debug.error(ex);
        }
    }

    /** Read in a database of all the hosts we've explored. Just a list
     ** of IP addresses, one per line.
     **/
    private void readHostsExplored(Reader r) {
        BufferedReader br = new BufferedReader(r);
        String s;
        try {
            while ((s = br.readLine()) != null) {
                Host h = null;
                try {
                    h = new Host(s);
                } catch (UnknownHostException ex) {
                    Debug.error(ex);
                }
                noteHostExplored(h);
            }
        } catch (Exception ex) {
            Debug.error(ex);
        }
    }

    /** Startup method - build a list of jobs. Look through all hosts,
     ** add those that haven't been explored.
     **/
    public void buildJobsList() {
        // Walk the host info, add in the unexplored.
        Enumeration e = knownHosts.keys();
        while (e.hasMoreElements()) {
            Host host = (Host)e.nextElement();
            HostInfo hi = (HostInfo)knownHosts.get(host);
            if (!hi.explored)
                stratumJobs[hi.stratum].addElement(host);
        }

        Debug.notice("Jobs: " + getJobReport());
    }

    /** Make a string showing how many jobs there are. **/
    private String getJobReport() {
        StringBuffer sb = new StringBuffer();
        for (int i = 0; i < stratumJobs.length; i++) {
            sb.append(" " + stratumJobs[i].size());
        }
        return sb.toString();
    }
    
    /** Data structure method - note that we've explored a host. **/
    private synchronized void noteHostExplored(Host h) {
        HostInfo hi = (HostInfo)knownHosts.get(h);
        if (hi == null) {
            Debug.error("Told we explored " + hi + " but didn't know it!");
        } else {
            hi.explored = true;
        }
    }

    /** Data structure method - note that a new host has arrived. **/
    private synchronized void noteNewHost(Host h, int sv) {
        if (knownHosts.get(h) != null)
            Debug.warning("Warning, host " + h + " was already known. Overwriting.");
        HostInfo hi = new HostInfo(sv, false);
        knownHosts.put(h, hi);
    }

    /** Find the next job on the queue. **/
    private synchronized Host takeNextJob() {
        for (int i = 0; i < stratumJobs.length; i++) {
            int size = stratumJobs[i].size();
            // Is there a job here? Great, grab it.
            if (size > 0) {
                // Take it from the job list
                Host job = (Host)stratumJobs[i].elementAt(0);
                stratumJobs[i].removeElementAt(0);

                // Mark it as explored
                noteHostExplored(job);

                // Return it.
                return job;
            }
        }
        
        // oops, no jobs!
        return null;
    }
    
    /** Run the spider! **/
    public void go() {
        // Fork a thread to monitor queues for us
        new Thread(new Runnable() {
                public void run() {
                    while (true) {
                        try {
                            Thread.sleep(30000);
                        } catch (InterruptedException ex) { }
                        Debug.notice("Jobs: " + getJobReport());
                    }
                }
            }).start();
                
        while (true) {
            Host job = takeNextJob();
            if (job == null) {
                Debug.notice("No jobs! Sleeping 5 seconds.");
                try {
                    Thread.sleep(5000);
                } catch (InterruptedException ex) { }
            } else {
                // Debug.notice("Scheduling job " + job);
                workerPool.doJob(job.toString());
            }
        }
    }
    
    public static void main(String args[]) {
        int poolSize = 1;
        if (args.length != 1) {
            System.err.println("Usage: NTPSpider poolsize");
        }
        try {
            poolSize = Integer.parseInt(args[0]);
        } catch (Exception ex) {
            Debug.error(ex);
        }
        
        WorkerPool wp = new WorkerPool(poolSize);
        NTPSpider spider = new NTPSpider(wp, "TOEXPLORE", "EXPLORED");
        spider.go();
    }    
}


class HostInfo {
    public byte stratum;                           // estimated stratum
    public boolean explored;                      // have we tested this?

    public HostInfo(int s, boolean t) {
        if (s < 0 || s > 16) {
            Debug.error("Changing out of range stratum " + s);
            s = 15;
        }
        stratum = (byte)s;
        explored = t;
    }
}