/*
 * The MIT License (MIT)
 *
 * Copyright (c) 2020 Ziver Koc
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

package zutil.net.mqtt;

import zutil.ObjectUtil;
import zutil.log.LogUtil;
import zutil.net.mqtt.packet.*;
import zutil.net.mqtt.packet.MqttPacketSubscribe.MqttSubscribePayload;
import zutil.net.mqtt.packet.MqttPacketSubscribeAck.MqttSubscribeAckPayload;
import zutil.net.mqtt.packet.MqttPacketUnsubscribe.MqttUnsubscribePayload;
import zutil.net.threaded.ThreadedTCPNetworkServer;
import zutil.net.threaded.ThreadedTCPNetworkServerThread;
import zutil.parser.binary.BinaryStructInputStream;
import zutil.parser.binary.BinaryStructOutputStream;

import java.io.BufferedInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.Socket;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * TODO:
 *
 * @see <a href="http://docs.oasis-open.org/mqtt/mqtt/v3.1.1/mqtt-v3.1.1.html">MQTT v3.1.1 Spec</a>
 */
public class MqttBroker extends ThreadedTCPNetworkServer {
    private static final Logger logger = LogUtil.getLogger();

    public static final int MQTT_PORT = 1883;
    public static final int MQTT_PORT_TLS = 8883;
    public static final int MQTT_PROTOCOL_VERSION = 0x04; // MQTT 3.1.1

    private List<MqttSubscriptionListener> globalListeners = new ArrayList<>();
    private Map<String, MqttTopic> topics = new ConcurrentHashMap<>();

    protected static class MqttTopic {
        public String name;
        /** If set this will be sent as a publish message as soon as a client subscribes to this topic */
        public byte[] retainedPayload;

        public List<MqttSubscriptionListener> subscribers = new ArrayList<>(0); // Stat with low capacity to save memory
    }


    public MqttBroker() throws IOException {
        super(MQTT_PORT);
    }

    public MqttBroker(int port) throws IOException {
        super(port);
    }


    @Override
    protected ThreadedTCPNetworkServerThread getThreadInstance(Socket s) throws IOException {
        return new MqttConnectionThread(this, s);
    }

    /**
     * @param topic the topic name
     * @return a new topic if it does not exist otherwise returns the existing topic.
     */
    private MqttTopic createTopic(String topic) {
        if (!topics.containsKey(topic)) {
            logger.finest("Creating new topic: " + topic);
            MqttTopic topicData = new MqttTopic();
            topicData.name = topic;
            topics.put(topic, topicData);
        }

        return topics.get(topic);
    }

    /**
     * Will check if the topic is empty, if so its data will be cleared.
     * A topic is empty if: <ul>
     *     <li>No subscribers</li>
     *     <li>No retained message</li>
     * </ul>
     *
     * @param topic the topic name
     */
    private void removeTopicIfEmpty(String topic) {
        MqttTopic topicData = topics.get(topic);
        if (topicData == null)
            return;

        if (topicData.retainedPayload == null &&
                topicData.subscribers.isEmpty()) {
            logger.finest("Removing empty topic: " + topic);
            topics.remove(topic);
        }
    }

    /**
     * @return the subscriber count for the specific topic, -1 if
     *         topic does not exist or has not been created yet.
     */
    public int getSubscriberCount(String topicName) {
        MqttTopic topicData = topics.get(topicName);

        if (topicData != null)
            return topicData.subscribers.size();
        return -1;
    }

    /**
     * Assign a listener that will be called for any topic.
     *
     * @param listener the listener that will be called.
     */
    public void addGlobalSubscriber(MqttSubscriptionListener listener) {
        globalListeners.add(listener);
    }

    /**
     * Remove the given listener from global subscribers.
     *
     * @param listener the listener that should not be called anymore.
     */
    public void removeGlobalSubscriber(MqttSubscriptionListener listener) {
        globalListeners.remove(listener);
    }

