#pragma once

//处理websocket相关数据的握手，拆解和拼装
//add by freeeyes

#include <stdint.h>
#include <chrono>
#include <random>
#include <string>
#include "Base64.hpp"
#include "SHA1.hpp"

class WebSocketFormat
{
public:
    enum class WebSocketFrameType
    {
        ERROR_FRAME = 0xff,
        CONTINUATION_FRAME = 0x00,
        TEXT_FRAME = 0x01,
        BINARY_FRAME = 0x02,
        CLOSE_FRAME = 0x08,
        PING_FRAME = 0x09,
        PONG_FRAME = 0x0A
    };

    static std::string wsHandshake(std::string secKey)
    {
        secKey.append("258EAFA5-E914-47DA-95CA-C5AB0DC85B11");

        CSHA1 s1;
        s1.Update((unsigned char*)secKey.c_str(), static_cast<unsigned int>(secKey.size()));
        s1.Final();
        unsigned char puDest[20];
        s1.GetHash(puDest);

        std::string base64Str = brynet::base::crypto::base64_encode((const unsigned char*)puDest, 20);
        return base64Str;
    }

    static bool wsFrameBuild(const char* payload,
        size_t payloadLen,
        std::string& frame,
        WebSocketFrameType frame_type = WebSocketFrameType::TEXT_FRAME,
        bool isFin = true,
        bool masking = false)
    {
        const auto unixTime = std::chrono::system_clock::now().time_since_epoch().count();
        static std::mt19937 random(static_cast<unsigned int>(unixTime));

        static_assert(std::is_same<std::string::value_type, char>::value, "");

        const uint8_t head = static_cast<uint8_t>(frame_type) | (isFin ? 0x80 : 0x00);

        frame.clear();
        frame.push_back(static_cast<char>(head));
        if (payloadLen <= 125)
        {
            // mask << 7 | payloadLen, mask = 0
            frame.push_back(static_cast<uint8_t>(payloadLen));
        }
        else if (payloadLen <= 0xFFFF)
        {
            // 126 + 16bit len
            frame.push_back(126);
            frame.push_back((payloadLen & 0xFF00) >> 8);
            frame.push_back(payloadLen & 0x00FF);
        }
        else
        {
            // 127 + 64bit len
            frame.push_back(127);
            // assume payload len is less than u_int32_max
            frame.push_back(0x00);
            frame.push_back(0x00);
            frame.push_back(0x00);
            frame.push_back(0x00);
            frame.push_back(static_cast<char>((payloadLen & 0xFF000000) >> 24));
            frame.push_back(static_cast<char>((payloadLen & 0x00FF0000) >> 16));
            frame.push_back(static_cast<char>((payloadLen & 0x0000FF00) >> 8));
            frame.push_back(static_cast<char>(payloadLen & 0x000000FF));
        }

        if (masking)
        {
            frame[1] = ((uint8_t)frame[1]) | 0x80;
            uint8_t mask[4];
            for (auto& m : mask)
            {
                m = static_cast<uint8_t>(random());
                frame.push_back(m);
            }

            frame.reserve(frame.size() + payloadLen);

            for (size_t i = 0; i < payloadLen; i++)
            {
                frame.push_back(static_cast<uint8_t>(payload[i]) ^ mask[i % 4]);
            }
        }
        else
        {
            frame.append(payload, payloadLen);
        }

        return true;
    }

    static bool wsFrameBuild(const std::string& payload,
        std::string& frame,
        WebSocketFrameType frame_type = WebSocketFrameType::TEXT_FRAME,
        bool isFin = true,
        bool masking = false)
    {
        return wsFrameBuild(payload.c_str(),
            payload.size(),
            frame,
            frame_type,
            isFin,
            masking);
    }

    static bool wsFrameExtractBuffer(const char* inbuffer,
        const size_t bufferSize,
        std::string& payload,
        WebSocketFrameType& outopcode,
        size_t& frameSize,
        bool& outfin)
    {
        const auto buffer = (const unsigned char*)inbuffer;

        if (bufferSize < 2)
        {
            return false;
        }

        outfin = (buffer[0] & 0x80) != 0;
        outopcode = (WebSocketFrameType)(buffer[0] & 0x0F);
        const bool isMasking = (buffer[1] & 0x80) != 0;
        uint32_t payloadlen = buffer[1] & 0x7F;

        uint32_t pos = 2;
        if (payloadlen == 126)
        {
            if (bufferSize < 4)
            {
                return false;
            }

            payloadlen = (buffer[2] << 8) + buffer[3];
            pos = 4;
        }
        else if (payloadlen == 127)
        {
            if (bufferSize < 10)
            {
                return false;
            }

            if (buffer[2] != 0 ||
                buffer[3] != 0 ||
                buffer[4] != 0 ||
                buffer[5] != 0)
            {
                return false;
            }

            if ((buffer[6] & 0x80) != 0)
            {
                return false;
            }

            payloadlen = (buffer[6] << 24) +
                (buffer[7] << 16) +
                (buffer[8] << 8) +
                buffer[9];
            pos = 10;
        }

        uint8_t mask[4];
        if (isMasking)
        {
            if (bufferSize < (pos + 4))
            {
                return false;
            }

            mask[0] = buffer[pos++];
            mask[1] = buffer[pos++];
            mask[2] = buffer[pos++];
            mask[3] = buffer[pos++];
        }

        if (bufferSize < (pos + payloadlen))
        {
            return false;
        }

        if (isMasking)
        {
            payload.reserve(payloadlen);
            for (size_t j = 0; j < payloadlen; j++)
                payload.push_back(buffer[pos + j] ^ mask[j % 4]);
        }
        else
        {
            payload.append((const char*)(buffer + pos), payloadlen);
        }

        frameSize = payloadlen + pos;

        return true;
    }

    static bool wsFrameExtractString(const std::string& buffer,
        std::string& payload,
        WebSocketFrameType& opcode,
        size_t& frameSize, bool& isFin)
    {
        return wsFrameExtractBuffer(buffer.c_str(),
            buffer.size(),
            payload,
            opcode,
            frameSize,
            isFin);
    }
};

