#include "cdp_ws.h"

#include "crude_json.h"

#include <chrono>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <random>
#include <sstream>

#ifdef _WIN32
#  include <ws2tcpip.h>
#endif

namespace navegator {

namespace {

// Base64 encoder (raw).
std::string b64encode(const uint8_t* data, size_t len) {
    static const char* tbl =
        "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    std::string out;
    out.reserve(((len + 2) / 3) * 4);
    for (size_t i = 0; i < len; i += 3) {
        uint32_t v = (uint32_t)data[i] << 16;
        if (i + 1 < len) v |= (uint32_t)data[i + 1] << 8;
        if (i + 2 < len) v |= (uint32_t)data[i + 2];
        out.push_back(tbl[(v >> 18) & 0x3F]);
        out.push_back(tbl[(v >> 12) & 0x3F]);
        out.push_back(i + 1 < len ? tbl[(v >> 6) & 0x3F] : '=');
        out.push_back(i + 2 < len ? tbl[v & 0x3F] : '=');
    }
    return out;
}

std::string make_ws_key() {
    uint8_t buf[16];
    std::random_device rd;
    for (int i = 0; i < 16; i += 2) {
        uint32_t v = rd();
        buf[i]     = (uint8_t)(v & 0xFF);
        buf[i + 1] = (uint8_t)((v >> 8) & 0xFF);
    }
    return b64encode(buf, 16);
}

#ifdef _WIN32
bool socket_set_timeout(SOCKET s, int ms) {
    DWORD t = (DWORD)ms;
    return setsockopt(s, SOL_SOCKET, SO_RCVTIMEO, (const char*)&t, sizeof(t)) == 0
        && setsockopt(s, SOL_SOCKET, SO_SNDTIMEO, (const char*)&t, sizeof(t)) == 0;
}
#endif

// CDP envia siempre JSON UTF-8. Json escape rapido.
std::string json_escape_str(const std::string& s) {
    std::string out;
    out.reserve(s.size() + 2);
    out.push_back('"');
    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:
                if ((unsigned char)c < 0x20) {
                    char b[8];
                    std::snprintf(b, sizeof(b), "\\u%04x", (unsigned)(uint8_t)c);
                    out += b;
                } else out.push_back(c);
        }
    }
    out.push_back('"');
    return out;
}

} // namespace

bool CdpWs::parse_ws_url(const std::string& url, std::string& host, int& port, std::string& path) {
    const std::string scheme = "ws://";
    if (url.compare(0, scheme.size(), scheme) != 0) return false;
    std::string rest = url.substr(scheme.size());
    size_t slash = rest.find('/');
    std::string authority = (slash == std::string::npos) ? rest : rest.substr(0, slash);
    path = (slash == std::string::npos) ? "/" : rest.substr(slash);
    size_t colon = authority.find(':');
    if (colon == std::string::npos) {
        host = authority;
        port = 80;
    } else {
        host = authority.substr(0, colon);
        try { port = std::stoi(authority.substr(colon + 1)); }
        catch (...) { return false; }
    }
    return true;
}

CdpWs::~CdpWs() { close(); }

void CdpWs::set_error(const std::string& e) {
    std::lock_guard<std::mutex> lk(err_mu_);
    last_err_ = e;
}

#ifdef _WIN32

bool CdpWs::send_all(const char* data, size_t len) {
    size_t sent = 0;
    while (sent < len) {
        int n = ::send(sock_, data + sent, (int)(len - sent), 0);
        if (n <= 0) return false;
        sent += (size_t)n;
        bytes_out_.fetch_add((uint64_t)n);
    }
    return true;
}

bool CdpWs::recv_n(char* out, size_t n) {
    size_t got = 0;
    while (got < n) {
        int r = ::recv(sock_, out + got, (int)(n - got), 0);
        if (r <= 0) return false;
        got += (size_t)r;
        bytes_in_.fetch_add((uint64_t)r);
    }
    return true;
}

