#ifndef _WIN32_WINNT
#define _WIN32_WINNT 0x0602
#endif
#ifndef WIN32_LEAN_AND_MEAN
#define WIN32_LEAN_AND_MEAN
#endif
#ifndef NOMINMAX
#define NOMINMAX
#endif

#include "AiWebSocketClient.h"
#include <winhttp.h>
#include <sstream>
#include <cstring>

#pragma comment(lib, "winhttp.lib")

static void AiAppendUtf8(std::vector<unsigned char>& out, const std::string& text)
{
    out.insert(out.end(), text.begin(), text.end());
}

static std::string AiJsonEscape(const std::string& s)
{
    std::string out;
    out.reserve(s.size() + 16);
    for (char c : s) {
        switch (c) {
        case '\\': out += "\\\\"; break;
        case '"': out += "\\\""; break;
        case '\n': out += "\\n"; break;
        case '\r': out += "\\r"; break;
        case '\t': out += "\\t"; break;
        default: out += c; break;
        }
    }
    return out;
}

static bool AiStartsWithI(const std::wstring& s, const wchar_t* prefix)
{
    size_t n = wcslen(prefix);
    if (s.size() < n) return false;
    return _wcsnicmp(s.c_str(), prefix, n) == 0;
}

static bool AiCrackWsUrl(const std::wstring& url,
                         std::wstring* host,
                         std::wstring* path,
                         INTERNET_PORT* port,
                         bool* secure)
{
    if (!host || !path || !port || !secure) return false;
    host->clear();
    path->clear();
    *port = 0;
    *secure = false;

    size_t schemeLen = 0;
    if (AiStartsWithI(url, L"ws://")) {
        schemeLen = 5;
        *secure = false;
        *port = 80;
    } else if (AiStartsWithI(url, L"wss://")) {
        schemeLen = 6;
        *secure = true;
        *port = 443;
    } else if (AiStartsWithI(url, L"http://")) {
        schemeLen = 7;
        *secure = false;
        *port = 80;
    } else if (AiStartsWithI(url, L"https://")) {
        schemeLen = 8;
        *secure = true;
        *port = 443;
    } else {
        return false;
    }

    size_t hostStart = schemeLen;
    size_t pathStart = url.find(L'/', hostStart);
    std::wstring hostPort = (pathStart == std::wstring::npos) ? url.substr(hostStart) : url.substr(hostStart, pathStart - hostStart);
    *path = (pathStart == std::wstring::npos) ? L"/" : url.substr(pathStart);
    if (hostPort.empty()) return false;

    // IPv6 literal support is intentionally limited for this local demo.
    size_t colon = hostPort.rfind(L':');
    if (colon != std::wstring::npos && colon + 1 < hostPort.size()) {
        *host = hostPort.substr(0, colon);
        int p = _wtoi(hostPort.substr(colon + 1).c_str());
        if (p > 0 && p <= 65535) *port = (INTERNET_PORT)p;
    } else {
        *host = hostPort;
    }
    return !host->empty() && !path->empty();
}

static std::wstring AiLastErrorText(const wchar_t* prefix)
{
    DWORD err = GetLastError();
    std::wstringstream ss;
    ss << prefix << L" GetLastError=" << err;
    return ss.str();
}

AiWebSocketClient::AiWebSocketClient()
    : session_(NULL), connect_(NULL), socket_(NULL)
{
}

AiWebSocketClient::~AiWebSocketClient()
{
    Close();
}

bool AiWebSocketClient::IsConnected() const
{
    return socket_ != NULL;
}

void AiWebSocketClient::Close()
{
    if (socket_) {
        WinHttpWebSocketClose((HINTERNET)socket_, WINHTTP_WEB_SOCKET_SUCCESS_CLOSE_STATUS, NULL, 0);
        WinHttpCloseHandle((HINTERNET)socket_);
        socket_ = NULL;
    }
    if (connect_) {
        WinHttpCloseHandle((HINTERNET)connect_);
        connect_ = NULL;
    }
    if (session_) {
        WinHttpCloseHandle((HINTERNET)session_);
        session_ = NULL;
    }
    endpoint_.clear();
}

