/*

 * Copyright (c) 2015, 2017, Oracle and/or its affiliates. All rights reserved.

 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.

 *

 * This code is free software; you can redistribute it and/or modify it

 * under the terms of the GNU General Public License version 2 only, as

 * published by the Free Software Foundation.

 *

 * This code is distributed in the hope that it will be useful, but WITHOUT

 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or

 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License

 * version 2 for more details (a copy is included in the LICENSE file that

 * accompanied this code).

 *

 * You should have received a copy of the GNU General Public License version

 * 2 along with this work; if not, write to the Free Software Foundation,

 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.

 *

 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA

 * or visit www.oracle.com if you need additional information or have any

 * questions.

 */

import java.io.IOException;

import java.net.*;

import java.util.*;

import java.util.concurrent.ExecutorService;

import java.util.concurrent.Executors;

import java.util.concurrent.ThreadFactory;

import java.util.concurrent.atomic.AtomicReference;

import java.util.function.Consumer;

import javax.net.ServerSocketFactory;

import javax.net.ssl.SSLContext;

import javax.net.ssl.SSLParameters;

import javax.net.ssl.SSLServerSocket;

import javax.net.ssl.SSLServerSocketFactory;

import javax.net.ssl.SNIServerName;

/**

 * Waits for incoming TCP connections from a client and establishes

 * a HTTP2 connection. Two threads are created per connection. One for reading

 * and one for writing. Incoming requests are dispatched to the supplied

 * Http2Handler on additional threads. All threads

 * obtained from the supplied ExecutorService.

 */

public class Http2TestServer implements AutoCloseable {

    final ServerSocket server;

    volatile boolean secure;

    final ExecutorService exec;

    volatile boolean stopping = false;

    final Map<String,Http2Handler> handlers;

    final SSLContext sslContext;

    final String serverName;

    final HashMap<InetSocketAddress,Http2TestServerConnection> connections;

    private static ThreadFactory defaultThreadFac =

        (Runnable r) -> {

            Thread t = new Thread(r);

            t.setName("Test-server-pool");

            return t;

        };

    private static ExecutorService getDefaultExecutor() {

        return Executors.newCachedThreadPool(defaultThreadFac);

    }

    public Http2TestServer(String serverName, boolean secure, int port) throws Exception {

        this(serverName, secure, port, getDefaultExecutor(), null);

    }

    public Http2TestServer(boolean secure, int port) throws Exception {

        this(null, secure, port, getDefaultExecutor(), null);

    }

    public InetSocketAddress getAddress() {

        return (InetSocketAddress)server.getLocalSocketAddress();

    }

    public Http2TestServer(boolean secure,

                           SSLContext context) throws Exception {

        this(null, secure, 0, null, context);

    }

    public Http2TestServer(String serverName, boolean secure,

                           SSLContext context) throws Exception {

        this(serverName, secure, 0, null, context);

    }

    public Http2TestServer(boolean secure,

                           int port,

                           ExecutorService exec,

                           SSLContext context) throws Exception {

        this(null, secure, port, exec, context);

    }

    /**

     * Create a Http2Server listening on the given port. Currently needs

     * to know in advance whether incoming connections are plain TCP "h2c"

     * or TLS "h2"/

     *

     * @param serverName SNI servername

     * @param secure https or http

     * @param port listen port

     * @param exec executor service (cached thread pool is used if null)

     * @param context the SSLContext used when secure is true

     */

    public Http2TestServer(String serverName,

                           boolean secure,

                           int port,

                           ExecutorService exec,

                           SSLContext context)

        throws Exception

    {

        this.serverName = serverName;

        if (secure) {

            server = initSecure(port);

        } else {

            server = initPlaintext(port);

        }

        this.secure = secure;

        this.exec = exec == null ? getDefaultExecutor() : exec;

        this.handlers = Collections.synchronizedMap(new HashMap<>());

        this.sslContext = context;

        this.connections = new HashMap<>();

    }

