#include "base.h"

#include <sstream>
#include <strings.h>
#include <unistd.h>
#include <sys/socket.h>
#include <memory>
#include <iostream>
#include <algorithm>
#include <errno.h>
#include <openssl/sha.h>
#include <openssl/evp.h>
#include <openssl/bio.h>
#include <openssl/buffer.h>
#include <set>
#include <mutex>
#include <libwebsockets.h>

// Forward declaration for WebSocket user data
struct WebSocketUserData {
    std::string pendingMessage;
    std::string token;
    std::string userName;
    std::string userColor;
    std::string currentRoom;
    bool authenticated;
};

namespace Yc {
    namespace Lib {

    static std::set<int> g_wsSockets;
    static std::mutex g_wsMutex;

        // Rekursive Entfernung aller "token"-Felder
        static void stripTokensRecursiveImpl(Json::Value &v) {
            if (v.isObject()) {
                if (v.isMember("token")) {
                    v.removeMember("token");
                }
                for (const auto &name : v.getMemberNames()) {
                    stripTokensRecursiveImpl(v[name]);
                }
            } else if (v.isArray()) {
                for (Json::ArrayIndex i = 0; i < v.size(); ++i) {
                    stripTokensRecursiveImpl(v[i]);
                }
            }
        }

        // Öffentliche Helper
        void Base::sanitizeTokens(Json::Value& v) {
            stripTokensRecursiveImpl(v);
        }

        void Base::sanitizeTokensInString(std::string& s) {
            if (s.find("\"token\"") == std::string::npos) return;
            Json::Value v = getJsonTree(s);
            if (!v.isNull()) {
                sanitizeTokens(v);
                s = getJsonString(v);
            }
        }

        std::string Base::getJsonString(const Json::Value& json) {
            std::stringstream outStream;
            outStream << json;
            return outStream.str();
        }

        void Base::send(int socket, std::string out) {
            #ifdef YC_DEBUG
            std::cout << "[Debug] Base::send called with socket: " << socket << ", length: " << out.length() << std::endl;
            #endif
            
            // Token-Felder aus String-Payloads entfernen (falls JSON)
            sanitizeTokensInString(out);
            if (isWebSocket(socket)) {
                #ifdef YC_DEBUG
                std::cout << "[Debug] Socket " << socket << " is WebSocket, using sendWebSocketMessage" << std::endl;
                #endif
                sendWebSocketMessage(socket, out);
                return;
            }
            
            // Socket-Validierung vor dem Senden
            if (socket < 0) {
                #ifdef YC_DEBUG
                std::cout << "[Debug] Invalid socket (< 0), skipping send" << std::endl;
                #endif
                return;
            }
            
            ssize_t result = write(socket, out.c_str(), out.length());
            if (result < 0) {
                int error = errno;
                // Nur kritische Fehler loggen, nicht alle "Bad file descriptor" Fehler
                if (error == EBADF) {
                    #ifdef YC_DEBUG
                    std::cout << "[Debug] Socket " << socket << " closed by client, skipping send" << std::endl;
                    #endif
                } else if (error == EPIPE) {
                    #ifdef YC_DEBUG
                    std::cout << "[Debug] Broken pipe on socket " << socket << ", connection closed" << std::endl;
                    #endif
                } else if (error == ECONNRESET) {
                    #ifdef YC_DEBUG
                    std::cout << "[Debug] Connection reset by peer on socket " << socket << std::endl;
                    #endif
                } else {
                    // Nur unbekannte Fehler als Error loggen
                    std::cout << "[Error] Fehler beim Senden auf Socket " << socket << ": " << strerror(error) << std::endl;
                }
            }
        }

        void Base::send(int socket, Json::Value out) {
            #ifdef YC_DEBUG
            std::cout << "[Debug] Base::send(Json::Value) called with socket: " << socket << std::endl;
            #endif
            
            // Entferne alle Token-Felder rekursiv aus Antworten
            sanitizeTokens(out);
            std::string outString = getJsonString(out);
            send(socket, outString);
        }

        std::string Base::readSocket(int socket)
        {
            if (isWebSocket(socket)) {
                return readWebSocketMessage(socket);
            }
            std::string msg;
            char buffer[1024];
            while (true) {
                ssize_t received = recv(socket, buffer, sizeof(buffer), 0);
                if (received > 0) {
                    msg.append(buffer, static_cast<size_t>(received));
                    if (received < static_cast<ssize_t>(sizeof(buffer))) {
                        break;
                    }
                } else if (received == 0) {
                    break;
                } else {
                    if (errno == EAGAIN || errno == EWOULDBLOCK) {
                        break;
                    }
                    break;
                }
            }
            return msg;
        }