bool CdpWs::connect(const CdpWsConfig& cfg, std::string* err) {
    if (running_.load()) return true;

    WSADATA wsa;
    static bool wsa_ok = (WSAStartup(MAKEWORD(2, 2), &wsa) == 0);
    if (!wsa_ok) { if (err) *err = "WSAStartup failed"; return false; }

    sock_ = ::socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (sock_ == INVALID_SOCKET) { if (err) *err = "socket() failed"; return false; }
    socket_set_timeout(sock_, cfg.timeout_ms);

    sockaddr_in addr{};
    addr.sin_family = AF_INET;
    addr.sin_port = htons((u_short)cfg.port);
    if (cfg.host == "127.0.0.1" || cfg.host == "localhost") {
        addr.sin_addr.s_addr = htonl(0x7F000001);
    } else {
        // Generico (LAN) — gethostbyname.
        addrinfo hints{}; addrinfo* res = nullptr;
        hints.ai_family = AF_INET;
        hints.ai_socktype = SOCK_STREAM;
        if (getaddrinfo(cfg.host.c_str(), nullptr, &hints, &res) != 0 || !res) {
            if (err) *err = "getaddrinfo failed";
            closesocket(sock_); sock_ = INVALID_SOCKET; return false;
        }
        addr.sin_addr = ((sockaddr_in*)res->ai_addr)->sin_addr;
        freeaddrinfo(res);
    }
    if (::connect(sock_, (sockaddr*)&addr, sizeof(addr)) != 0) {
        if (err) *err = "tcp connect failed";
        closesocket(sock_); sock_ = INVALID_SOCKET; return false;
    }

    std::string key = make_ws_key();
    std::ostringstream req;
    req << "GET " << cfg.path << " HTTP/1.1\r\n"
        << "Host: " << cfg.host << ":" << cfg.port << "\r\n"
        << "Upgrade: websocket\r\n"
        << "Connection: Upgrade\r\n"
        << "Sec-WebSocket-Key: " << key << "\r\n"
        << "Sec-WebSocket-Version: 13\r\n"
        << "Origin: http://localhost\r\n"
        << "\r\n";
    std::string raw_req = req.str();
    if (!send_all(raw_req.data(), raw_req.size())) {
        if (err) *err = "ws upgrade send failed";
        closesocket(sock_); sock_ = INVALID_SOCKET; return false;
    }

    // Leer respuesta hasta \r\n\r\n. recv pequeño porque solo cabeceras.
    std::string resp;
    char buf[1024];
    while (resp.find("\r\n\r\n") == std::string::npos && resp.size() < 8192) {
        int n = ::recv(sock_, buf, sizeof(buf), 0);
        if (n <= 0) {
            if (err) *err = "ws upgrade recv failed";
            closesocket(sock_); sock_ = INVALID_SOCKET; return false;
        }
        resp.append(buf, (size_t)n);
        bytes_in_.fetch_add((uint64_t)n);
    }
    auto sp1 = resp.find(' ');
    int status = 0;
    if (sp1 != std::string::npos) {
        try { status = std::stoi(resp.substr(sp1 + 1, 3)); } catch (...) {}
    }
    if (status != 101) {
        if (err) *err = "ws upgrade not 101 (status " + std::to_string(status) + ")";
        closesocket(sock_); sock_ = INVALID_SOCKET; return false;
    }

    running_.store(true);
    stop_.store(false);
    reader_ = std::thread(&CdpWs::reader_loop, this);
    return true;
}

void CdpWs::close() {
    if (!running_.exchange(false)) return;
    stop_.store(true);
    send_close_frame();
    if (sock_ != INVALID_SOCKET) {
        shutdown(sock_, SD_BOTH);
        closesocket(sock_);
        sock_ = INVALID_SOCKET;
    }
    if (reader_.joinable()) reader_.join();
    resp_cv_.notify_all();
}

