import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.SocketException;
import java.util.LinkedList;
import java.util.List;
import java.util.Random;

/**
 * Class for master process in simple master/worker example using socket
 * communication.  In this example, the master process generates a fixed number
 * of (fake) tasks with randomly-generated times and waits for workers to
 * connect and process them. 
 */
public class SocketsMaster {

    /**
     * Constructs object with specified parameters.
     */
    public SocketsMaster(int port, int tasks, int maxTaskTime) 
        throws IOException {
        // generate tasks and put on queue
        Random randGen = new Random();
        for (int i = 0; i < tasks; ++i) {
            // generate integer in range 1 .. maxTaskTime
            taskQueue.add(new Task(randGen.nextInt(maxTaskTime-1)+1));
        }
        // set up server socket
        serverSocket = new ServerSocket(port);
    }

    /**
     * Starts master process (i.e., enables it to accept connections from
     * worker processes).  As workers connect, a new thread is started for
     * each worker.
     */
    public void start() {

        // put this in a separate thread so we can wait for the queue to
        // to be empty in another thread (and close the server socket when
        // it is) -- otherwise we would block at the call to accept()

        Thread acceptConnections = new Thread(new Runnable() {

            public void run() {
                try {
                    // runs until another thread closes serverSocket,
                    // generating a SocketException
                    while (true) {
                        Socket s = serverSocket.accept();
                        try {
                            WorkerTalker w = new WorkerTalker(s);
                            new Thread(w).start();
                        }
                        catch (IOException e) {
                            System.err.println(
                                "Error starting worker thread for client on " +
                                clientAddress(s) + ":  " + e);
                            s.close();
                        }
                    }
                }
                catch (SocketException e) {
                    if (!done) {
                        System.err.println("Server error:  " + e);
                    }
                    else {
                        System.out.println("Server shut down");
                    }
                }
                catch (IOException e) {
                    System.err.println("Server error:  " + e);
                }
            }
        });
        acceptConnections.start();
    }

    /**
     * Waits until the queue is empty and all workers have reported
     * results.
     */
    public void waitForComplete() {
        // wait for all tasks to be complete
        synchronized (taskQueue) {
            while (!taskQueue.isEmpty()) {
                try {
                    taskQueue.wait();
                }
                catch (InterruptedException e) {
                    System.err.println("should not happen?:  " + e);
                }
            }
        }
        // wait for all workers to report results
        synchronized (workers) {
            while (!workers.isEmpty()) {
                try {
                    workers.wait();
                }
                catch (InterruptedException e) {
                    System.err.println("should not happen?:  " + e);
                }
            }
        }
    }

    /**
     * Stops master process (i.e., closes server socket).
     */
    public void stop() throws IOException {
        done = true;
        serverSocket.close();
    }

    /**
     * Gets task from queue (thread-safe).  Notifies waiting process if
     * queue is empty.
     */
    public Task getTask() {
        synchronized (taskQueue) {
            if (taskQueue.isEmpty()) {
                taskQueue.notifyAll();
                return null;
            }
            else {
                return taskQueue.remove(0);
            }
        }
    }

    /**
     * Adds worker (thread-safe).
     */
    private void addWorker(WorkerTalker w) {
        synchronized (workers) {
            workers.add(w);
        }
    }

    /**
     * Removes worker (thread-safe).  Notifies waiting process if no workers.
     */
    private void removeWorker(WorkerTalker w) {
        synchronized (workers) {
            workers.remove(w);
            if (workers.isEmpty())
                workers.notifyAll();
        }
    }

    /**
     * Prints worker results.
     */
    private synchronized void reportResults(WorkerTalker w) {
        System.out.println("Worker at " + w.hostName() + 
                " completed " + w.numTasks() + " tasks" +
                ", total time = " + w.totalTime());
    }

    /**
     * Gets hostname for socket.
     */
    private static String clientAddress(Socket s) {
        return s.getInetAddress().getHostName();
    }

    /**
     * Main program -- constructs an object and starts it up. 
     * @param args command-line arguments -- port number, total tasks, 
     *   maximum task time
     */
    public static void main(String[] args) {
        String usageMsg = "Arguments are portNum, totalTasks, maxTime";
        if (args.length < 3) {
            System.err.println(usageMsg);
            System.exit(1);
        }
        try {
            int port = Integer.parseInt(args[0]);
            int tasks = Integer.parseInt(args[1]);
            int maxtime = Integer.parseInt(args[2]);

            SocketsMaster taskMaster = new SocketsMaster(port, tasks, maxtime);
            taskMaster.start();
            System.out.println("Master waiting for input");
            taskMaster.waitForComplete();
            System.out.println("Master done");
            taskMaster.stop();
        }
        catch (NumberFormatException e) {
            System.err.println(usageMsg);
        }
        catch (IOException e) {
            System.err.println("I/O error in server:"  + e);
        }
    }

    // ---- variables ----

    private List<Task> taskQueue = new LinkedList<Task>();
    private List<WorkerTalker> workers = new LinkedList<WorkerTalker>();
    private ServerSocket serverSocket;
    private volatile boolean done = false;

    // ---- classes ----

    /**
     * Class for threads to communicate with workers.
     */
    private class WorkerTalker implements Runnable {

        /**
         * Constructs object.
         */
        public WorkerTalker(Socket socket) throws IOException {

            this.socket = socket;
            System.out.println("New worker for " + hostName());
            output = new ObjectOutputStream(socket.getOutputStream());
            input = new ObjectInputStream(socket.getInputStream());
            addWorker(this);
        }

        /**
         * Contains code to be executed by thread.
         */
        public void run() {

            try {
                Task t;
                TaskResult response;
                // get and process tasks until no more in queue
                while ((t = getTask()) != null) {
                    // send task to worker process, receive result
                    System.out.println("Sending to worker at " + hostName() + 
                            ":  task with time " + t.time);
                    output.writeObject(t);
                    response = (TaskResult) input.readObject();
                    ++tasks;
                    totalTime += response.actualTime;
                }
                // shut things down
                socket.shutdownOutput();
                output.close();
                input.close();
                socket.close();
            }
            catch (ClassNotFoundException e) { 
                System.err.println("Error in client for " + hostName() + 
                        ":  " + e);
            }
            catch (IOException e) { 
                System.err.println("Error in client for " + hostName() +
                        ":  " + e);
            }
            reportResults(this);
            removeWorker(this);
        }

        /**
         * Gets hostname for worker.
         */
        public String hostName() {
            return clientAddress(socket);
        }

        /**
         * Gets total number of tasks completed by worker.
         */
        public int numTasks() {
            return tasks;
        }

        /**
         * Gets total time used by worker.
         */
        public int totalTime() {
            return totalTime;
        }

        // ---- variables ----

        private Socket socket;
        private ObjectInputStream input;
        private ObjectOutputStream output;
        private int tasks = 0;
        private int totalTime = 0;
    }
}