        static std::string base64Encode(const unsigned char* input, size_t len) {
            BIO *bmem = BIO_new(BIO_s_mem());
            BIO *b64 = BIO_new(BIO_f_base64());
            BIO_set_flags(b64, BIO_FLAGS_BASE64_NO_NL);
            b64 = BIO_push(b64, bmem);
            BIO_write(b64, input, (int)len);
            BIO_flush(b64);
            BUF_MEM *bptr;
            BIO_get_mem_ptr(b64, &bptr);
            std::string out(bptr->data, bptr->length);
            BIO_free_all(b64);
            return out;
        }

        static std::string wsAcceptKey(const std::string& clientKey) {
            const std::string magic = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
            std::string in = clientKey + magic;
            unsigned char sha[SHA_DIGEST_LENGTH];
            SHA1(reinterpret_cast<const unsigned char*>(in.data()), in.size(), sha);
            return base64Encode(sha, SHA_DIGEST_LENGTH);
        }

        std::string Base::webSocketAcceptKey(const std::string& clientKey) {
            return wsAcceptKey(clientKey);
        }

        void Base::markWebSocket(int socket) {
            std::lock_guard<std::mutex> lock(g_wsMutex);
            g_wsSockets.insert(socket);
        }
        void Base::unmarkWebSocket(int socket) {
            std::lock_guard<std::mutex> lock(g_wsMutex);
            g_wsSockets.erase(socket);
        }
        bool Base::isWebSocket(int socket) {
            std::lock_guard<std::mutex> lock(g_wsMutex);
            return g_wsSockets.count(socket) > 0;
        }

        std::string Base::readWebSocketMessage(int socket) {
            std::string accumulated;
            bool haveText = false;
            unsigned char firstOpcode = 0;
            while (true) {
                unsigned char hdr[2];
                ssize_t r = recv(socket, hdr, 2, MSG_WAITALL);
                if (r != 2) return std::string();
                bool fin = (hdr[0] & 0x80) != 0;
                unsigned char opcode = hdr[0] & 0x0F;
                bool masked = (hdr[1] & 0x80) != 0;
                uint64_t len = hdr[1] & 0x7F;
                if (len == 126) {
                    unsigned char ext[2];
                    if (recv(socket, ext, 2, MSG_WAITALL) != 2) return {};
                    len = (ext[0] << 8) | ext[1];
                } else if (len == 127) {
                    unsigned char ext[8];
                    if (recv(socket, ext, 8, MSG_WAITALL) != 8) return {};
                    len = 0; for (int i=0;i<8;++i) { len = (len<<8) | ext[i]; }
                }
                unsigned char mask[4] = {0,0,0,0};
                if (masked) {
                    if (recv(socket, mask, 4, MSG_WAITALL) != 4) return {};
                }
                std::string payload; payload.resize((size_t)len);
                if (len > 0) {
                    if (recv(socket, payload.data(), len, MSG_WAITALL) != (ssize_t)len) return {};
                    if (masked) {
                        for (uint64_t i=0;i<len;++i) payload[i] ^= mask[i % 4];
                    }
                }
                if (opcode == 0x8) {
                    // Close frame
                    return std::string();
                } else if (opcode == 0x9) {
                    // Ping -> reply Pong with same payload
                    std::string frame;
                    unsigned char oh[2]; oh[0] = 0x80 | 0xA; // FIN + pong
                    size_t plen = payload.size();
                    if (plen < 126) { oh[1] = (unsigned char)plen; frame.append((char*)oh, 2); }
                    else if (plen <= 0xFFFF) { oh[1] = 126; frame.append((char*)oh, 2); unsigned char ext[2] = {(unsigned char)((plen>>8)&0xFF),(unsigned char)(plen&0xFF)}; frame.append((char*)ext,2);} 
                    else { oh[1] = 127; frame.append((char*)oh,2); unsigned char ext[8]; uint64_t l=plen; for(int i=7;i>=0;--i){ext[i]=(unsigned char)(l&0xFF); l>>=8;} frame.append((char*)ext,8);} 
                    frame.append(payload);
                    ::send(socket, frame.data(), frame.size(), 0);
                    // continue reading next frame
                    continue;
                } else if (opcode == 0xA) {
                    // Pong -> ignore
                    continue;
                } else if (opcode == 0x1 || (opcode == 0x0 && haveText)) {
                    if (!haveText) { firstOpcode = opcode; haveText = true; }
                    accumulated += payload;
                    if (fin) {
                        return accumulated;
                    } else {
                        // read next continuation frame
                        continue;
                    }
                } else {
                    // unsupported/binary -> ignore content, continue
                    if (fin) continue;
                }
            }
        }