bool CdpWs::send_frame_text(const std::string& payload) {
    std::lock_guard<std::mutex> lk(send_mu_);
    if (!running_.load()) return false;

    uint8_t header[14];
    size_t  hlen = 0;
    header[hlen++] = 0x81; // FIN=1, opcode=1 (text)

    size_t plen = payload.size();
    if (plen < 126) {
        header[hlen++] = (uint8_t)(0x80 | plen);
    } else if (plen <= 0xFFFF) {
        header[hlen++] = (uint8_t)(0x80 | 126);
        header[hlen++] = (uint8_t)((plen >> 8) & 0xFF);
        header[hlen++] = (uint8_t)(plen & 0xFF);
    } else {
        header[hlen++] = (uint8_t)(0x80 | 127);
        for (int i = 7; i >= 0; --i) header[hlen++] = (uint8_t)((plen >> (i * 8)) & 0xFF);
    }
    uint8_t mask[4];
    std::random_device rd;
    uint32_t mr = rd();
    mask[0] = (uint8_t)(mr & 0xFF);
    mask[1] = (uint8_t)((mr >> 8) & 0xFF);
    mask[2] = (uint8_t)((mr >> 16) & 0xFF);
    mask[3] = (uint8_t)((mr >> 24) & 0xFF);
    std::memcpy(header + hlen, mask, 4); hlen += 4;

    if (!send_all((const char*)header, hlen)) return false;

    // Enmascarar payload en bloques.
    std::string masked(plen, '\0');
    for (size_t i = 0; i < plen; ++i) masked[i] = payload[i] ^ mask[i & 3];
    return send_all(masked.data(), plen);
}

bool CdpWs::send_close_frame() {
    if (sock_ == INVALID_SOCKET) return false;
    std::lock_guard<std::mutex> lk(send_mu_);
    uint8_t f[6] = { 0x88, 0x80, 0, 0, 0, 0 }; // close, masked, len 0
    return ::send(sock_, (const char*)f, 6, 0) > 0;
}

bool CdpWs::recv_frame(uint8_t& opcode, std::string& payload) {
    char hdr2[2];
    if (!recv_n(hdr2, 2)) return false;
    bool fin     = (hdr2[0] & 0x80) != 0;
    opcode       = (uint8_t)(hdr2[0] & 0x0F);
    bool masked  = (hdr2[1] & 0x80) != 0;
    uint64_t len = (uint64_t)(hdr2[1] & 0x7F);
    if (len == 126) {
        char lb[2];
        if (!recv_n(lb, 2)) return false;
        len = ((uint64_t)(uint8_t)lb[0] << 8) | (uint8_t)lb[1];
    } else if (len == 127) {
        char lb[8];
        if (!recv_n(lb, 8)) return false;
        len = 0;
        for (int i = 0; i < 8; ++i) len = (len << 8) | (uint8_t)lb[i];
    }
    uint8_t mask[4] = {0,0,0,0};
    if (masked) {
        if (!recv_n((char*)mask, 4)) return false;
    }
    payload.assign((size_t)len, '\0');
    if (len > 0) {
        if (!recv_n(payload.data(), (size_t)len)) return false;
        if (masked) {
            for (size_t i = 0; i < (size_t)len; ++i) payload[i] = (char)((uint8_t)payload[i] ^ mask[i & 3]);
        }
    }
    if (!fin) {
        // Continuation: append until FIN. CDP rara vez fragmenta pero por compat.
        std::string tail;
        uint8_t op2 = 0;
        while (true) {
            char h2[2];
            if (!recv_n(h2, 2)) return false;
            bool fin2     = (h2[0] & 0x80) != 0;
            op2           = (uint8_t)(h2[0] & 0x0F);
            bool masked2  = (h2[1] & 0x80) != 0;
            uint64_t len2 = (uint64_t)(h2[1] & 0x7F);
            if (len2 == 126) {
                char lb[2];
                if (!recv_n(lb, 2)) return false;
                len2 = ((uint64_t)(uint8_t)lb[0] << 8) | (uint8_t)lb[1];
            } else if (len2 == 127) {
                char lb[8];
                if (!recv_n(lb, 8)) return false;
                len2 = 0;
                for (int i = 0; i < 8; ++i) len2 = (len2 << 8) | (uint8_t)lb[i];
            }
            uint8_t m2[4] = {0,0,0,0};
            if (masked2 && !recv_n((char*)m2, 4)) return false;
            std::string p2((size_t)len2, '\0');
            if (len2 > 0) {
                if (!recv_n(p2.data(), (size_t)len2)) return false;
                if (masked2) for (size_t i = 0; i < (size_t)len2; ++i) p2[i] = (char)((uint8_t)p2[i] ^ m2[i & 3]);
            }
            payload.append(p2);
            if (fin2) break;
        }
    }
    return true;
}