bool AiWebSocketClient::Connect(const std::wstring& endpoint, std::wstring* errorText)
{
    Close();

    std::wstring host, path;
    INTERNET_PORT port = 0;
    bool secure = false;
    if (!AiCrackWsUrl(endpoint, &host, &path, &port, &secure)) {
        if (errorText) *errorText = L"Invalid WebSocket URL.";
        return false;
    }

    HINTERNET session = NULL;
    HINTERNET connect = NULL;
    HINTERNET request = NULL;
    HINTERNET socket = NULL;

    session = WinHttpOpen(L"HWS AI WebSocket Transport/1.0",
                          WINHTTP_ACCESS_TYPE_DEFAULT_PROXY,
                          WINHTTP_NO_PROXY_NAME,
                          WINHTTP_NO_PROXY_BYPASS,
                          0);
    if (!session) {
        if (errorText) *errorText = AiLastErrorText(L"WinHttpOpen failed.");
        return false;
    }
    WinHttpSetTimeouts(session, 3000, 3000, 5000, 10000);

    connect = WinHttpConnect(session, host.c_str(), port, 0);
    if (!connect) {
        if (errorText) *errorText = AiLastErrorText(L"WinHttpConnect failed.");
        WinHttpCloseHandle(session);
        return false;
    }

    DWORD flags = secure ? WINHTTP_FLAG_SECURE : 0;
    request = WinHttpOpenRequest(connect, L"GET", path.c_str(), NULL,
                                 WINHTTP_NO_REFERER, WINHTTP_DEFAULT_ACCEPT_TYPES, flags);
    if (!request) {
        if (errorText) *errorText = AiLastErrorText(L"WinHttpOpenRequest failed.");
        WinHttpCloseHandle(connect);
        WinHttpCloseHandle(session);
        return false;
    }

    BOOL opt = WinHttpSetOption(request, WINHTTP_OPTION_UPGRADE_TO_WEB_SOCKET, NULL, 0);
    if (!opt) {
        if (errorText) *errorText = AiLastErrorText(L"WINHTTP_OPTION_UPGRADE_TO_WEB_SOCKET failed.");
        WinHttpCloseHandle(request);
        WinHttpCloseHandle(connect);
        WinHttpCloseHandle(session);
        return false;
    }

    BOOL sent = WinHttpSendRequest(request,
                                   WINHTTP_NO_ADDITIONAL_HEADERS, 0,
                                   WINHTTP_NO_REQUEST_DATA, 0,
                                   0, 0);
    if (!sent || !WinHttpReceiveResponse(request, NULL)) {
        if (errorText) *errorText = AiLastErrorText(L"WebSocket upgrade request failed.");
        WinHttpCloseHandle(request);
        WinHttpCloseHandle(connect);
        WinHttpCloseHandle(session);
        return false;
    }

    socket = WinHttpWebSocketCompleteUpgrade(request, 0);
    WinHttpCloseHandle(request);
    if (!socket) {
        if (errorText) *errorText = AiLastErrorText(L"WinHttpWebSocketCompleteUpgrade failed.");
        WinHttpCloseHandle(connect);
        WinHttpCloseHandle(session);
        return false;
    }

    session_ = session;
    connect_ = connect;
    socket_ = socket;
    endpoint_ = endpoint;
    return true;
}

bool AiWebSocketClient::SendBinary(const std::vector<unsigned char>& payload, std::wstring* errorText)
{
    if (!socket_) {
        if (errorText) *errorText = L"WebSocket is not connected.";
        return false;
    }
    DWORD err = WinHttpWebSocketSend((HINTERNET)socket_,
                                     WINHTTP_WEB_SOCKET_BINARY_MESSAGE_BUFFER_TYPE,
                                     payload.empty() ? NULL : (void*)payload.data(),
                                     (DWORD)payload.size());
    if (err != ERROR_SUCCESS) {
        std::wstringstream ss;
        ss << L"WinHttpWebSocketSend failed. error=" << err;
        if (errorText) *errorText = ss.str();
        Close();
        return false;
    }
    return true;
}

bool AiWebSocketClient::ReceiveText(std::string* responseUtf8, std::wstring* errorText)
{
    if (responseUtf8) responseUtf8->clear();
    if (!socket_) {
        if (errorText) *errorText = L"WebSocket is not connected.";
        return false;
    }

    std::string out;
    std::vector<unsigned char> buffer(64 * 1024);
    for (;;) {
        DWORD bytesRead = 0;
        WINHTTP_WEB_SOCKET_BUFFER_TYPE type = WINHTTP_WEB_SOCKET_CLOSE_BUFFER_TYPE;
        DWORD err = WinHttpWebSocketReceive((HINTERNET)socket_, buffer.data(), (DWORD)buffer.size(), &bytesRead, &type);
        if (err != ERROR_SUCCESS) {
            std::wstringstream ss;
            ss << L"WinHttpWebSocketReceive failed. error=" << err;
            if (errorText) *errorText = ss.str();
            Close();
            return false;
        }

        if (type == WINHTTP_WEB_SOCKET_CLOSE_BUFFER_TYPE) {
            if (errorText) *errorText = L"WebSocket closed by AI service.";
            Close();
            return false;
        }

        if (bytesRead > 0) {
            out.append((const char*)buffer.data(), (size_t)bytesRead);
        }

        if (type == WINHTTP_WEB_SOCKET_UTF8_MESSAGE_BUFFER_TYPE ||
            type == WINHTTP_WEB_SOCKET_BINARY_MESSAGE_BUFFER_TYPE) {
            break;
        }

        if (type != WINHTTP_WEB_SOCKET_UTF8_FRAGMENT_BUFFER_TYPE &&
            type != WINHTTP_WEB_SOCKET_BINARY_FRAGMENT_BUFFER_TYPE) {
            break;
        }
    }

    if (responseUtf8) *responseUtf8 = out;
    return !out.empty();
}

