// Modifications Copyright 2018-current Getnamo. All Rights Reserved

//
//  sio_packet.cpp
//
//  Created by Melo Yao on 3/22/15.
//

/* This disables two things:
   1) error 4503 where MSVC complains about
      decorated names being too long. There's no way around
      this.
   2) We also disable a security error triggered by
      websocketpp not using checked iterators.
*/
#ifdef _MSC_VER
#pragma warning(disable : 4503)
#define _SCL_SECURE_NO_WARNINGS
#endif

/* For this code, we will use standalone ASIO
   and websocketpp in C++11 mode only */
#define ASIO_STANDALONE
#define _WEBSOCKETPP_CPP11_STL_

#include "sio_packet.h"
#include <rapidjson/document.h>
#include <rapidjson/encodedstream.h>
#include <rapidjson/writer.h>
#include <cassert>

#define kBIN_PLACE_HOLDER "_placeholder"

#include "Runtime/Core/Public/Misc/CString.h"
#include "Runtime/Core/Public/Containers/UnrealString.h"

namespace sio
{
    using namespace rapidjsonsocketio;
    using namespace std;
    void accept_message(message const& msg, Value& val, Document& doc, vector<shared_ptr<const string> >& buffers);

    void accept_bool_message(bool_message const& msg, Value& val)
    {
        val.SetBool(msg.get_bool());
    }

    void accept_null_message(Value& val)
    {
        val.SetNull();
    }

    void accept_int_message(int_message const& msg, Value& val)
    {
        val.SetInt64(msg.get_int());
    }

    void accept_double_message(double_message const& msg, Value& val)
    {
        val.SetDouble(msg.get_double());
    }

    void accept_string_message(string_message const& msg, Value& val)
    {
        val.SetString(msg.get_string().data(), (SizeType)msg.get_string().length());
    }


    void accept_binary_message(binary_message const& msg, Value& val, Document& doc, vector<shared_ptr<const string> >& buffers)
    {
        val.SetObject();
        Value boolVal;
        boolVal.SetBool(true);
        val.AddMember(kBIN_PLACE_HOLDER, boolVal, doc.GetAllocator());
        Value numVal;
        numVal.SetInt((int)buffers.size());
        val.AddMember("num", numVal, doc.GetAllocator());
        buffers.push_back(msg.get_binary());
    }

    void accept_array_message(array_message const& msg, Value& val, Document& doc, vector<shared_ptr<const string> >& buffers)
    {
        val.SetArray();
        for (vector<message::ptr>::const_iterator it = msg.get_vector().begin(); it != msg.get_vector().end(); ++it) {
            Value child;
            accept_message(*(*it), child, doc, buffers);
            val.PushBack(child, doc.GetAllocator());
        }
    }

    void accept_object_message(object_message const& msg, Value& val, Document& doc, vector<shared_ptr<const string> >& buffers)
    {
        val.SetObject();
        for (map<string, message::ptr>::const_iterator it = msg.get_map().begin(); it != msg.get_map().end(); ++it) {
            Value nameVal;
            nameVal.SetString(it->first.data(), (SizeType)it->first.length(), doc.GetAllocator());
            Value valueVal;
            accept_message(*(it->second), valueVal, doc, buffers);
            val.AddMember(nameVal, valueVal, doc.GetAllocator());
        }
    }

    void accept_message(message const& msg, Value& val, Document& doc, vector<shared_ptr<const string> >& buffers)
    {
        const message* msg_ptr = &msg;
        switch (msg.get_flag())
        {
        case message::flag_integer:
        {
            accept_int_message(*(static_cast<const int_message*>(msg_ptr)), val);
            break;
        }
        case message::flag_double:
        {
            accept_double_message(*(static_cast<const double_message*>(msg_ptr)), val);
            break;
        }
        case message::flag_string:
        {
            accept_string_message(*(static_cast<const string_message*>(msg_ptr)), val);
            break;
        }
        case message::flag_boolean:
        {
            accept_bool_message(*(static_cast<const bool_message*>(msg_ptr)), val);
            break;
        }
        case message::flag_null:
        {
            accept_null_message(val);
            break;
        }
        case message::flag_binary:
        {
            accept_binary_message(*(static_cast<const binary_message*>(msg_ptr)), val, doc, buffers);
            break;
        }
        case message::flag_array:
        {
            accept_array_message(*(static_cast<const array_message*>(msg_ptr)), val, doc, buffers);
            break;
        }
        case message::flag_object:
        {
            accept_object_message(*(static_cast<const object_message*>(msg_ptr)), val, doc, buffers);
            break;
        }
        default:
            break;
        }
    }

