package org.apache.activemq.transport.nio;

import com.sun.xml.internal.ws.api.pipe.Engine;
import org.apache.activemq.util.IOExceptionSupport;
import org.apache.activemq.wireformat.WireFormat;

import javax.net.SocketFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLEngine;
import javax.net.ssl.SSLEngineResult;
import javax.net.ssl.SSLSession;
import java.io.IOException;
import java.net.Socket;
import java.net.URI;
import java.net.UnknownHostException;
import java.nio.ByteBuffer;
import java.sql.SQLOutput;

import static javax.net.ssl.SSLEngineResult.HandshakeStatus.NEED_UNWRAP;
import static javax.net.ssl.SSLEngineResult.HandshakeStatus.NEED_TASK;
import static javax.net.ssl.SSLEngineResult.HandshakeStatus.NEED_WRAP;
import static javax.net.ssl.SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING;
import static javax.net.ssl.SSLEngineResult.Status.BUFFER_OVERFLOW;

public class NIOSSLTransport extends NIOTransport  {

    protected SSLContext sslContext;
    protected SSLEngine sslEngine;
    protected SSLSession sslSession;

    private ByteBuffer readBuffer;
    private ByteBuffer writeBuffer;
    private ByteBuffer plain;

    public NIOSSLTransport(WireFormat wireFormat, SocketFactory socketFactory, URI remoteLocation, URI localLocation) throws UnknownHostException, IOException {
        super(wireFormat, socketFactory, remoteLocation, localLocation);
    }

    public NIOSSLTransport(WireFormat wireFormat, Socket socket) throws IOException {
        super(wireFormat, socket);
    }

    public void setSslContext(SSLContext sslContext) {
        this.sslContext = sslContext;
    }

    @Override
    protected void initializeStreams() throws IOException {
        super.initializeStreams();
        try {
            if (sslContext == null) {
                sslContext = SSLContext.getDefault();
            }
            sslEngine = sslContext.createSSLEngine();
            sslEngine.setUseClientMode(false);
            sslSession = sslEngine.getSession();
            readBuffer = ByteBuffer.allocate(sslSession.getPacketBufferSize());
            writeBuffer = ByteBuffer.allocate(sslSession.getPacketBufferSize());
            sslEngine.beginHandshake();
            plain = ByteBuffer.allocate(sslSession.getApplicationBufferSize());
        } catch (Exception e) {
            throw new IOException(e);
        }
    }

    boolean handshakeInProgress = false;

    protected void doHandshake() throws Exception {
        handshakeInProgress = true;
        while (true) {
            switch (sslEngine.getHandshakeStatus()) {
                case NEED_UNWRAP:
                    readBuffer = ByteBuffer.allocate(sslSession.getPacketBufferSize());
                    int bytesRead = channel.read(readBuffer);
                    readBuffer.flip();

                    SSLEngineResult result = sslEngine.unwrap(readBuffer, plain);
                    readBuffer.clear();
                    plain.flip();
                case NEED_TASK:
                    //TODO use the pool
                    Runnable task;
                    while ((task = sslEngine.getDelegatedTask()) != null) {
                        task.run();
                    }
                case NEED_WRAP:
                    result = sslEngine.wrap(ByteBuffer.allocate(0), writeBuffer);
                    writeBuffer.flip();
                    channel.write(writeBuffer);
                case FINISHED:
                    handshakeInProgress = false;
                    break;
            }
        }
    }

    protected void serviceRead() {
        try {
            while (true) {
                if (sslEngine.getHandshakeStatus() != NOT_HANDSHAKING) {
                    doHandshake();
                }

            }
        } catch (IOException e) {
            e.printStackTrace();
            onException(e);
        } catch (Throwable e) {
            e.printStackTrace();
            onException(IOExceptionSupport.create(e));
        }
    }
}