#include "WebSocket.h"

#include <iostream>

#include "HttpRequest.h"

using namespace Network;
using namespace WebSocket;
using namespace Framework;

__declspec(dllexport) Frame& Frame::operator+=(
    const Frame& b) // baut frames zusammen welche zu einer nachricht gehören
{
    fin = b.fin;
    if (opcode == 0) opcode = b.opcode;
    dataLength += b.dataLength;
    if (dataLength)
    {
        char* data = new char[(int)dataLength];
        if (data) memcpy(data, this->data, (int)(dataLength - b.dataLength));
        if (b.data) memcpy(data + dataLength, b.data, (int)b.dataLength);
        delete[] this->data;
        this->data = data;
    }
    else
    {
        delete[] this->data;
        this->data = 0;
    }
    return *this;
}

__declspec(dllexport) WebSocketClient::WebSocketClient(
    const char* path, const char* host, unsigned short port)
    : Thread()
{
    queue = new Array<Frame>();
    callback = 0;
    klient = 0;
    this->path = path;
    this->host = host;
    this->port = port;
    lastPingFrame = 0;
    nextClose = 0;
}

WebSocketClient::~WebSocketClient()
{
    disconnect();
    while (queue->getEintragAnzahl())
    {
        Frame f = queue->get(0);
        delete[] f.data;
        queue->remove(0);
    }
    queue->release();
}

__declspec(dllexport) void WebSocketClient::setMessageCallback(
    std::function<void(
        WebSocketClient*, __int64 size, const char* data, DataType typ)>
        callback)
{
    this->callback = callback;
}

__declspec(dllexport) bool WebSocketClient::connect()
{
    char allowedKeyChars[] = {'a',
        'b',
        'c',
        'd',
        'e',
        'f',
        'g',
        'h',
        'i',
        'j',
        'k',
        'l',
        'm',
        'n',
        'o',
        'p',
        'q',
        'r',
        's',
        't',
        'u',
        'v',
        'w',
        'x',
        'y',
        'z',
        '0',
        '1',
        '2',
        '3',
        '4',
        '5',
        '6',
        '7',
        '8',
        '9',
        'A',
        'B',
        'C',
        'D',
        'E',
        'F',
        'G',
        'H',
        'I',
        'J',
        'K',
        'L',
        'M',
        'N',
        'O',
        'P',
        'Q',
        'R',
        'S',
        'T',
        'U',
        'V',
        'W',
        'X',
        'Y',
        'Z',
        0};
    __int64 numKeyChars = strlen(allowedKeyChars);
    if (!klient)
    {
        Text message = "GET ";
        message += path;
        message += " HTTP/1.1\r\n";
        message += "Host: ";
        message += host;
        message += "\r\nUpgrade: websocket\r\n";
        message += "Connection: Upgrade\r\n";
        char* key = new char[25];
        key[24] = 0;
        key[23] = '=';
        key[22] = '=';
        for (int i = 0; i < 22; i++)
            key[i] = allowedKeyChars[(int)(gen.rand() * numKeyChars)];
        message += "Sec-WebSocket-Key: ";
        message += key;
        delete[] key;
        message += "\r\nSec-WebSocket-Version: 13\r\n\r\n";
        klient = new Klient();
        if (!klient->verbinde(port, host))
        {
            klient = (Klient*)klient->release();
            return 0;
        }
        klient->sende(message, message.getLength());
        Text answer;
        int br = 0;
        do
        {
            char buff[2];
            buff[1] = 0;
            if (klient->getNachricht(buff, 1))
                answer += buff;
            else
                break;
            if (buff[0] == '\n')
                br++;
            else if (buff[0] != '\r')
                br = 0;
        } while (br < 2 && klient->hatNachricht(1000));
        HTTP::Answer* handshakeResponse = new HTTP::Answer(answer);
        if (handshakeResponse->getStatusCode() != 101)
        {
            handshakeResponse->release();
            klient = (Klient*)klient->release();
            return 0;
        }
        handshakeResponse->release();
        start();
        return 1;
    }
    else
        return 1;
}

