package jxm.JWS;

import jxm.JWS.EX.InvalidDataException;
import jxm.JWS.FR.CloseFrame;
import jxm.JWS.FR.Framedata;
import jxm.JWS.HS.ClientHandshake;
import jxm.JWS.HS.Handshakedata;
import jxm.JWS.DF.Draft;
import jxm.JWS.HS.ServerHandshakeBuilder;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.UnknownHostException;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.util.*;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;


/**
 * <tt>WebSocketServer</tt> is an abstract class that only takes care of the
 * HTTP handshake portion of WebSockets. It's up to a subclass to add
 * functionality/purpose to the server.
 */
abstract class WebSocketServer extends WebSocketAdapter implements Runnable {

    public static int DECODERS = Runtime.getRuntime().availableProcessors();

    /**
     * Holds the list of active WebSocket connections. "Active" means WebSocket
     * handshake is complete and socket can be written to, or read from.
     */
    private final Collection<WebSocket> connections;
    /**
     * The port number that this WebSocket server should listen on. Default is
     * WebSocket.DEFAULT_PORT.
     */
    private final InetSocketAddress address;
    /**
     * The socket channel for this WebSocket server.
     */
    private ServerSocketChannel server;
    /**
     * The 'Selector' used to get event keys from the underlying socket.
     */
    private Selector selector;
    /**
     * The Draft of the WebSocket protocol the Server is adhering to.
     */
    private List<Draft> drafts;

    private Thread selectorthread;

    private volatile AtomicBoolean isclosed = new AtomicBoolean(false);

    private List<WebSocketWorker> decoders;

    private List<WebSocketImpl> iqueue;
    private BlockingQueue<ByteBuffer> buffers;
    private int queueinvokes = 0;
    private AtomicInteger queuesize = new AtomicInteger(0);

    private WebSocketServerFactory wsf = new DefaultWebSocketServerFactory();

    /**
     * Creates a WebSocketServer that will attempt to
     * listen on port <var>WebSocket.DEFAULT_PORT</var>.
     *
     * @see #WebSocketServer(InetSocketAddress, int, List, Collection) more details here
     */
    public WebSocketServer() throws UnknownHostException {
        this(new InetSocketAddress(WebSocket.DEFAULT_PORT), DECODERS, null);
    }

    /**
     * Creates a WebSocketServer that will attempt to bind/listen on the given <var>address</var>.
     *
     * @see #WebSocketServer(InetSocketAddress, int, List, Collection) more details here
     */
    public WebSocketServer(InetSocketAddress address) {
        this(address, DECODERS, null);
    }

    /**
     * @see #WebSocketServer(InetSocketAddress, int, List, Collection) more details here
     */
    public WebSocketServer(InetSocketAddress address, int decoders) {
        this(address, decoders, null);
    }

    /**
     * @see #WebSocketServer(InetSocketAddress, int, List, Collection) more details here
     */
    public WebSocketServer(InetSocketAddress address, List<Draft> drafts) {
        this(address, DECODERS, drafts);
    }

    /**
     * @see #WebSocketServer(InetSocketAddress, int, List, Collection) more details here
     */
    public WebSocketServer(InetSocketAddress address, int decodercount, List<Draft> drafts) {
        this(address, decodercount, drafts, new HashSet<WebSocket>());
    }

    /**
     * Creates a WebSocketServer that will attempt to bind/listen on the given <var>address</var>,
     * and comply with <tt>Draft</tt> version <var>draft</var>.
     *
     * @param address              The address (host:port) this server should listen on.
     * @param decodercount         The number of {@link WebSocketWorker}s that will be used to process the incoming network data. By default this will be <code>Runtime.getRuntime().availableProcessors()</code>
     * @param drafts               The versions of the WebSocket protocol that this server
     *                             instance should comply to. Clients that use an other protocol version will be rejected.
     * @param connectionscontainer Allows to specify a collection that will be used to store the websockets in. <br>
     *                             If you plan to often iterate through the currently connected websockets you may want to use a collection that does not require synchronization like a {@link CopyOnWriteArraySet}. In that case make sure that you overload {@link #removeConnection(WebSocket)} and {@link #addConnection(WebSocket)}.<br>
     *                             By default a {@link HashSet} will be used.
     * @see #removeConnection(WebSocket) for more control over syncronized operation
     * @see <a href="https://github.com/TooTallNate/Java-WebSocket/wiki/Drafts" > more about drafts
     */
    public WebSocketServer(InetSocketAddress address, int decodercount, List<Draft> drafts, Collection<WebSocket> connectionscontainer) {
        if (address == null || decodercount < 1 || connectionscontainer == null) {
            throw new IllegalArgumentException("address and connectionscontainer must not be null and you need at least 1 decoder");
        }

        if (drafts == null)
            this.drafts = Collections.emptyList();
        else
            this.drafts = drafts;

        this.address = address;
        this.connections = connectionscontainer;

        iqueue = new LinkedList<WebSocketImpl>();

        decoders = new ArrayList<WebSocketWorker>(decodercount);
        buffers = new LinkedBlockingQueue<ByteBuffer>();
        for (int i = 0; i < decodercount; i++) {
            WebSocketWorker ex = new WebSocketWorker();
            decoders.add(ex);
            ex.start();
        }
    }

