/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.activemq.transport.ws;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.UUID;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.TimeUnit;

import org.apache.activemq.transport.mqtt.MQTTWireFormat;
import org.apache.activemq.util.ByteSequence;
import org.eclipse.jetty.websocket.api.Session;
import org.eclipse.jetty.websocket.api.WebSocketAdapter;
import org.eclipse.jetty.websocket.api.WebSocketListener;
import org.fusesource.hawtbuf.UTF8Buffer;
import org.fusesource.mqtt.codec.CONNACK;
import org.fusesource.mqtt.codec.CONNECT;
import org.fusesource.mqtt.codec.DISCONNECT;
import org.fusesource.mqtt.codec.MQTTFrame;
import org.fusesource.mqtt.codec.PINGREQ;
import org.fusesource.mqtt.codec.PINGRESP;
import org.fusesource.mqtt.codec.PUBACK;
import org.fusesource.mqtt.codec.PUBCOMP;
import org.fusesource.mqtt.codec.PUBLISH;
import org.fusesource.mqtt.codec.PUBREC;
import org.fusesource.mqtt.codec.PUBREL;
import org.fusesource.mqtt.codec.SUBACK;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Implements a simple WebSocket based MQTT Client that can be used for unit testing.
 */
public class MQTTWSConnection extends WebSocketAdapter implements WebSocketListener {

    private static final Logger LOG = LoggerFactory.getLogger(MQTTWSConnection.class);

    private static final MQTTFrame PING_RESP_FRAME = new PINGRESP().encode();

    private Session connection;
    private final CountDownLatch connectLatch = new CountDownLatch(1);
    private final MQTTWireFormat wireFormat = new MQTTWireFormat();

    private final BlockingQueue<MQTTFrame> prefetch = new LinkedBlockingDeque<>();

    private boolean writePartialFrames;
    private int closeCode = -1;
    private String closeMessage;

    @Override
    public boolean isConnected() {
        return connection != null ? connection.isOpen() : false;
    }

    public void close() {
        if (connection != null) {
            connection.close();
        }
    }

    protected Session getConnection() {
        return connection;
    }

    //----- Connection and Disconnection methods -----------------------------//

    public void connect() throws Exception {
        connect(UUID.randomUUID().toString());
    }

    public void connect(String clientId) throws Exception {
        CONNECT command = new CONNECT();
        command.clientId(new UTF8Buffer(clientId));
        command.cleanSession(false);
        command.version(3);
        command.keepAlive((short) 0);
        connect(command);
    }

    public void connect(CONNECT command) throws Exception {
        checkConnected();

        sendBytes(wireFormat.marshal(command.encode()));

        MQTTFrame incoming = receive(15, TimeUnit.SECONDS);

        if (incoming == null || incoming.messageType() != CONNACK.TYPE) {
            throw new IOException("Failed to connect to remote service.");
        } else {
            CONNACK connack = new CONNACK().decode(incoming);
            if (!connack.code().equals(CONNACK.Code.CONNECTION_ACCEPTED)) {
                throw new IOException("Failed to connect to remote service: " + connack.code());
            }
        }
    }

    public void disconnect() throws Exception {
        if (!isConnected()) {
            return;
        }

        DISCONNECT command = new DISCONNECT();
        sendBytes(wireFormat.marshal(command.encode()));
    }

    //---- Send methods ------------------------------------------------------//

    public void sendFrame(MQTTFrame frame) throws Exception {
        checkConnected();
        sendBytes(wireFormat.marshal(frame));
    }

    public void keepAlive() throws Exception {
        checkConnected();
        sendBytes(wireFormat.marshal(new PINGREQ().encode()));
    }

    //----- Receive methods --------------------------------------------------//

    public MQTTFrame receive() throws Exception {
        checkConnected();
        return prefetch.take();
    }

    public MQTTFrame receive(long timeout, TimeUnit unit) throws Exception {
        checkConnected();
        return prefetch.poll(timeout, unit);
    }

    public MQTTFrame receiveNoWait() throws Exception {
        checkConnected();
        return prefetch.poll();
    }

    //---- Blocking state change calls ---------------------------------------//

    public void awaitConnection() throws InterruptedException {
        connectLatch.await();
    }