        void Base::sendWebSocketMessage(int socket, const std::string& out) {
            #ifdef YC_DEBUG
            std::cout << "[Debug] Base::sendWebSocketMessage called with socket: " << socket << ", length: " << out.length() << std::endl;
            #endif
            
            // Socket-Validierung vor dem Senden
            if (socket < 0) {
                #ifdef YC_DEBUG
                std::cout << "[Debug] Invalid WebSocket socket (< 0), skipping send" << std::endl;
                #endif
                return;
            }
            
            std::string frame;
            unsigned char hdr[2];
            hdr[0] = 0x80 | 0x1; // FIN + text
            size_t len = out.size();
            if (len < 126) {
                hdr[1] = (unsigned char)len;
                frame.append(reinterpret_cast<char*>(hdr), 2);
            } else if (len <= 0xFFFF) {
                hdr[1] = 126;
                frame.append(reinterpret_cast<char*>(hdr), 2);
                unsigned char ext[2] = { (unsigned char)((len>>8)&0xFF), (unsigned char)(len & 0xFF) };
                frame.append(reinterpret_cast<char*>(ext), 2);
            } else {
                hdr[1] = 127;
                frame.append(reinterpret_cast<char*>(hdr), 2);
                unsigned char ext[8];
                uint64_t l = len;
                for (int i=7;i>=0;--i) { ext[i] = (unsigned char)(l & 0xFF); l >>= 8; }
                frame.append(reinterpret_cast<char*>(ext), 8);
            }
            frame.append(out);
            
            ssize_t result = ::send(socket, frame.data(), frame.size(), 0);
            if (result < 0) {
                int error = errno;
                // Nur kritische Fehler loggen, nicht alle "Bad file descriptor" Fehler
                if (error == EBADF) {
                    #ifdef YC_DEBUG
                    std::cout << "[Debug] WebSocket " << socket << " closed by client, skipping send" << std::endl;
                    #endif
                } else if (error == EPIPE) {
                    #ifdef YC_DEBUG
                    std::cout << "[Debug] Broken pipe on WebSocket " << socket << ", connection closed" << std::endl;
                    #endif
                } else if (error == ECONNRESET) {
                    #ifdef YC_DEBUG
                    std::cout << "[Debug] Connection reset by peer on WebSocket " << socket << std::endl;
                    #endif
                } else {
                    // Nur unbekannte Fehler als Error loggen
                    std::cout << "[Error] Fehler beim Senden auf WebSocket " << socket << ": " << strerror(error) << std::endl;
                }
            }
        }

    Json::Value Base::getJsonTree(const std::string& msg) {
            Json::Value inputTree;
            Json::CharReaderBuilder rbuilder;
            std::unique_ptr<Json::CharReader> const reader(rbuilder.newCharReader());
            JSONCPP_STRING inputJsonString(msg);
            reader->parse(inputJsonString.data(), inputJsonString.data() + inputJsonString.size(), &inputTree, NULL);
            return inputTree;
        }

    std::string Base::generateToken() {
        static const char charset[] = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";
        static const size_t charset_size = sizeof(charset) - 1;
        
        std::string token;
        token.reserve(32);
        
        // Einfacher Zufallsgenerator für Token
        for (int i = 0; i < 32; ++i) {
            token += charset[rand() % charset_size];
        }
        
        return token;
    }

    void Base::sendWebSocketMessage(void* wsi, const std::string& out) {
        #ifdef YC_DEBUG
        std::cout << "[Debug] Base::sendWebSocketMessage(void*) called with wsi: " << wsi << ", length: " << out.length() << std::endl;
        #endif
        
        if (!wsi) {
            #ifdef YC_DEBUG
            std::cout << "[Debug] Invalid WebSocket wsi (nullptr), skipping send" << std::endl;
            #endif
            return;
        }
        
        // Cast to lws* and use libwebsockets API
        struct lws* lws_wsi = static_cast<struct lws*>(wsi);
        
        // Store message in user data for sending
        auto* ud = reinterpret_cast<WebSocketUserData*>(lws_wsi_user(lws_wsi));
        if (ud) {
            ud->pendingMessage = out;
            lws_callback_on_writable(lws_wsi);
        } else {
            #ifdef YC_DEBUG
            std::cout << "[Debug] No user data found for WebSocket, cannot send message" << std::endl;
            #endif
        }
    }

    void Base::sendWebSocketMessage(void* wsi, const Json::Value& out) {
        #ifdef YC_DEBUG
        std::cout << "[Debug] Base::sendWebSocketMessage(void*, Json::Value) called with wsi: " << wsi << std::endl;
        #endif
        
        // Convert JSON to string and call the string version
        std::string outString = getJsonString(out);
        sendWebSocketMessage(wsi, outString);
    }

    } // namespace Lib
} // namespace Yc