    /**
     * Starts the server selectorthread that binds to the currently set port number and
     * listeners for WebSocket connection requests. Creates a fixed thread pool with the size {@link WebSocketServer#DECODERS}<br>
     * May only be called once.
     * <p/>
     * Alternatively you can call {@link WebSocketServer#run()} directly.
     *
     * @throws IllegalStateException
     */
    public void start() {
        if (selectorthread != null)
            throw new IllegalStateException(getClass().getName() + " can only be started once.");
        new Thread(this).start();
        ;
    }

    /**
     * Closes all connected clients sockets, then closes the underlying
     * ServerSocketChannel, effectively killing the server socket selectorthread,
     * freeing the port the server was bound to and stops all internal workerthreads.
     * <p/>
     * If this method is called before the server is started it will never start.
     *
     * @param timeout Specifies how many milliseconds the overall close handshaking may take altogether before the connections are closed without proper close handshaking.<br>
     * @throws IOException          When {@link ServerSocketChannel}.close throws an IOException
     * @throws InterruptedException
     */
    public void stop(int timeout) throws InterruptedException {
        if (!isclosed.compareAndSet(false, true)) { // this also makes sure that no further connections will be added to this.connections
            return;
        }

        List<WebSocket> socketsToClose = null;

        // copy the connections in a list (prevent callback deadlocks)
        synchronized (connections) {
            socketsToClose = new ArrayList<WebSocket>(connections);
        }

        for (WebSocket ws : socketsToClose) {
            ws.close(CloseFrame.GOING_AWAY);
        }

        synchronized (this) {
            if (selectorthread != null) {
                if (Thread.currentThread() != selectorthread) {

                }
                if (selectorthread != Thread.currentThread()) {
                    if (socketsToClose.size() > 0)
                        selectorthread.join(timeout);// isclosed will tell the selectorthread to go down after the last connection was closed
                    selectorthread.interrupt();// in case the selectorthread did not terminate in time we send the interrupt
                    selectorthread.join();
                }
            }
        }
    }

    public void stop() throws IOException, InterruptedException {
        stop(0);
    }

    /**
     * Returns a WebSocket[] of currently connected clients.
     * Its iterators will be failfast and its not judicious
     * to modify it.
     *
     * @return The currently connected clients.
     */
    public Collection<WebSocket> connections() {
        return this.connections;
    }

    public InetSocketAddress getAddress() {
        return this.address;
    }

    /**
     * Gets the port number that this server listens on.
     *
     * @return The port number.
     */
    public int getPort() {
        int port = getAddress().getPort();
        if (port == 0 && server != null) {
            port = server.socket().getLocalPort();
        }
        return port;
    }

    public List<Draft> getDraft() {
        return Collections.unmodifiableList(drafts);
    }