    /**
     * Add the listener as a subscriber of the specific topic.
     *
     * @param topicName the topic that will be subscribed to.
     * @param listener  the listener that will be called.
     */
    public synchronized void subscribe(String topicName, MqttSubscriptionListener listener) {
        if (topicName == null || topicName.isEmpty() || listener == null)
            return;

        MqttTopic topicData = createTopic(topicName);

        if (!topicData.subscribers.contains(listener)) {
            topicData.subscribers.add(listener);
            logger.finer("New subscriber on topic: " + topicName + " (subscriber count: " + topicData.subscribers.size() + ")");

            // Send any retained messages for the topic
            if (topicData.retainedPayload != null) {
                listener.dataPublished(topicName, topicData.retainedPayload);
            }
        }
    }

    /**
     * Publish data to the specific topic.
     *
     * @param topicName the topic where the data should be published to.
     * @param data      the data that should be published. String will be converted to UTF-8 byte array.
     */
    public void publish(String topicName, String data) {
        publish(topicName, data.getBytes(StandardCharsets.UTF_8));
    }

    /**
     * Publish data to the specific topic.
     *
     * @param topicName the topic where the data should be published to.
     * @param data      the data that should be published.
     */
    public void publish(String topicName, byte[] data) {
        logger.finest("Data has been published to topic: " + topicName);

        if (globalListeners != null) {
            publishToSubscribers(globalListeners, topicName, data);
        }

        MqttTopic topicData = topics.get(topicName);
        if (topicData != null) {
            publishToSubscribers(topicData.subscribers, topicName, data);
        }
    }

    private static void publishToSubscribers(List<MqttSubscriptionListener> subscriberList, String topicName, byte[] data) {
        for (MqttSubscriptionListener subscriber : subscriberList) {
            try {
                subscriber.dataPublished(topicName, data);
            } catch (Exception e) {
                logger.log(Level.WARNING, "MqttSubscriptionListener threw exception.", e);
            }
        }
    }

    /**
     * Will retain the given publish message for the specific topic, the message will be sent to any new client subscribing to the topic.
     * Note, only a single message is saved for each topic.
     *
     * @param topicName the topic where the message should be retained.
     * @param data      the payload to be stored for the topic. Set to null to clear retained message.
     */
    public void retain(String topicName, byte[] data) {
        MqttTopic topicData = createTopic(topicName);

        if (topicData != null) {
            if (ObjectUtil.isEmpty(data)) {
                logger.finer("Retained message removed for topic: " + topicName);
                topicData.retainedPayload = null;
            } else {
                logger.finer("New message retained for topic: " + topicName);
                topicData.retainedPayload = data;
            }
        }
    }

    /**
     * Unsubscribe the listener from all available MQTT topics.
     *
     * @param listener the listener that should be unsubscribed.
     */
    public synchronized void unsubscribe(MqttSubscriptionListener listener) {
        if (listener == null)
            return;

        for (MqttTopic topic : topics.values()) {
            unsubscribe(topic.name, listener);
        }
    }

    /**
     * Unsubscribe the listener from the specific MQTT topic.
     *
     * @param topicName the specific topic that should be unsubscribed from.
     * @param listener  the target listener that should be unsubscribed.
     */
    public synchronized void unsubscribe(String topicName, MqttSubscriptionListener listener) {
        if (topicName == null || topicName.isEmpty() || listener == null)
            return;

        MqttTopic topicData = topics.get(topicName);
        if (topicData != null) {
            if (topicData.subscribers.remove(listener)) {
                logger.finer("Subscriber unsubscribed from topic " + topicName + " (subscriber count: " + topicData.subscribers.size() + ")");

                removeTopicIfEmpty(topicName);
            }
        }
    }


    protected static class MqttConnectionThread implements ThreadedTCPNetworkServerThread, MqttSubscriptionListener {
        private final MqttBroker broker;
        private Socket socket;
        private BinaryStructInputStream in;
        private BinaryStructOutputStream out;
        private String clientId;

        private boolean disconnected = false;
        /** A message that should be sent in case the connection to client is abnormally disconnected */
        private MqttPacketPublish willPacket = null;
        /** The maximum amount of time(seconds) to wait for activity from client, 0 means no timeout */
        private int connectionTimeoutTime = 0;