void CdpWs::reader_loop() {
    while (running_.load() && !stop_.load()) {
        uint8_t opcode = 0;
        std::string payload;
        if (!recv_frame(opcode, payload)) {
            if (running_.load()) set_error("recv_frame failed");
            break;
        }
        frames_in_.fetch_add(1);

        if (opcode == 0x1) {
            // Text frame. Si trae "id":N y "result|error", suelta wait_response.
            // Hacemos parse parcial barato: buscar substring '"id":'. Si esta,
            // parsear id y guardar.
            int found_id = -1;
            auto idp = payload.find("\"id\":");
            if (idp != std::string::npos) {
                size_t p = idp + 5;
                while (p < payload.size() && (payload[p] == ' ' || payload[p] == '\t')) ++p;
                int sign = 1;
                if (p < payload.size() && payload[p] == '-') { sign = -1; ++p; }
                int v = 0; bool any = false;
                while (p < payload.size() && payload[p] >= '0' && payload[p] <= '9') {
                    v = v * 10 + (payload[p] - '0'); ++p; any = true;
                }
                if (any) found_id = sign * v;
            }
            if (found_id >= 0 &&
                (payload.find("\"result\"") != std::string::npos ||
                 payload.find("\"error\"")  != std::string::npos)) {
                {
                    std::lock_guard<std::mutex> lk(resp_mu_);
                    responses_[found_id] = payload;
                }
                resp_cv_.notify_all();
            }
            // Empujar a la cola pase lo que pase: el panel decide si filtrar.
            {
                std::lock_guard<std::mutex> lk(queue_mu_);
                if (queue_.size() < 100000) queue_.push(std::move(payload));
            }
        } else if (opcode == 0x9) {
            // Ping -> pong.
            std::lock_guard<std::mutex> lk(send_mu_);
            uint8_t pong[6] = { 0x8A, 0x80, 0, 0, 0, 0 };
            ::send(sock_, (const char*)pong, 6, 0);
        } else if (opcode == 0x8) {
            // Close: termina.
            break;
        }
        // opcode 0xA (pong) o binary: ignorar.
    }
    running_.store(false);
}

#else // !_WIN32 stubs

bool CdpWs::send_all(const char*, size_t) { return false; }
bool CdpWs::recv_n(char*, size_t) { return false; }
bool CdpWs::connect(const CdpWsConfig&, std::string* err) { if (err) *err = "Windows-only"; return false; }
void CdpWs::close() {}
bool CdpWs::send_frame_text(const std::string&) { return false; }
bool CdpWs::send_close_frame() { return false; }
bool CdpWs::recv_frame(uint8_t&, std::string&) { return false; }
void CdpWs::reader_loop() {}

#endif

int CdpWs::send_command(const std::string& method, const std::string& params_json) {
    if (!running_.load()) return -1;
    int id = next_id_.fetch_add(1);
    std::ostringstream os;
    os << "{\"id\":" << id << ",\"method\":" << json_escape_str(method);
    if (!params_json.empty()) {
        os << ",\"params\":" << params_json;
    }
    os << "}";
    if (!send_frame_text(os.str())) return -1;
    return id;
}

std::vector<std::string> CdpWs::drain(size_t max) {
    std::vector<std::string> out;
    std::lock_guard<std::mutex> lk(queue_mu_);
    while (!queue_.empty() && out.size() < max) {
        out.push_back(std::move(queue_.front()));
        queue_.pop();
    }
    return out;
}

bool CdpWs::wait_response(int id, std::string& out_json, int timeout_ms) {
    std::unique_lock<std::mutex> lk(resp_mu_);
    auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(timeout_ms);
    while (true) {
        auto it = responses_.find(id);
        if (it != responses_.end()) {
            out_json = std::move(it->second);
            responses_.erase(it);
            return true;
        }
        if (!running_.load()) return false;
        if (resp_cv_.wait_until(lk, deadline) == std::cv_status::timeout) return false;
    }
}

} // namespace navegator