    // Runnable IMPLEMENTATION /////////////////////////////////////////////////
    public void run() {
        synchronized (this) {
            if (selectorthread != null)
                throw new IllegalStateException(getClass().getName() + " can only be started once.");
            selectorthread = Thread.currentThread();
            if (isclosed.get()) {
                return;
            }
        }
        selectorthread.setName("WebsocketSelector" + selectorthread.getId());
        try {
            server = ServerSocketChannel.open();
            server.configureBlocking(false);
            ServerSocket socket = server.socket();
            socket.setReceiveBufferSize(WebSocketImpl.RCVBUF);
            socket.bind(address);
            selector = Selector.open();
            server.register(selector, server.validOps());
        } catch (IOException ex) {
            handleFatal(null, ex);
            return;
        }
        try {
            while (!selectorthread.isInterrupted()) {
                SelectionKey key = null;
                WebSocketImpl conn = null;
                try {
                    selector.select();
                    Set<SelectionKey> keys = selector.selectedKeys();
                    Iterator<SelectionKey> i = keys.iterator();

                    while (i.hasNext()) {
                        key = i.next();

                        if (!key.isValid()) {
                            // Object o = key.attachment();
                            continue;
                        }

                        if (key.isAcceptable()) {
                            if (!onConnect(key)) {
                                key.cancel();
                                continue;
                            }

                            SocketChannel channel = server.accept();
                            channel.configureBlocking(false);
                            WebSocketImpl w = wsf.createWebSocket(this, drafts, channel.socket());
                            w.key = channel.register(selector, SelectionKey.OP_READ, w);
                            w.channel = wsf.wrapChannel(channel, w.key);
                            i.remove();
                            allocateBuffers(w);
                            continue;
                        }

                        if (key.isReadable()) {
                            conn = (WebSocketImpl) key.attachment();
                            ByteBuffer buf = takeBuffer();
                            try {
                                if (SocketChannelIOHelper.read(buf, conn, conn.channel)) {
                                    if (buf.hasRemaining()) {
                                        conn.inQueue.put(buf);
                                        queue(conn);
                                        i.remove();
                                        if (conn.channel instanceof WrappedByteChannel) {
                                            if (((WrappedByteChannel) conn.channel).isNeedRead()) {
                                                iqueue.add(conn);
                                            }
                                        }
                                    } else
                                        pushBuffer(buf);
                                } else {
                                    pushBuffer(buf);
                                }
                            } catch (IOException e) {
                                pushBuffer(buf);
                                throw e;
                            }
                        }
                        if (key.isWritable()) {
                            conn = (WebSocketImpl) key.attachment();
                            if (SocketChannelIOHelper.batch(conn, conn.channel)) {
                                if (key.isValid())
                                    key.interestOps(SelectionKey.OP_READ);
                            }
                        }
                    }
                    while (!iqueue.isEmpty()) {
                        conn = iqueue.remove(0);
                        WrappedByteChannel c = ((WrappedByteChannel) conn.channel);
                        ByteBuffer buf = takeBuffer();
                        try {
                            if (SocketChannelIOHelper.readMore(buf, conn, c))
                                iqueue.add(conn);
                            if (buf.hasRemaining()) {
                                conn.inQueue.put(buf);
                                queue(conn);
                            } else {
                                pushBuffer(buf);
                            }
                        } catch (IOException e) {
                            pushBuffer(buf);
                            throw e;
                        }

                    }
                } catch (CancelledKeyException e) {
                    // an other thread may cancel the key
                } catch (ClosedByInterruptException e) {
                    return; // do the same stuff as when InterruptedException is thrown
                } catch (IOException ex) {
                    if (key != null)
                        key.cancel();
                    handleIOException(key, conn, ex);
                } catch (InterruptedException e) {
                    return;// FIXME controlled shutdown (e.g. take care of buffermanagement)
                }
            }

        } catch (RuntimeException e) {
            // should hopefully never occur
            handleFatal(null, e);
        } finally {
            if (decoders != null) {
                for (WebSocketWorker w : decoders) {
                    w.interrupt();
                }
            }
            if (server != null) {
                try {
                    server.close();
                } catch (IOException e) {
                    onError(null, e);
                }
            }
        }
    }

    protected void allocateBuffers(WebSocket c) throws InterruptedException {
        if (queuesize.get() >= 2 * decoders.size() + 1) {
            return;
        }
        queuesize.incrementAndGet();
        buffers.put(createBuffer());
    }

    protected void releaseBuffers(WebSocket c) throws InterruptedException {
        // queuesize.decrementAndGet();
        // takeBuffer();
    }

    public ByteBuffer createBuffer() {
        return ByteBuffer.allocate(WebSocketImpl.RCVBUF);
    }

    private void queue(WebSocketImpl ws) throws InterruptedException {
        if (ws.workerThread == null) {
            ws.workerThread = decoders.get(queueinvokes % decoders.size());
            queueinvokes++;
        }
        ws.workerThread.put(ws);
    }

    private ByteBuffer takeBuffer() throws InterruptedException {
        return buffers.take();
    }

    private void pushBuffer(ByteBuffer buf) throws InterruptedException {
        if (buffers.size() > queuesize.intValue())
            return;
        buffers.put(buf);
    }

    private void handleIOException(SelectionKey key, WebSocket conn, IOException ex) {
        // onWebsocketError( conn, ex );// conn may be null here
        if (conn != null) {
            conn.closeConnection(CloseFrame.ABNORMAL_CLOSE, ex.getMessage());
        } else if (key != null) {
            SelectableChannel channel = key.channel();
            if (channel != null && channel.isOpen()) { // this could be the case if the IOException ex is a SSLException
                try {
                    channel.close();
                } catch (IOException e) {
                    // there is nothing that must be done here
                }
                if (WebSocketImpl.DEBUG)
                    System.out.println("Connection closed because of" + ex);
            }
        }
    }

