/**
 * @file ws_client.cpp
 * @brief WebSocket 客户端实现 - 轻量级版本
 */

#include "ws_client.h"
#include "../utils/logger.h"
#include <cstring>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <fcntl.h>
#include <openssl/sha.h>
#include <openssl/bio.h>
#include <openssl/evp.h>
#include <openssl/buffer.h>

// WebSocket 魔术字符串（用于握手验证，当前未使用）
// static const char* WS_MAGIC_STRING = "258EAFAA-5914-47DA-95CA-C5AB0DC85B11";

WSClient::WSClient(const std::string& host, int port, const std::string& path)
    : m_host(host), m_port(port), m_path(path),
      m_sockfd(-1), m_connected(false), m_running(false) {
}

WSClient::~WSClient() {
    disconnect();
}

bool WSClient::connect() {
    if (m_connected) return true;
    
    // 创建 socket
    m_sockfd = socket(AF_INET, SOCK_STREAM, 0);
    if (m_sockfd < 0) {
        LOG_ERROR("[WS] Failed to create socket: %s", strerror(errno));
        return false;
    }
    
    // 解析主机地址
    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(m_port);
    
    // 尝试直接解析 IP
    if (inet_pton(AF_INET, m_host.c_str(), &server_addr.sin_addr) <= 0) {
        // 不是IP地址，尝试 DNS 解析
        struct hostent* he = gethostbyname(m_host.c_str());
        if (!he) {
            LOG_ERROR("[WS] Failed to resolve host: %s", m_host.c_str());
            close(m_sockfd);
            m_sockfd = -1;
            return false;
        }
        memcpy(&server_addr.sin_addr, he->h_addr_list[0], he->h_length);
    }
    
    // 连接到服务器
    if (::connect(m_sockfd, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
        LOG_ERROR("[WS] Failed to connect to %s:%d: %s", 
                  m_host.c_str(), m_port, strerror(errno));
        close(m_sockfd);
        m_sockfd = -1;
        return false;
    }
    
    // 增大发送缓冲区到 256KB（默认通常只有 16KB）
    int sndbuf = 256 * 1024;
    if (setsockopt(m_sockfd, SOL_SOCKET, SO_SNDBUF, &sndbuf, sizeof(sndbuf)) < 0) {
        LOG_WARN("[WS] Failed to set send buffer size: %s", strerror(errno));
    }
    
    // 增大接收缓冲区到 256KB
    int rcvbuf = 256 * 1024;
    if (setsockopt(m_sockfd, SOL_SOCKET, SO_RCVBUF, &rcvbuf, sizeof(rcvbuf)) < 0) {
        LOG_WARN("[WS] Failed to set recv buffer size: %s", strerror(errno));
    }
    
    LOG_INFO("[WS] TCP connected to %s:%d (buffers: 256KB)", m_host.c_str(), m_port);
    
    // 执行 WebSocket 握手
    if (!perform_handshake()) {
        close(m_sockfd);
        m_sockfd = -1;
        return false;
    }
    
    m_connected = true;
    m_running = true;
    
    // 修复: 在创建新线程前，必须确保旧线程已被 join，否则会触发 std::terminate
    if (m_recv_thread.joinable()) {
        m_recv_thread.join();
    }
    
    // 启动接收线程
    m_recv_thread = std::thread(&WSClient::recv_loop, this);
    
    LOG_INFO("[WS] WebSocket connected to ws://%s:%d%s", 
             m_host.c_str(), m_port, m_path.c_str());
    
    return true;
}

void WSClient::disconnect() {
    if (!m_connected && m_sockfd < 0) return;
    
    m_running = false;
    m_connected = false;
    
    // 发送关闭帧
    if (m_sockfd >= 0) {
        uint8_t close_frame[] = {0x88, 0x00}; // FIN + CLOSE, 无payload
        send(m_sockfd, close_frame, sizeof(close_frame), 0);
        
        // 关闭 socket 读写以中断任何阻塞的 recv()/send()
        shutdown(m_sockfd, SHUT_RDWR);
    }
    
    // 等待接收线程退出
    if (m_recv_thread.joinable()) {
        m_recv_thread.join();
    }
    
    if (m_sockfd >= 0) {
        close(m_sockfd);
        m_sockfd = -1;
    }
    
    LOG_INFO("[WS] Disconnected");
}

bool WSClient::send_text(const std::string& text) {
    if (!m_connected) return false;
    return send_frame((const uint8_t*)text.c_str(), text.length(), OP_TEXT);
}

bool WSClient::send_binary(const uint8_t* data, size_t size) {
    if (!m_connected) return false;
    
    // Day 12: 检查是否需要发送心跳
    time_t now = time(nullptr);
    if (m_last_send_time > 0 && (now - m_last_send_time) > HEARTBEAT_INTERVAL) {
        send_ping();
    }
    
    return send_frame(data, size, OP_BINARY);
}

void WSClient::poll_messages(MessageCallback callback) {
    std::lock_guard<std::mutex> lock(m_recv_mutex);
    
    while (!m_recv_queue.empty()) {
        callback(m_recv_queue.front());
        m_recv_queue.pop();
    }
}

void WSClient::poll_binary_messages(std::function<void(const uint8_t*, size_t)> callback) {
    std::lock_guard<std::mutex> lock(m_binary_mutex);
    
    while (!m_binary_queue.empty()) {
        const std::vector<uint8_t>& data = m_binary_queue.front();
        callback(data.data(), data.size());
        m_binary_queue.pop();
    }
}

// ===== 内部实现 =====

bool WSClient::perform_handshake() {
    // 生成随机 Sec-WebSocket-Key (base64编码的16字节随机数)
    uint8_t random_bytes[16];
    for (int i = 0; i < 16; ++i) {
        random_bytes[i] = rand() % 256;
    }
    
    // Base64 编码
    BIO *bio, *b64;
    BUF_MEM *bufferPtr;
    
    b64 = BIO_new(BIO_f_base64());
    bio = BIO_new(BIO_s_mem());
    bio = BIO_push(b64, bio);
    
    BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);
    BIO_write(bio, random_bytes, 16);
    BIO_flush(bio);
    BIO_get_mem_ptr(bio, &bufferPtr);
    
    std::string ws_key(bufferPtr->data, bufferPtr->length);
    BIO_free_all(bio);
    
    // 构造 HTTP 升级请求
    char request[1024];
    snprintf(request, sizeof(request),
        "GET %s HTTP/1.1\r\n"
        "Host: %s:%d\r\n"
        "Upgrade: websocket\r\n"
        "Connection: Upgrade\r\n"
        "Sec-WebSocket-Key: %s\r\n"
        "Sec-WebSocket-Version: 13\r\n"
        "\r\n",
        m_path.c_str(), m_host.c_str(), m_port, ws_key.c_str());
    
    // 发送请求
    if (send(m_sockfd, request, strlen(request), 0) < 0) {
        LOG_ERROR("[WS] Failed to send handshake request");
        return false;
    }
    
    // 接收响应
    char response[2048];
    int n = recv(m_sockfd, response, sizeof(response) - 1, 0);
    if (n <= 0) {
        LOG_ERROR("[WS] Failed to receive handshake response");
        return false;
    }
    response[n] = '\0';
    
    // 简单验证响应（检查 "101 Switching Protocols"）
    if (strstr(response, "101") == NULL || 
        strstr(response, "Switching Protocols") == NULL) {
        LOG_ERROR("[WS] Handshake failed: %s", response);
        return false;
    }
    
    LOG_DEBUG("[WS] Handshake successful");
    return true;
}