__declspec(dllexport) bool WebSocketClient::send(
    __int64 size, const char* data, DataType typ)
{
    Frame f;
    f.fin = 1;
    f.rsv1 = 0;
    f.rsv2 = 0;
    f.rsv3 = 0;
    f.mask = 1;
    if (typ == TEXT)
        f.opcode = 1;
    else
        f.opcode = 2;
    f.dataLength = size;
    f.data = new char[(int)f.dataLength];
    memcpy(f.data, data, (int)f.dataLength);
    f.key[0] = (unsigned char)(gen.rand() * 256);
    f.key[1] = (unsigned char)(gen.rand() * 256);
    f.key[2] = (unsigned char)(gen.rand() * 256);
    f.key[3] = (unsigned char)(gen.rand() * 256);
    c.lock();
    queue->add(f);
    c.unlock();
    return 1;
}

__declspec(dllexport) void WebSocketClient::thread()
{
    while (klient)
    {
        c2.lock();
        if (!klient)
        {
            c2.unlock();
            return;
        }
        if (klient->hatNachricht(100))
        {
            bool ok = 1;
            c2.unlock();
            bool first = 1;
            Frame m;
            unsigned char byte;
            do
            {
                c2.lock();
                if (!klient)
                {
                    c2.unlock();
                    return;
                }
                Frame message;
                ok &= klient->getNachricht((char*)&byte, 1);
                message.fin = (byte & 0x80) != 0;
                message.rsv1 = (byte & 0x40) != 0;
                message.rsv2 = (byte & 0x20) != 0;
                message.rsv3 = (byte & 0x10) != 0;
                message.opcode = byte & 0xF;
                ok &= klient->getNachricht((char*)&byte, 1);
                message.mask = (byte & 0x80) != 0;
                message.dataLength = byte & 0x7F;
                if (message.dataLength == 126)
                {
                    ok &= klient->getNachricht((char*)&byte, 1);
                    message.dataLength = byte << 8;
                    ok &= klient->getNachricht((char*)&byte, 1);
                    message.dataLength |= byte;
                }
                else if (message.dataLength == 127)
                {
                    ok &= klient->getNachricht((char*)&byte, 1);
                    message.dataLength = (__int64)byte << 56;
                    ok &= klient->getNachricht((char*)&byte, 1);
                    message.dataLength |= (__int64)byte << 48;
                    ok &= klient->getNachricht((char*)&byte, 1);
                    message.dataLength |= (__int64)byte << 40;
                    ok &= klient->getNachricht((char*)&byte, 1);
                    message.dataLength |= (__int64)byte << 32;
                    ok &= klient->getNachricht((char*)&byte, 1);
                    message.dataLength |= (__int64)byte << 24;
                    ok &= klient->getNachricht((char*)&byte, 1);
                    message.dataLength |= (__int64)byte << 16;
                    ok &= klient->getNachricht((char*)&byte, 1);
                    message.dataLength |= (__int64)byte << 8;
                    ok &= klient->getNachricht((char*)&byte, 1);
                    message.dataLength |= (__int64)byte;
                }
                if (message.mask)
                {
                    ok &= klient->getNachricht((char*)&byte, 1);
                    message.key[0] = byte;
                    ok &= klient->getNachricht((char*)&byte, 1);
                    message.key[1] = byte;
                    ok &= klient->getNachricht((char*)&byte, 1);
                    message.key[2] = byte;
                    ok &= klient->getNachricht((char*)&byte, 1);
                    message.key[3] = byte;
                }
                if (!ok) message.dataLength = 1;
                message.data = 0;
                if (message.dataLength)
                    message.data = new char[(int)message.dataLength];
                for (int i = 0; i < message.dataLength; i++)
                {
                    ok &= klient->getNachricht((char*)&byte, 1);
                    if (message.mask)
                        message.data[i] = byte ^ message.key[i % 4];
                    else
                        message.data[i] = byte;
                }
                c2.unlock();
                if (first)
                    m = message;
                else
                {
                    m += message;
                    delete[] message.data;
                }
                first = 0;
                if (!ok) break;
            } while (!m.fin);
            if (!ok)
            {
                delete[] m.data;
                return;
            }
            if (m.opcode == 0x9)
            {
                m.opcode = 0xA;
                m.key[0] = (unsigned char)(gen.rand() * 256);
                m.key[1] = (unsigned char)(gen.rand() * 256);
                m.key[2] = (unsigned char)(gen.rand() * 256);
                m.key[3] = (unsigned char)(gen.rand() * 256);
                queue->add(m);
            }
            else if (m.opcode == 0xA)
            {
                delete[] m.data;
            }
            else if (m.opcode == 0x8)
            {
                if (nextClose)
                {
                    delete[] m.data;
                    c2.lock();
                    klient->trenne();
                    klient = (Klient*)klient->release();
                    c2.unlock();
                    return;
                }
                else
                {
                    m.key[0] = (unsigned char)(gen.rand() * 256);
                    m.key[1] = (unsigned char)(gen.rand() * 256);
                    m.key[2] = (unsigned char)(gen.rand() * 256);
                    m.key[3] = (unsigned char)(gen.rand() * 256);
                    queue->add(m);
                    nextClose = 1;
                }
            }
            else if (callback)
            {
                callback(
                    this, m.dataLength, m.data, m.opcode == 1 ? TEXT : BINARY);
                delete[] m.data;
            }
        }
        else
        {
            c2.unlock();
            c.lock();
            while (queue->getEintragAnzahl())
            {
                Frame f = queue->get(0);
                c.unlock();
                unsigned char byte = (f.fin ? 1 : 0) << 7;
                byte |= (f.rsv1 ? 1 : 0) << 6;
                byte |= (f.rsv2 ? 1 : 0) << 5;
                byte |= (f.rsv3 ? 1 : 0) << 4;
                byte |= f.opcode & 0xF;
                c2.lock();
                if (!klient)
                {
                    c2.unlock();
                    return;
                }
                klient->sende((char*)&byte, 1);
                byte = (f.mask ? 1 : 0) << 7;
                if (f.dataLength < 126)
                    byte |= (unsigned char)f.dataLength & 0x7F;
                else if (f.dataLength <= 32767)
                    byte |= 126;
                else
                    byte |= 127;
                klient->sende((char*)&byte, 1);
                if (f.dataLength >= 126)
                {
                    if (f.dataLength <= 32767)
                    {
                        byte = (unsigned char)(f.dataLength >> 8);
                        klient->sende((char*)&byte, 1);
                        byte = (unsigned char)f.dataLength & 0xFF;
                        klient->sende((char*)&byte, 1);
                    }
                    else
                    {
                        byte = (unsigned char)(f.dataLength >> 56);
                        klient->sende((char*)&byte, 1);
                        byte = (unsigned char)(f.dataLength >> 48);
                        klient->sende((char*)&byte, 1);
                        byte = (unsigned char)(f.dataLength >> 40);
                        klient->sende((char*)&byte, 1);
                        byte = (unsigned char)(f.dataLength >> 32);
                        klient->sende((char*)&byte, 1);
                        byte = (unsigned char)(f.dataLength >> 24);
                        klient->sende((char*)&byte, 1);
                        byte = (unsigned char)(f.dataLength >> 16);
                        klient->sende((char*)&byte, 1);
                        byte = (unsigned char)(f.dataLength >> 8);
                        klient->sende((char*)&byte, 1);
                        byte = (unsigned char)f.dataLength & 0xFF;
                        klient->sende((char*)&byte, 1);
                    }
                }
                if (f.mask)
                {
                    byte = (unsigned char)f.key[0];
                    klient->sende((char*)&byte, 1);
                    byte = (unsigned char)f.key[1];
                    klient->sende((char*)&byte, 1);
                    byte = (unsigned char)f.key[2];
                    klient->sende((char*)&byte, 1);
                    byte = (unsigned char)f.key[3];
                    klient->sende((char*)&byte, 1);
                }
                if (f.dataLength)
                {
                    for (int i = 0; i < f.dataLength; i++)
                    {
                        if (f.mask)
                            byte = (unsigned char)f.data[i] ^ f.key[i % 4];
                        else
                            byte = (unsigned char)f.data[i];
                        klient->sende((char*)&byte, 1);
                    }
                }
                c2.unlock();
                delete[] f.data;
                c.lock();
                queue->remove(0);
                if (f.opcode == 0x8)
                {
                    if (nextClose)
                    {
                        c.unlock();
                        c2.lock();
                        klient->trenne();
                        klient = (Klient*)klient->release();
                        c2.unlock();
                        return;
                    }
                    else
                    {
                        nextClose = 1;
                    }
                }
            }
            c.unlock();
        }
    }
}

__declspec(dllexport) void WebSocketClient::disconnect()
{
    if (!klient) return;
    c2.lock();
    klient->trenne();
    klient = (Klient*)klient->release();
    c2.unlock();
    warteAufThread(1000);
}

__declspec(dllexport) bool WebSocketClient::isConnected() const
{
    return klient != 0;
}