    private void handleFatal(WebSocket conn, Exception e) {
        onError(conn, e);
        try {
            stop();
        } catch (IOException e1) {
            onError(null, e1);
        } catch (InterruptedException e1) {
            Thread.currentThread().interrupt();
            onError(null, e1);
        }
    }

    /**
     * Gets the XML string that should be returned if a client requests a Flash
     * security policy.
     * <p/>
     * The default implementation allows access from all remote domains, but
     * only on the port that this WebSocketServer is listening on.
     * <p/>
     * This is specifically implemented for gitime's WebSocket client for Flash:
     * http://github.com/gimite/web-socket-js
     *
     * @return An XML String that comforms to Flash's security policy. You MUST
     *         not include the null char at the end, it is appended automatically.
     */
    protected String getFlashSecurityPolicy() {
        return "<cross-domain-policy><allow-access-from domain=\"*\" to-ports=\"" + getPort() + "\" /></cross-domain-policy>";
    }

    @Override
    public final void onWebsocketMessage(WebSocket conn, String message) {
        onMessage(conn, message);
    }

    @Override
    @Deprecated
    public/*final*/void onWebsocketMessageFragment(WebSocket conn, Framedata frame) {// onFragment should be overloaded instead
        onFragment(conn, frame);
    }

    @Override
    public final void onWebsocketMessage(WebSocket conn, ByteBuffer blob) {
        onMessage(conn, blob);
    }

    @Override
    public final void onWebsocketOpen(WebSocket conn, Handshakedata handshake) {
        if (addConnection(conn)) {
            onOpen(conn, (ClientHandshake) handshake);
        }
    }