    public boolean awaitConnection(long time, TimeUnit unit) throws InterruptedException {
        return connectLatch.await(time, unit);
    }

    //----- Property Accessors -----------------------------------------------//

    public int getCloseCode() {
        return closeCode;
    }

    public String getCloseMessage() {
        return closeMessage;
    }

    public boolean isWritePartialFrames() {
        return writePartialFrames;
    }

    public MQTTWSConnection setWritePartialFrames(boolean value) {
        this.writePartialFrames = value;
        return this;
    }

    //----- WebSocket callback handlers --------------------------------------//

    @Override
    public void onWebSocketBinary(byte[] data, int offset, int length) {
        if (data ==null || length <= 0) {
            return;
        }

        MQTTFrame frame = null;

        try {
            frame = (MQTTFrame)wireFormat.unmarshal(new ByteSequence(data, offset, length));
        } catch (IOException e) {
            LOG.error("Could not decode incoming MQTT Frame: {}", e.getMessage());
            connection.close();
        }

        try {
            switch (frame.messageType()) {
            case PINGREQ.TYPE:
                PINGREQ ping = new PINGREQ().decode(frame);
                LOG.info("WS-Client read frame: {}", ping);
                sendFrame(PING_RESP_FRAME);
                break;
            case PINGRESP.TYPE:
                LOG.info("WS-Client ping response received.");
                break;
            case CONNACK.TYPE:
                CONNACK connAck = new CONNACK().decode(frame);
                LOG.info("WS-Client read frame: {}", connAck);
                prefetch.put(frame);
                break;
            case SUBACK.TYPE:
                SUBACK subAck = new SUBACK().decode(frame);
                LOG.info("WS-Client read frame: {}", subAck);
                prefetch.put(frame);
                break;
            case PUBLISH.TYPE:
                PUBLISH publish = new PUBLISH().decode(frame);
                LOG.info("WS-Client read frame: {}", publish);
                prefetch.put(frame);
                break;
            case PUBACK.TYPE:
                PUBACK pubAck = new PUBACK().decode(frame);
                LOG.info("WS-Client read frame: {}", pubAck);
                prefetch.put(frame);
                break;
            case PUBREC.TYPE:
                PUBREC pubRec = new PUBREC().decode(frame);
                LOG.info("WS-Client read frame: {}", pubRec);
                prefetch.put(frame);
                break;
            case PUBREL.TYPE:
                PUBREL pubRel = new PUBREL().decode(frame);
                LOG.info("WS-Client read frame: {}", pubRel);
                prefetch.put(frame);
                break;
            case PUBCOMP.TYPE:
                PUBCOMP pubComp = new PUBCOMP().decode(frame);
                LOG.info("WS-Client read frame: {}", pubComp);
                prefetch.put(frame);
                break;
            default:
                LOG.error("Unknown MQTT  Frame received.");
                connection.close();
            }
        } catch (Exception e) {
            LOG.error("Could not decode incoming MQTT Frame: {}", e.getMessage());
            connection.close();
        }
    }

    //----- Internal implementation ------------------------------------------//

    private void sendBytes(ByteSequence payload) throws IOException {
        if (!isWritePartialFrames()) {
            connection.getRemote().sendBytes(ByteBuffer.wrap(payload.data, payload.offset, payload.length));
        } else {
            connection.getRemote().sendBytes(ByteBuffer.wrap(
                payload.data, payload.offset, payload.length / 2));
            connection.getRemote().sendBytes(ByteBuffer.wrap(
                payload.data, payload.offset + payload.length / 2, payload.length / 2));
        }
    }

    private void checkConnected() throws IOException {
        if (!isConnected()) {
            throw new IOException("MQTT WS Connection is closed.");
        }
    }

    @Override
    public void onWebSocketClose(int statusCode, String reason) {
        LOG.trace("MQTT WS Connection closed, code:{} message:{}", statusCode, reason);

        this.connection = null;
        this.closeCode = statusCode;
        this.closeMessage = reason;

    }

    @Override
    public void onWebSocketConnect(org.eclipse.jetty.websocket.api.Session session) {
        this.connection = session;
        this.connection.setIdleTimeout(Duration.ZERO);
        this.connectLatch.countDown();
    }
}