    message::ptr from_json(Value const& value, vector<shared_ptr<const string> > const& buffers)
    {
        if (value.IsInt64())
        {
            return int_message::create(value.GetInt64());
        }
        else if (value.IsDouble())
        {
            return double_message::create(value.GetDouble());
        }
        else if (value.IsString())
        {
            string str(value.GetString(), value.GetStringLength());
            return string_message::create(str);
        }
        else if (value.IsArray())
        {
            message::ptr ptr = array_message::create();
            for (SizeType i = 0; i < value.Size(); ++i) {
                static_cast<array_message*>(ptr.get())->get_vector().push_back(from_json(value[i], buffers));
            }
            return ptr;
        }
        else if (value.IsObject())
        {
            //binary placeholder
            auto mem_it = value.FindMember(kBIN_PLACE_HOLDER);
            if (mem_it != value.MemberEnd() && mem_it->value.GetBool()) {

                int num = value["num"].GetInt();
                if (num >= 0 && num < static_cast<int>(buffers.size()))
                {
                    return binary_message::create(buffers[num]);
                }
                return message::ptr();
            }
            //real object message.
            message::ptr ptr = object_message::create();
            for (auto it = value.MemberBegin(); it != value.MemberEnd(); ++it)
            {
                if (it->name.IsString())
                {
                    string key(it->name.GetString(), it->name.GetStringLength());
                    static_cast<object_message*>(ptr.get())->get_map()[key] = from_json(it->value, buffers);
                }
            }
            return ptr;
        }
        else if (value.IsBool())
        {
            return bool_message::create(value.GetBool());
        }
        else if (value.IsNull())
        {
            return null_message::create();
        }
        return message::ptr();
    }

    packet::packet(string const& nsp, message::ptr const& msg, int pack_id, bool isAck) :
        _frame(frame_message),
        _type((isAck ? type_ack : type_event) | type_undetermined),
        _nsp(nsp),
        _pack_id(pack_id),
        _message(msg),
        _pending_buffers(0)
    {
        assert((!isAck
            || (isAck && pack_id >= 0)));
    }

    packet::packet(type type, string const& nsp, message::ptr const& msg) :
        _frame(frame_message),
        _type(type),
        _nsp(nsp),
        _pack_id(-1),
        _message(msg),
        _pending_buffers(0)
    {

    }

    packet::packet(packet::frame_type frame) :
        _frame(frame),
        _type(type_undetermined),
        _pack_id(-1),
        _pending_buffers(0)
    {

    }

    packet::packet() :
        _type(type_undetermined),
        _pack_id(-1),
        _pending_buffers(0)
    {

    }


    bool packet::is_binary_message(string const& payload_ptr)
    {
        return payload_ptr.size() > 0 && payload_ptr[0] == frame_message;
    }

    bool packet::is_text_message(string const& payload_ptr)
    {
        return payload_ptr.size() > 0 && payload_ptr[0] == (frame_message + '0');
    }

    bool packet::is_message(string const& payload_ptr)
    {
        return is_binary_message(payload_ptr) || is_text_message(payload_ptr);
    }

    bool packet::parse_buffer(const string& buf_payload)
    {
        if (_pending_buffers > 0) {
            //assert(is_binary_message(buf_payload));//this is ensured by outside.
            _buffers.push_back(std::make_shared<string>(buf_payload.data(), buf_payload.size()));
            _pending_buffers--;
            if (_pending_buffers == 0) {

                Document doc;
                doc.Parse<0>(_buffers.front()->data());
                _buffers.erase(_buffers.begin());
                _message = from_json(doc, _buffers);
                _buffers.clear();
                return false;
            }
            return true;
        }
        return false;
    }

    bool packet::parse(const string& payload_ptr)
    {
        assert(!is_binary_message(payload_ptr)); //this is ensured by outside
        _frame = (packet::frame_type)(payload_ptr[0] - '0');
        _message.reset();
        _pack_id = -1;
        _buffers.clear();
        _pending_buffers = 0;
        size_t pos = 1;
        if (_frame == frame_message) {
            _type = (packet::type)(payload_ptr[pos] - '0');
            if (_type < type_min || _type > type_max)
            {
                return false;
            }
            pos++;
            if (_type == type_binary_event || _type == type_binary_ack) {
                size_t score_pos = payload_ptr.find('-');
                _pending_buffers = FCString::Atoi64(*FString(payload_ptr.substr(pos, score_pos - pos).c_str()));
                pos = score_pos + 1;
            }
        }

        size_t nsp_json_pos = payload_ptr.find_first_of("{[\"/", pos, 4);
        if (nsp_json_pos == string::npos)//no namespace and no message,the end.
        {
            _nsp = "/";
            return false;
        }
        size_t json_pos = nsp_json_pos;
        if (payload_ptr[nsp_json_pos] == '/')//nsp_json_pos is start of nsp
        {
            size_t comma_pos = payload_ptr.find_first_of(",");//end of nsp
            if (comma_pos == string::npos)//packet end with nsp
            {
                _nsp = payload_ptr.substr(nsp_json_pos);
                return false;
            }
            else//we have a message, maybe the message have an id.
            {
                _nsp = payload_ptr.substr(nsp_json_pos, comma_pos - nsp_json_pos);
                pos = comma_pos + 1;//start of the message
                json_pos = payload_ptr.find_first_of("\"[{", pos, 3);//start of the json part of message
                if (json_pos == string::npos)
                {
                    //no message,the end
                    //assume if there's no message, there's no message id.
                    return false;
                }
            }
        }
        else
        {
            _nsp = "/";
        }

        if (pos < json_pos)//we've got pack id.
        {
            _pack_id = FCString::Atoi(*FString(payload_ptr.substr(pos, json_pos - pos).c_str()));
        }
        if (_frame == frame_message && (_type == type_binary_event || _type == type_binary_ack)) {
            //parse later when all buffers are arrived.
            _buffers.push_back(make_shared<string>(payload_ptr.data() + json_pos, payload_ptr.length() - json_pos));
            return true;
        }
        else
        {
            Document doc;
            doc.Parse<0>(payload_ptr.data() + json_pos);
            _message = from_json(doc, vector<shared_ptr<const string> >());
            return false;
        }

    }