    @Override
    public final void onWebsocketClose(WebSocket conn, int code, String reason, boolean remote) {
        selector.wakeup();
        try {
            if (removeConnection(conn)) {
                onClose(conn, code, reason, remote);
            }
        } finally {
            try {
                releaseBuffers(conn);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }

    }

    /**
     * This method performs remove operations on the connection and therefore also gives control over whether the operation shall be synchronized
     * <p/>
     * {@link #WebSocketServer(InetSocketAddress, int, List, Collection)} allows to specify a collection which will be used to store current connections in.<br>
     * Depending on the type on the connection, modifications of that collection may have to be synchronized.
     */
    protected boolean removeConnection(WebSocket ws) {
        boolean removed;
        synchronized (connections) {
            removed = this.connections.remove(ws);
            assert (removed);
        }
        if (isclosed.get() && connections.size() == 0) {
            selectorthread.interrupt();
        }
        return removed;
    }

    @Override
    public ServerHandshakeBuilder onWebsocketHandshakeReceivedAsServer(WebSocket conn, Draft draft, ClientHandshake request) throws InvalidDataException {
        return super.onWebsocketHandshakeReceivedAsServer(conn, draft, request);
    }

    /**
     * @see #removeConnection(WebSocket)
     */
    protected boolean addConnection(WebSocket ws) {
        if (!isclosed.get()) {
            synchronized (connections) {
                boolean succ = this.connections.add(ws);
                assert (succ);
                return succ;
            }
        } else {
            // This case will happen when a new connection gets ready while the server is already stopping.
            ws.close(CloseFrame.GOING_AWAY);
            return true;// for consistency sake we will make sure that both onOpen will be called
        }
    }

    /**
     * @param conn may be null if the error does not belong to a single connection
     */
    @Override
    public final void onWebsocketError(WebSocket conn, Exception ex) {
        onError(conn, ex);
    }

    @Override
    public final void onWriteDemand(WebSocket w) {
        WebSocketImpl conn = (WebSocketImpl) w;
        try {
            conn.key.interestOps(SelectionKey.OP_READ | SelectionKey.OP_WRITE);
        } catch (CancelledKeyException e) {
            // the thread which cancels key is responsible for possible cleanup
            conn.outQueue.clear();
        }
        selector.wakeup();
    }

    @Override
    public void onWebsocketCloseInitiated(WebSocket conn, int code, String reason) {
        onCloseInitiated(conn, code, reason);
    }

    @Override
    public void onWebsocketClosing(WebSocket conn, int code, String reason, boolean remote) {
        onClosing(conn, code, reason, remote);

    }

    public void onCloseInitiated(WebSocket conn, int code, String reason) {
    }

    public void onClosing(WebSocket conn, int code, String reason, boolean remote) {

    }

    public final void setWebSocketFactory(WebSocketServerFactory wsf) {
        this.wsf = wsf;
    }

    public final WebSocketFactory getWebSocketFactory() {
        return wsf;
    }

    /**
     * Returns whether a new connection shall be accepted or not.<br>
     * Therefore method is well suited to implement some kind of connection limitation.<br>
     *
     * @see {@link #onOpen(WebSocket, ClientHandshake)}, {@link #onWebsocketHandshakeReceivedAsServer(WebSocket, Draft, ClientHandshake)}
     */
    protected boolean onConnect(SelectionKey key) {
        return true;
    }

    private Socket getSocket(WebSocket conn) {
        WebSocketImpl impl = (WebSocketImpl) conn;
        return ((SocketChannel) impl.key.channel()).socket();
    }

    @Override
    public InetSocketAddress getLocalSocketAddress(WebSocket conn) {
        return (InetSocketAddress) getSocket(conn).getLocalSocketAddress();
    }

    @Override
    public InetSocketAddress getRemoteSocketAddress(WebSocket conn) {
        return (InetSocketAddress) getSocket(conn).getRemoteSocketAddress();
    }

    /**
     * Called after an opening handshake has been performed and the given websocket is ready to be written on.
     */
    public abstract void onOpen(WebSocket conn, ClientHandshake handshake);

    /**
     * Called after the websocket connection has been closed.
     *
     * @param code   The codes can be looked up here: {@link CloseFrame}
     * @param reason Additional information string
     * @param remote Returns whether or not the closing of the connection was initiated by the remote host.
     */
    public abstract void onClose(WebSocket conn, int code, String reason, boolean remote);

    /**
     * Callback for string messages received from the remote host
     *
     * @see #onMessage(WebSocket, ByteBuffer)
     */
    public abstract void onMessage(WebSocket conn, String message);

    /**
     * Called when errors occurs. If an error causes the websocket connection to fail {@link #onClose(WebSocket, int, String, boolean)} will be called additionally.<br>
     * This method will be called primarily because of IO or protocol errors.<br>
     * If the given exception is an RuntimeException that probably means that you encountered a bug.<br>
     *
     * @param con Can be null if there error does not belong to one specific websocket. For example if the servers port could not be bound.
     */
    public abstract void onError(WebSocket conn, Exception ex);

    /**
     * Callback for binary messages received from the remote host
     *
     * @see #onMessage(WebSocket, String)
     */
    public void onMessage(WebSocket conn, ByteBuffer message) {
    }

    /**
     * @see WebSocket#sendFragmentedFrame(jxm._.JWS.FR.Framedata.Opcode, ByteBuffer, boolean)
     */
    public void onFragment(WebSocket conn, Framedata fragment) {
    }

    public class WebSocketWorker extends Thread {

        private BlockingQueue<WebSocketImpl> iqueue;

        public WebSocketWorker() {
            iqueue = new LinkedBlockingQueue<WebSocketImpl>();
            setName("WebSocketWorker-" + getId());
            setUncaughtExceptionHandler(new UncaughtExceptionHandler() {
                @Override
                public void uncaughtException(Thread t, Throwable e) {
                    getDefaultUncaughtExceptionHandler().uncaughtException(t, e);
                }
            });
        }

        public void put(WebSocketImpl ws) throws InterruptedException {
            iqueue.put(ws);
        }

        @Override
        public void run() {
            WebSocketImpl ws = null;
            try {
                while (true) {
                    ByteBuffer buf = null;
                    ws = iqueue.take();
                    buf = ws.inQueue.poll();
                    assert (buf != null);
                    try {
                        ws.decode(buf);
                    } finally {
                        pushBuffer(buf);
                    }
                }
            } catch (InterruptedException e) {
            } catch (RuntimeException e) {
                handleFatal(ws, e);
            }
        }
    }

    public interface WebSocketServerFactory extends WebSocketFactory {
        @Override
        public WebSocketImpl createWebSocket(WebSocketAdapter a, Draft d, Socket s);

        public WebSocketImpl createWebSocket(WebSocketAdapter a, List<Draft> drafts, Socket s);

        /**
         * Allows to wrap the Socketchannel( key.channel() ) to insert a protocol layer( like ssl or proxy authentication) beyond the ws layer.
         *
         * @param key a SelectionKey of an open SocketChannel.
         * @return The channel on which the read and write operations will be performed.<br>
         */
        public ByteChannel wrapChannel(SocketChannel channel, SelectionKey key) throws IOException;
    }
}