void WSClient::recv_loop() {
    LOG_INFO("[WS] Receive loop started");
    
    while (m_running && m_connected) {
        std::vector<uint8_t> payload;
        uint8_t opcode;
        
        if (recv_frame(payload, opcode)) {
            if (opcode == OP_TEXT) {
                // 文本消息
                std::string msg((char*)payload.data(), payload.size());
                std::lock_guard<std::mutex> lock(m_recv_mutex);
                m_recv_queue.push(msg);
                LOG_INFO("[WS] 📩 Received text: %s", msg.c_str());
            } else if (opcode == OP_BINARY) {
                // 二进制消息 - 保存到队列
                std::lock_guard<std::mutex> lock(m_binary_mutex);
                m_binary_queue.push(payload);
                // Day 12: 提升日志级别，确认服务器是否发送 TTS 音频
                LOG_INFO("[WS] 🔔 Received binary: %zu bytes (queue: %zu)", payload.size(), m_binary_queue.size());
            } else if (opcode == OP_PING) {
                // 回复 PONG
                send_frame(payload.data(), payload.size(), OP_PONG);
            } else if (opcode == OP_CLOSE) {
                LOG_WARN("[WS] Server closed connection");
                m_connected = false;
                break;
            }
        } else {
            if (m_connected) {
                LOG_ERROR("[WS] Failed to receive frame, disconnecting");
                m_connected = false;
            }
            break;
        }
    }
    
    LOG_INFO("[WS] Receive loop stopped");
}

