#include "websocket.h"

#include <string.h>
#include <stdio.h>
#include <unistd.h>

#include <esp_log.h>
#include <esp_random.h>
#include <lwip/inet.h>
#include <lwip/sockets.h>

#define TAG "PipecatVoiceWs"
#define WS_HEADER_LIMIT 1024U
#define WS_CONTROL_LIMIT 125U

static int send_all(int fd, const uint8_t* data, size_t length) {
    size_t sent = 0;
    while (sent < length) {
        int result = lwip_send(fd, data + sent, length - sent, 0);
        if (result <= 0) return -1;
        sent += (size_t)result;
    }
    return 0;
}

static int recv_all(int fd, uint8_t* data, size_t length) {
    size_t received = 0;
    while (received < length) {
        int result = lwip_recv(fd, data + received, length - received, 0);
        if (result <= 0) return -1;
        received += (size_t)result;
    }
    return 0;
}

static int discard(int fd, uint64_t length) {
    uint8_t buffer[256];
    while (length > 0) {
        size_t chunk = length > sizeof(buffer) ? sizeof(buffer) : (size_t)length;
        if (recv_all(fd, buffer, chunk) < 0) return -1;
        length -= chunk;
    }
    return 0;
}

static int send_frame(int fd, uint8_t opcode, const uint8_t* payload, size_t length) {
    if (length > 65535U || ((opcode & 0x08U) && length > WS_CONTROL_LIMIT)) return -1;
    uint8_t header[8];
    size_t header_length = 2;
    header[0] = 0x80U | opcode;
    if (length < 126U) {
        header[1] = 0x80U | (uint8_t)length;
    } else {
        header[1] = 0x80U | 126U;
        header[2] = (uint8_t)(length >> 8U);
        header[3] = (uint8_t)length;
        header_length = 4;
    }
    uint8_t mask[4];
    uint32_t random = esp_random();
    memcpy(mask, &random, sizeof(mask));
    memcpy(header + header_length, mask, sizeof(mask));
    header_length += sizeof(mask);
    if (send_all(fd, header, header_length) < 0) return -1;

    uint8_t chunk[512];
    size_t offset = 0;
    while (offset < length) {
        size_t count = length - offset > sizeof(chunk) ? sizeof(chunk) : length - offset;
        for (size_t i = 0; i < count; ++i) chunk[i] = payload[offset + i] ^ mask[(offset + i) % sizeof(mask)];
        if (send_all(fd, chunk, count) < 0) return -1;
        offset += count;
    }
    return 0;
}

int ws_connect(const char* host, int port, const char* path, const char* device_id, const char* api_key) {
    if (host == NULL || path == NULL || device_id == NULL || api_key == NULL || port < 1 || port > 65535) return -1;
    int fd = lwip_socket(AF_INET, SOCK_STREAM, 0);
    if (fd < 0) {
        ESP_LOGW(TAG, "socket create failed");
        return -1;
    }
    struct sockaddr_in address = {0};
    address.sin_family = AF_INET;
    address.sin_port = htons((uint16_t)port);
    address.sin_addr.s_addr = ipaddr_addr(host);
    if (address.sin_addr.s_addr == IPADDR_NONE) {
        ESP_LOGW(TAG, "endpoint address parse failed");
        close(fd);
        return -1;
    }
    if (lwip_connect(fd, (struct sockaddr*)&address, sizeof(address)) < 0) {
        ESP_LOGW(TAG, "TCP connect failed");
        close(fd);
        return -1;
    }
    struct timeval timeout = {.tv_sec = 15, .tv_usec = 0};
    lwip_setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));
    char request[WS_HEADER_LIMIT];
    int request_length = snprintf(request, sizeof(request),
        "GET %s HTTP/1.1\r\nHost: %s:%d\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n"
        "Sec-WebSocket-Key: MDEyMzQ1Njc4OWFiY2RlZg==\r\nSec-WebSocket-Version: 13\r\n"
        "Authorization: Bearer %s\r\nX-Device-ID: %s\r\n\r\n",
        path, host, port, api_key, device_id);
    if (request_length < 0 || (size_t)request_length >= sizeof(request) || send_all(fd, (const uint8_t*)request, (size_t)request_length) < 0) {
        ESP_LOGW(TAG, "WebSocket upgrade request failed");
        close(fd);
        return -1;
    }
    char response[WS_HEADER_LIMIT];
    size_t length = 0;
    while (length + 1 < sizeof(response)) {
        if (recv_all(fd, (uint8_t*)&response[length], 1) < 0) {
            ESP_LOGW(TAG, "WebSocket upgrade response failed");
            close(fd);
            return -1;
        }
        response[++length] = '\0';
        if (length >= 4 && memcmp(response + length - 4, "\r\n\r\n", 4) == 0) break;
    }
    if (length + 1 >= sizeof(response) || strstr(response, " 101 ") == NULL) {
        ESP_LOGW(TAG, "WebSocket upgrade rejected");
        close(fd);
        return -1;
    }
    ESP_LOGI(TAG, "WebSocket upgrade accepted");
    return fd;
}

int ws_send(int fd, const uint8_t* data, size_t length, bool binary) {
    if (fd < 0 || data == NULL || length == 0) return -1;
    return send_frame(fd, binary ? 0x02U : 0x01U, data, length);
}

int ws_recv(int fd, int* opcode, bool* final, uint8_t* payload, size_t maximum) {
    uint8_t header[2];
    if (fd < 0 || recv_all(fd, header, sizeof(header)) < 0) return -1;
    uint64_t length = header[1] & 0x7fU;
    if (length == 126U) {
        uint8_t extended[2];
        if (recv_all(fd, extended, sizeof(extended)) < 0) return -1;
        length = ((uint64_t)extended[0] << 8U) | extended[1];
    } else if (length == 127U) {
        uint8_t extended[8];
        if (recv_all(fd, extended, sizeof(extended)) < 0) return -1;
        length = 0;
        for (size_t i = 0; i < sizeof(extended); ++i) length = (length << 8U) | extended[i];
    }
    bool masked = (header[1] & 0x80U) != 0;
    uint8_t mask[4] = {0};
    if (masked && recv_all(fd, mask, sizeof(mask)) < 0) return -1;
    uint8_t frame_opcode = header[0] & 0x0fU;
    if (((frame_opcode & 0x08U) && (length > WS_CONTROL_LIMIT || !(header[0] & 0x80U))) || length > maximum) {
        if (discard(fd, length) < 0) return -1;
        return -2;
    }
    if (length > 0 && recv_all(fd, payload, (size_t)length) < 0) return -1;
    if (masked) for (size_t i = 0; i < (size_t)length; ++i) payload[i] ^= mask[i % sizeof(mask)];
    if (opcode) *opcode = frame_opcode;
    if (final) *final = (header[0] & 0x80U) != 0;
    return (int)length;
}

int ws_send_pong(int fd, const uint8_t* payload, size_t length) { return send_frame(fd, 0x0aU, payload, length); }
int ws_send_close(int fd) { return send_frame(fd, 0x08U, NULL, 0); }
void ws_close(int fd) { if (fd >= 0) close(fd); }