        /**
         * Test constructor
         */
        protected MqttConnectionThread(MqttBroker b) {
            broker = b;
        }

        public MqttConnectionThread(MqttBroker b, Socket s) throws IOException {
            this(b);
            socket = s;
            clientId = socket.getInetAddress().toString();

            InputStream baseInputstream = socket.getInputStream();
            if (!baseInputstream.markSupported())
                baseInputstream = new BufferedInputStream(baseInputstream);

            in = new BinaryStructInputStream(baseInputstream);
            out = new BinaryStructOutputStream(socket.getOutputStream());
        }


        @Override
        public void run() {
            try {
                // Setup connection

                logger.fine("[" + socket.getInetAddress() + "] New MQTT client connected.");
                MqttPacketHeader connectPacket = MqttPacket.read(in);
                handleConnect(connectPacket);

                // Connected

                while (!disconnected) {
                    MqttPacketHeader packet = MqttPacket.read(in);
                    if (packet == null)
                        return;

                    handlePacket(packet);
                }

            } catch (IOException e) {
                logger.log(Level.WARNING, "[" + socket.getInetAddress() + "] There was a issue with the client connection.", e);
            } finally {
                try {
                    sendWillPacket();

                    logger.fine("[" + socket.getInetAddress() + "] MQTT client disconnected.");
                    socket.close();
                    broker.unsubscribe(this);
                } catch (IOException e) {
                    logger.log(Level.SEVERE, null, e);
                }
            }
        }

        private void handleConnect(MqttPacketHeader connectPacket) throws IOException {
            // Unexpected packet?
            if (!(connectPacket instanceof MqttPacketConnect))
                throw new IOException("Expected MqttPacketConnect but received " + connectPacket.getClass());
            MqttPacketConnect conn = (MqttPacketConnect) connectPacket;

            // Reply
            MqttPacketConnectAck connectAck = new MqttPacketConnectAck();

            // ----------------------------------
            // Handling Header
            // ----------------------------------

            // Incorrect protocol version?
            if (conn.protocolLevel != MQTT_PROTOCOL_VERSION) {
                connectAck.returnCode = MqttPacketConnectAck.RETCODE_PROT_VER_ERROR;
                sendPacket(connectAck);
                return;
            } else {
                connectAck.returnCode = MqttPacketConnectAck.RETCODE_OK;
            }

            // Is reserved field properly set? should be false
            if (conn.flagReserved) {
                disconnected = true;
                return;
            }

            // Handle Session
            if (conn.flagCleanSession) {
                // TODO: Remove session
                connectAck.flagSessionPresent = false;
            } else {
                // TODO: Restore or create new session
                throw new UnsupportedOperationException("Sessions currently not supported.");
            }

            // Handle will message
            if (conn.flagWillFlag) {
                logger.fine("[" + socket.getInetAddress() + "] Registered will packet for topic: " + conn.willTopic);
                willPacket = new MqttPacketPublish();
                willPacket.topicName = conn.willTopic;
                willPacket.setFlagQoS(conn.flagWillQoS);
                willPacket.setFlagRetain(conn.flagWillRetain);
            } else {
                willPacket = null;
            }

            // TODO: authenticate
            if (conn.flagUsername) {
                String username;

                if (conn.flagPassword) {
                    String password;
                }
            }

            connectionTimeoutTime = conn.keepAlive;

            // ----------------------------------
            // Handling Payload
            // ----------------------------------


            sendPacket(connectAck);
        }