bool WSClient::send_frame(const uint8_t* data, size_t len, uint8_t opcode) {
    if (m_sockfd < 0 || !m_connected) return false;
    
    std::vector<uint8_t> frame;
    
    // 第1字节: FIN=1, opcode
    frame.push_back(0x80 | opcode);
    
    // 第2字节及后续: MASK=1, payload长度
    uint8_t mask_key[4];
    generate_mask_key(mask_key);
    
    if (len < 126) {
        frame.push_back(0x80 | (uint8_t)len);
    } else if (len < 65536) {
        frame.push_back(0x80 | 126);
        frame.push_back((len >> 8) & 0xFF);
        frame.push_back(len & 0xFF);
    } else {
        frame.push_back(0x80 | 127);
        for (int i = 7; i >= 0; --i) {
            frame.push_back((len >> (i * 8)) & 0xFF);
        }
    }
    
    // Mask key
    for (int i = 0; i < 4; ++i) {
        frame.push_back(mask_key[i]);
    }
    
    // Masked payload
    for (size_t i = 0; i < len; ++i) {
        frame.push_back(data[i] ^ mask_key[i % 4]);
    }
    
    // 发送帧
    ssize_t sent = send(m_sockfd, frame.data(), frame.size(), 0);
    if (sent < 0 || (size_t)sent != frame.size()) {
        LOG_ERROR("[WS] Failed to send frame");
        return false;
    }
    
    // Day 12: 更新最后发送时间
    m_last_send_time = time(nullptr);
    
    return true;
}

bool WSClient::recv_frame(std::vector<uint8_t>& payload, uint8_t& opcode) {
    // 读取前2字节
    uint8_t header[2];
    if (recv(m_sockfd, header, 2, MSG_WAITALL) != 2) {
        return false;
    }
    
    // bool fin = (header[0] & 0x80) != 0;  // 当前未使用
    opcode = header[0] & 0x0F;
    bool masked = (header[1] & 0x80) != 0;
    uint64_t payload_len = header[1] & 0x7F;
    
    // 读取扩展长度
    if (payload_len == 126) {
        uint8_t len_bytes[2];
        if (recv(m_sockfd, len_bytes, 2, MSG_WAITALL) != 2) {
            return false;
        }
        payload_len = ((uint64_t)len_bytes[0] << 8) | len_bytes[1];
    } else if (payload_len == 127) {
        uint8_t len_bytes[8];
        if (recv(m_sockfd, len_bytes, 8, MSG_WAITALL) != 8) {
            return false;
        }
        payload_len = 0;
        for (int i = 0; i < 8; ++i) {
            payload_len = (payload_len << 8) | len_bytes[i];
        }
    }
    
    // 读取 mask key（服务器不应该发送masked数据）
    uint8_t mask_key[4] = {0};
    if (masked) {
        if (recv(m_sockfd, mask_key, 4, MSG_WAITALL) != 4) {
            return false;
        }
    }
    
    // 读取 payload
    payload.resize(payload_len);
    if (payload_len > 0) {
        ssize_t received = recv(m_sockfd, payload.data(), payload_len, MSG_WAITALL);
        if (received != (ssize_t)payload_len) {
            return false;
        }
        
        // 解码（如果有mask）
        if (masked) {
            for (size_t i = 0; i < payload_len; ++i) {
                payload[i] ^= mask_key[i % 4];
            }
        }
    }
    
    return true;
}

void WSClient::generate_mask_key(uint8_t* mask) {
    for (int i = 0; i < 4; ++i) {
        mask[i] = rand() % 256;
    }
}

// Day 12: 心跳 PING 帧发送
void WSClient::send_ping() {
    if (m_sockfd >= 0 && m_connected) {
        send_frame(nullptr, 0, OP_PING);
        LOG_INFO("[WS] 💓 Sent PING heartbeat");
    }
}