bool AiWebSocketClient::SendAnalyzeFrame(int slot,
                                         long long timestamp100ns,
                                         long long frameIndex,
                                         int width,
                                         int height,
                                         const std::vector<unsigned char>& jpeg,
                                         const std::string& modeUtf8,
                                         const std::string& promptUtf8,
                                         std::string* responseUtf8,
                                         std::wstring* errorText)
{
    if (jpeg.empty()) {
        if (errorText) *errorText = L"No JPEG data.";
        return false;
    }

    std::ostringstream hs;
    hs << "{"
       << "\"type\":\"video.jpeg_frame\"," 
       << "\"slot\":" << slot << ","
       << "\"timestamp100ns\":" << timestamp100ns << ","
       << "\"frameIndex\":" << frameIndex << ","
       << "\"width\":" << width << ","
       << "\"height\":" << height << ","
       << "\"format\":\"jpeg\"," 
       << "\"mode\":\"" << AiJsonEscape(modeUtf8) << "\"," 
       << "\"prompt\":\"" << AiJsonEscape(promptUtf8) << "\""
       << "}";
    std::string header = hs.str();

    if (header.size() > 1024 * 1024) {
        if (errorText) *errorText = L"WebSocket header too large.";
        return false;
    }

    std::vector<unsigned char> payload;
    payload.reserve(4 + header.size() + jpeg.size());
    unsigned int headerSize = (unsigned int)header.size();
    payload.push_back((unsigned char)(headerSize & 0xFF));
    payload.push_back((unsigned char)((headerSize >> 8) & 0xFF));
    payload.push_back((unsigned char)((headerSize >> 16) & 0xFF));
    payload.push_back((unsigned char)((headerSize >> 24) & 0xFF));
    AiAppendUtf8(payload, header);
    payload.insert(payload.end(), jpeg.begin(), jpeg.end());

    if (!SendBinary(payload, errorText)) return false;
    return ReceiveText(responseUtf8, errorText);
}

bool AiWebSocketClient::SendAudioChunk(int audioSource,
                                       long long timestamp100ns,
                                       long long chunkIndex,
                                       int sampleRate,
                                       int channels,
                                       int bitsPerSample,
                                       int sampleCount,
                                       const std::vector<unsigned char>& pcm16,
                                       const std::string& srcLangUtf8,
                                       const std::string& dstLangUtf8,
                                       std::string* responseUtf8,
                                       std::wstring* errorText)
{
    if (pcm16.empty()) {
        if (errorText) *errorText = L"No PCM audio data.";
        return false;
    }

    std::ostringstream hs;
    hs << "{"
       << "\"type\":\"audio.pcm_chunk\"," 
       << "\"audioSource\":" << audioSource << ","
       << "\"timestamp100ns\":" << timestamp100ns << ","
       << "\"chunkIndex\":" << chunkIndex << ","
       << "\"sampleRate\":" << sampleRate << ","
       << "\"channels\":" << channels << ","
       << "\"bitsPerSample\":" << bitsPerSample << ","
       << "\"sampleCount\":" << sampleCount << ","
       << "\"format\":\"pcm_s16le\"," 
       << "\"srcLang\":\"" << AiJsonEscape(srcLangUtf8) << "\"," 
       << "\"dstLang\":\"" << AiJsonEscape(dstLangUtf8) << "\""
       << "}";
    std::string header = hs.str();

    if (header.size() > 1024 * 1024) {
        if (errorText) *errorText = L"WebSocket audio header too large.";
        return false;
    }

    std::vector<unsigned char> payload;
    payload.reserve(4 + header.size() + pcm16.size());
    unsigned int headerSize = (unsigned int)header.size();
    payload.push_back((unsigned char)(headerSize & 0xFF));
    payload.push_back((unsigned char)((headerSize >> 8) & 0xFF));
    payload.push_back((unsigned char)((headerSize >> 16) & 0xFF));
    payload.push_back((unsigned char)((headerSize >> 24) & 0xFF));
    AiAppendUtf8(payload, header);
    payload.insert(payload.end(), pcm16.begin(), pcm16.end());

    if (!SendBinary(payload, errorText)) return false;
    return ReceiveText(responseUtf8, errorText);
}