        protected void handlePacket(MqttPacketHeader packet) throws IOException {
            // TODO: QOS 1
            // TODO: QOS 2
            // TODO: handle connection timeout

            switch (packet.type) {
                case MqttPacketHeader.PACKET_TYPE_PUBLISH:
                    handlePublish((MqttPacketPublish) packet);
                    break;

                case MqttPacketHeader.PACKET_TYPE_PUBREC:
                    break;

                case MqttPacketHeader.PACKET_TYPE_SUBSCRIBE:
                    handleSubscribe((MqttPacketSubscribe) packet);
                    break;

                case MqttPacketHeader.PACKET_TYPE_UNSUBSCRIBE:
                    handleUnsubscribe((MqttPacketUnsubscribe) packet);
                    break;

                // Ping
                case MqttPacketHeader.PACKET_TYPE_PINGREQ:
                    sendPacket(new MqttPacketPingResp());
                    break;

                // Close connection
                default:
                    logger.warning("[" + clientId + "] Received unknown packet type: " + packet.type + " (" + packet.getClass() + ")");
                    sendWillPacket(); // Only publish will message in case of faulty disconnect
                    /* FALLTHROUGH */
                case MqttPacketHeader.PACKET_TYPE_DISCONNECT:
                    willPacket = null;
                    disconnected = true;
                    break;
            }
        }


        private void handlePublish(MqttPacketPublish publishPacket) {
            if (publishPacket.getFlagQoS() != 0)
                throw new UnsupportedOperationException("QoS larger then 0 not supported.");
            if (publishPacket.getFlagQoS() >= 3)
                throw new IllegalArgumentException("QoS value of 3 is not valid.");

            logger.finer("[" + clientId + "] Publishing to topic: " + publishPacket.topicName);
            broker.publish(publishPacket.topicName, publishPacket.payload);

            if (publishPacket.getFlagRetain()) {
                broker.retain(publishPacket.topicName, publishPacket.payload);
            }
        }

        private void handleSubscribe(MqttPacketSubscribe subscribePacket) throws IOException {
            MqttPacketSubscribeAck subscribeAckPacket = new MqttPacketSubscribeAck();
            subscribeAckPacket.packetId = subscribePacket.packetId;

            for (MqttSubscribePayload payload : subscribePacket.payloads) {
                logger.finer("[" + clientId + "] Subscribing to topic: " + payload.topicFilter);
                broker.subscribe(payload.topicFilter, this);

                // Prepare response
                MqttSubscribeAckPayload ackPayload = new MqttSubscribeAckPayload();
                ackPayload.returnCode = MqttSubscribeAckPayload.RETCODE_SUCESS_MAX_QOS_0;
                subscribeAckPacket.payloads.add(ackPayload);
            }

            sendPacket(subscribeAckPacket);
        }

        private void handleUnsubscribe(MqttPacketUnsubscribe unsubscribePacket) throws IOException {
            for (MqttUnsubscribePayload payload : unsubscribePacket.payloads) {
                logger.finer("[" + clientId + "] Unsubscribing from topic: " + payload.topicFilter);
                broker.unsubscribe(payload.topicFilter, this);
            }

            // Prepare response
            MqttPacketUnsubscribeAck unsubscribeAckPacket = new MqttPacketUnsubscribeAck();
            unsubscribeAckPacket.packetId = unsubscribePacket.packetId;
            sendPacket(unsubscribeAckPacket);
        }


        @Override
        public void dataPublished(String topic, byte[] data) {
            try {
                // Prepare response
                MqttPacketPublish publishPacket = new MqttPacketPublish();
                publishPacket.setFlagQoS(MqttPacketPublish.PUBLISH_QOS_0);
                publishPacket.topicName = topic;
                publishPacket.payload = data;
                sendPacket(publishPacket);

                // TODO: QOS 1
                // TODO: QOS 2
            } catch (IOException e) {
                logger.log(Level.WARNING, "Was unable to publish topic to client " + clientId, e);
            }
        }

        private void sendWillPacket() throws IOException {
            if (willPacket != null) {
                logger.fine("[" + clientId + "] Publishing will packet.");
                // TODO: Change to handlePublish when QoS is implemented
                broker.publish(willPacket.topicName, willPacket.payload);

                if (willPacket.getFlagRetain()) {
                    broker.retain(willPacket.topicName, willPacket.payload);
                }

                willPacket = null;
            }
        }

        public synchronized void sendPacket(MqttPacketHeader packet) throws IOException {
            MqttPacket.write(out, packet);
        }


        public boolean isDisconnected() {
            return disconnected;
        }
    }
}