    bool packet::accept(string& payload_ptr, vector<shared_ptr<const string> >& buffers)
    {
        char frame_char = _frame + '0';
        payload_ptr.append(&frame_char, 1);
        if (_frame != frame_message) {
            return false;
        }
        bool hasMessage = false;
        Document doc;
        if (_message) {
            accept_message(*_message, doc, doc, buffers);
            hasMessage = true;
        }
        bool hasBinary = buffers.size() > 0;
        _type = _type & (~type_undetermined);
        if (_type == type_event)
        {
            _type = hasBinary ? type_binary_event : type_event;
        }
        else if (_type == type_ack)
        {
            _type = hasBinary ? type_binary_ack : type_ack;
        }
        ostringstream ss;
        ss.precision(8);
        ss << _type;
        if (hasBinary) {
            ss << buffers.size() << "-";
        }
        if (_nsp.size() > 0 && _nsp != "/")
        {
            ss << _nsp;
            if (hasMessage || _pack_id >= 0) {
                ss << ",";
            }
        }

        if (_pack_id >= 0)
        {
            ss << _pack_id;
        }

        payload_ptr.append(ss.str());
        if (hasMessage)
        {
            StringBuffer buffer;
            Writer<StringBuffer> writer(buffer);
            doc.Accept(writer);
            payload_ptr.append(buffer.GetString(), buffer.GetSize());
        }
        return hasBinary;
    }

    packet::frame_type packet::get_frame() const
    {
        return _frame;
    }

    packet::type packet::get_type() const
    {
        assert((_type & type_undetermined) == 0);
        return (type)_type;
    }

    string const& packet::get_nsp() const
    {
        return _nsp;
    }

    message::ptr const& packet::get_message() const
    {
        return _message;
    }

    unsigned packet::get_pack_id() const
    {
        return _pack_id;
    }


    void packet_manager::set_decode_callback(function<void(packet const&)> const& decode_callback)
    {
        m_decode_callback = decode_callback;
    }

    void packet_manager::set_encode_callback(function<void(bool, shared_ptr<const string> const&)> const& encode_callback)
    {
        m_encode_callback = encode_callback;
    }

    void packet_manager::reset()
    {
        m_partial_packet.reset();
    }

    void packet_manager::encode(packet& pack, encode_callback_function const& override_encode_callback) const
    {
        shared_ptr<string> ptr = make_shared<string>();
        vector<shared_ptr<const string> > buffers;
        const encode_callback_function* cb_ptr = &m_encode_callback;
        if (override_encode_callback)
        {
            cb_ptr = &override_encode_callback;
        }
        if (pack.accept(*ptr, buffers))
        {
            if ((*cb_ptr))
            {
                (*cb_ptr)(false, ptr);
            }
            for (auto it = buffers.begin(); it != buffers.end(); ++it)
            {
                if ((*cb_ptr))
                {
                    (*cb_ptr)(true, *it);
                }
            }
        }
        else
        {
            if ((*cb_ptr))
            {
                (*cb_ptr)(false, ptr);
            }
        }
    }

    void packet_manager::put_payload(string const& payload)
    {
        unique_ptr<packet> p;
        do
        {
            if (packet::is_text_message(payload))
            {
                p.reset(new packet());
                if (p->parse(payload))
                {
                    m_partial_packet = std::move(p);
                }
                else
                {
                    break;
                }
            }

            //In current socket.io version, it doesn't appear binary payloads come with a frame message.
            else if (m_partial_packet && payload.size() > 0)
            {
                if (!m_partial_packet->parse_buffer(payload))
                {
                    p = std::move(m_partial_packet);
                    break;
                }
            }
            else
            {
                p.reset(new packet());
                p->parse(payload);
                break;
            }
            return;
        } while (0);

        if (m_decode_callback)
        {
            m_decode_callback(*p);
        }
    }
}