    /**

     * Adds the given handler for the given path

     */

    public void addHandler(Http2Handler handler, String path) {

        handlers.put(path, handler);

    }

    volatile Http2TestExchangeSupplier exchangeSupplier = Http2TestExchangeSupplier.ofDefault();

    /**

     * Sets an explicit exchange handler to be used for all future connections.

     * Useful for testing scenarios where non-standard or specific server

     * behaviour is required, either direct control over the frames sent, "bad"

     * behaviour, or something else.

     */

    public void setExchangeSupplier(Http2TestExchangeSupplier exchangeSupplier) {

        this.exchangeSupplier = exchangeSupplier;

    }

    Http2Handler getHandlerFor(String path) {

        if (path == null || path.equals(""))

            path = "/";

        final String fpath = path;

        AtomicReference<String> bestMatch = new AtomicReference<>("");

        AtomicReference<Http2Handler> href = new AtomicReference<>();

        handlers.forEach((key, value) -> {

            if (fpath.startsWith(key) && key.length() > bestMatch.get().length()) {

                bestMatch.set(key);

                href.set(value);

            }

        });

        Http2Handler handler = href.get();

        if (handler == null)

            throw new RuntimeException("No handler found for path " + path);

        System.err.println("Using handler for: " + bestMatch.get());

        return handler;

    }

    final ServerSocket initPlaintext(int port) throws Exception {

        return new ServerSocket(port);

    }

    public synchronized void stop() {

        // TODO: clean shutdown GoAway

        stopping = true;

        System.err.printf("Server stopping %d connections\n", connections.size());

        for (Http2TestServerConnection connection : connections.values()) {

            connection.close(ErrorFrame.NO_ERROR);

        }

        try {

            server.close();

        } catch (IOException e) {}

        exec.shutdownNow();

    }

    final ServerSocket initSecure(int port) throws Exception {

        ServerSocketFactory fac;

        if (sslContext != null) {

            fac = sslContext.getServerSocketFactory();

        } else {

            fac = SSLServerSocketFactory.getDefault();

        }

        SSLServerSocket se = (SSLServerSocket) fac.createServerSocket(port);

        SSLParameters sslp = se.getSSLParameters();

        sslp.setApplicationProtocols(new String[]{"h2"});

        se.setSSLParameters(sslp);

        se.setEnabledCipherSuites(se.getSupportedCipherSuites());

        se.setEnabledProtocols(se.getSupportedProtocols());

        // other initialisation here

        return se;

    }

    public String serverName() {

        return serverName;

    }

    private synchronized void putConnection(InetSocketAddress addr, Http2TestServerConnection c) {

        if (!stopping)

            connections.put(addr, c);

    }

    private synchronized void removeConnection(InetSocketAddress addr, Http2TestServerConnection c) {

        connections.remove(addr, c);

    }

    /**

     * Starts a thread which waits for incoming connections.

     */

    public void start() {

        exec.submit(() -> {

            try {

                while (!stopping) {

                    Socket socket = server.accept();

                    InetSocketAddress addr = (InetSocketAddress) socket.getRemoteSocketAddress();

                    Http2TestServerConnection c =

                            new Http2TestServerConnection(this, socket, exchangeSupplier);

                    putConnection(addr, c);

                    try {

                        c.run();

                    } catch (Throwable e) {

                        // we should not reach here, but if we do

                        // the connection might not have been closed

                        // and if so then the client might wait

                        // forever.

                        removeConnection(addr, c);

                        c.close(ErrorFrame.PROTOCOL_ERROR);

                        System.err.println("TestServer: start exception: " + e);

                        //throw e;

                    }

                }

            } catch (Throwable e) {

                if (!stopping) {

                    System.err.println("TestServer: terminating, caught " + e);

                    e.printStackTrace();

                }

            }

        });

    }

    @Override

    public void close() throws Exception {

        stop();

    }

}