#include "websocket.h"

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

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

static int recv_all(int fd, void* buf, size_t len) {
    size_t total = 0;
    char* p = (char*)buf;
    while (total < len) {
        int r = lwip_recv(fd, p + total, len - total, 0);
        if (r <= 0) {
            return -1;
        }
        total += r;
    }
    return 0;
}

static uint16_t my_htons(uint16_t val) {
    return (uint16_t)(((val & 0xff) << 8) | ((val & 0xff00) >> 8));
}

int ws_connect(const char* host, int port, const char* path, const char* device_id, const char* auth_key) {
    int fd = lwip_socket(AF_INET, SOCK_STREAM, 0);
    if (fd < 0) return -1;
    
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = my_htons(port);
    addr.sin_addr.s_addr = ipaddr_addr(host);
    
    if (lwip_connect(fd, (struct sockaddr*)&addr, sizeof(addr)) < 0) {
        close(fd);
        return -1;
    }
    
    // Set socket receive timeout (e.g. 90 seconds) to prevent blocking indefinitely
    struct timeval tv;
    tv.tv_sec = 90;
    tv.tv_usec = 0;
    lwip_setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
    
    // Send HTTP upgrade handshake request
    char req[1024];
    snprintf(req, sizeof(req),
             "GET %s HTTP/1.1\r\n"
             "Host: %s:%d\r\n"
             "Upgrade: websocket\r\n"
             "Connection: Upgrade\r\n"
             "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n"
             "Sec-WebSocket-Version: 13\r\n"
             "Authorization: Bearer %s\r\n"
             "X-Device-ID: %s\r\n"
             "\r\n",
             path, host, port, auth_key, device_id);
             
    if (lwip_send(fd, req, strlen(req), 0) < 0) {
        close(fd);
        return -1;
    }
    
    // Read HTTP response headers until we hit "\r\n\r\n"
    char header_buf[1024];
    size_t header_len = 0;
    while (header_len < sizeof(header_buf) - 1) {
        char c;
        int r = lwip_recv(fd, &c, 1, 0);
        if (r <= 0) {
            close(fd);
            return -1;
        }
        header_buf[header_len++] = c;
        header_buf[header_len] = '\0';
        
        if (header_len >= 4 && strcmp(header_buf + header_len - 4, "\r\n\r\n") == 0) {
            break;
        }
    }
    
    // Verify HTTP 101 Switching Protocols response status
    if (strstr(header_buf, "HTTP/1.1 101") == NULL && strstr(header_buf, "HTTP/1.0 101") == NULL) {
        close(fd);
        return -1;
    }
    
    return fd;
}

int ws_send(int fd, const uint8_t* data, size_t len, bool binary) {
    uint8_t header[10];
    size_t header_len = 0;
    
    header[0] = binary ? 0x82 : 0x81;
    
    if (len < 126) {
        header[1] = 0x80 | (uint8_t)len;
        header_len = 2;
    } else {
        header[1] = 0x80 | 126;
        header[2] = (uint8_t)((len >> 8) & 0xFF);
        header[3] = (uint8_t)(len & 0xFF);
        header_len = 4;
    }
    
    // Use fixed client mask for performance: 0x12, 0x34, 0x56, 0x78
    uint8_t mask[4] = { 0x12, 0x34, 0x56, 0x78 };
    memcpy(header + header_len, mask, 4);
    header_len += 4;
    
    // Send WebSocket frame header
    int sent = lwip_send(fd, header, header_len, 0);
    if (sent < 0) return -1;
    
    // Mask the payload
    uint8_t* masked = malloc(len);
    if (masked == NULL) return -1;
    for (size_t i = 0; i < len; ++i) {
        masked[i] = data[i] ^ mask[i % 4];
    }
    
    // Send masked payload
    sent = lwip_send(fd, masked, len, 0);
    free(masked);
    
    return sent >= 0 ? 0 : -1;
}

int ws_recv(int fd, int* out_opcode, uint8_t* payload, size_t max_len) {
    uint8_t header[2];
    if (recv_all(fd, header, 2) < 0) {
        return -1;
    }
    
    int opcode = header[0] & 0x0F;
    if (out_opcode != NULL) {
        *out_opcode = opcode;
    }
    
    int masked = (header[1] & 0x80) != 0;
    size_t len = header[1] & 0x7F;
    
    if (len == 126) {
        uint8_t ext_len[2];
        if (recv_all(fd, ext_len, 2) < 0) return -1;
        len = ((size_t)ext_len[0] << 8) | ext_len[1];
    } else if (len == 127) {
        uint8_t ext_len[8];
        if (recv_all(fd, ext_len, 8) < 0) return -1;
        // Parse 64-bit length into size_t
        len = ((size_t)ext_len[4] << 24) | ((size_t)ext_len[5] << 16) | ((size_t)ext_len[6] << 8) | ext_len[7];
    }
    
    if (masked) {
        uint8_t mask[4];
        if (recv_all(fd, mask, 4) < 0) return -1;
        
        if (len > max_len) {
            ESP_LOGE("websocket", "ws_recv overflow (masked): len=%u, max_len=%u", (unsigned)len, (unsigned)max_len);
            // Buffer overflow, skip payload to align stream
            size_t to_discard = len;
            uint8_t discard_buf[256];
            while (to_discard > 0) {
                size_t chunk = to_discard < sizeof(discard_buf) ? to_discard : sizeof(discard_buf);
                if (recv_all(fd, discard_buf, chunk) < 0) return -1;
                to_discard -= chunk;
            }
            return -2;
        }
        
        if (recv_all(fd, payload, len) < 0) return -1;
        for (size_t i = 0; i < len; ++i) {
            payload[i] ^= mask[i % 4];
        }
    } else {
        if (len > max_len) {
            ESP_LOGE("websocket", "ws_recv overflow (unmasked): len=%u, max_len=%u", (unsigned)len, (unsigned)max_len);
            size_t to_discard = len;
            uint8_t discard_buf[256];
            while (to_discard > 0) {
                size_t chunk = to_discard < sizeof(discard_buf) ? to_discard : sizeof(discard_buf);
                if (recv_all(fd, discard_buf, chunk) < 0) return -1;
                to_discard -= chunk;
            }
            return -2;
        }
        
        if (recv_all(fd, payload, len) < 0) return -1;
    }
    
    return (int)len;
}

void ws_close(int fd) {
    if (fd >= 0) {
        close(fd);
    }
}

int ws_send_pong(int fd, const uint8_t* payload, size_t len) {
    uint8_t header[10];
    size_t header_len = 0;
    
    header[0] = 0x8A; // FIN | PONG (0x0A)
    
    if (len < 126) {
        header[1] = 0x80 | (uint8_t)len;
        header_len = 2;
    } else {
        header[1] = 0x80 | 126;
        header[2] = (uint8_t)((len >> 8) & 0xFF);
        header[3] = (uint8_t)(len & 0xFF);
        header_len = 4;
    }
    
    uint8_t mask[4] = { 0x12, 0x34, 0x56, 0x78 };
    memcpy(header + header_len, mask, 4);
    header_len += 4;
    
    int sent = lwip_send(fd, header, header_len, 0);
    if (sent < 0) return -1;
    
    if (len > 0 && payload != NULL) {
        uint8_t* masked = malloc(len);
        if (masked == NULL) return -1;
        for (size_t i = 0; i < len; ++i) {
            masked[i] = payload[i] ^ mask[i % 4];
        }
        sent = lwip_send(fd, masked, len, 0);
        free(masked);
    }
    
    return sent >= 0 ? 0 : -1;
}