module websocket

import encoding.utf8

const header_len_offset = 2 // offset for lengthpart of websocket header

const buffer_size = 256 // default buffer size

const extended_payload16_end_byte = 4 // header length with 16-bit extended payload

const extended_payload64_end_byte = 10

// Fragment represents a websocket data fragment
struct Fragment {
    data   []u8   // included data payload data in a fragment
    opcode OPCode // interpretation of the payload data
}

// Frame represents a data frame header
struct Frame {
mut:
    // length of the websocket header part
    header_len int = 2
    // size of total frame
    frame_size  int = 2
    fin         bool   // true if final fragment of message
    rsv1        bool   // reserved for future use in websocket RFC
    rsv2        bool   // reserved for future use in websocket RFC
    rsv3        bool   // reserved for future use in websocket RFC
    opcode      OPCode // interpretation of the payload data
    has_mask    bool   // true if the payload data is masked
    payload_len int    // payload length
    masking_key [4]u8  // all frames from client to server is masked with this key
}

const invalid_close_codes = [999, 1004, 1005, 1006, 1014, 1015, 1016, 1100, 2000, 2999, 5000, 65536]

// validate_client validates client frame rules from RFC6455
pub fn (mut ws Client) validate_frame(frame &Frame) ! {
    if frame.rsv1 || frame.rsv2 || frame.rsv3 {
        ws.close(1002, 'rsv cannot be other than 0, not negotiated')!
        return error('rsv cannot be other than 0, not negotiated')
    }
    if (int(frame.opcode) >= 3 && int(frame.opcode) <= 7)
        || (int(frame.opcode) >= 11 && int(frame.opcode) <= 15) {
        ws.close(1002, 'use of reserved opcode')!
        return error('use of reserved opcode')
    }
    if frame.has_mask && !ws.is_server {
        // server should never send masked frames
        // to client, close connection
        ws.close(1002, 'client got masked frame')!
        return error('client sent masked frame')
    }
    if is_control_frame(frame.opcode) {
        if !frame.fin {
            ws.close(1002, 'control message must not be fragmented')!
            return error('unexpected control frame with no fin')
        }
        if frame.payload_len > 125 {
            ws.close(1002, 'control frames must not exceed 125 bytes')!
            return error('unexpected control frame payload length')
        }
    }
    if frame.fin == false && ws.fragments.len == 0 && frame.opcode == .continuation {
        err_msg := 'unexecpected continuation, there are no frames to continue, ${frame}'
        ws.close(1002, err_msg)!
        return error(err_msg)
    }
}

// is_control_frame returns true if the frame is a control frame
fn is_control_frame(opcode OPCode) bool {
    return opcode !in [.text_frame, .binary_frame, .continuation]
}

// is_data_frame returns true if the frame is a control frame
fn is_data_frame(opcode OPCode) bool {
    return opcode in [.text_frame, .binary_frame]
}

// read_payload reads the message payload from the socket
fn (mut ws Client) read_payload(frame &Frame) ![]u8 {
    if frame.payload_len == 0 {
        return []u8{}
    }
    mut buffer := []u8{cap: frame.payload_len}
    mut read_buf := [1]u8{}
    mut bytes_read := 0
    for bytes_read < frame.payload_len {
        len := ws.socket_read_ptr(&read_buf[0], 1)!
        if len != 1 {
            return error('expected read all message, got zero')
        }
        bytes_read += len
        buffer << read_buf[0]
    }
    if bytes_read != frame.payload_len {
        return error('failed to read payload')
    }
    if frame.has_mask {
        for i in 0 .. frame.payload_len {
            buffer[i] ^= frame.masking_key[i % 4] & 0xff
        }
    }
    return buffer
}

// validate_utf_8 validates payload for valid utf8 encoding
// - Future implementation needs to support fail fast utf errors for strict autobahn conformance
fn (mut ws Client) validate_utf_8(opcode OPCode, payload []u8) ! {
    if opcode in [.text_frame, .close] && !utf8.validate(payload.data, payload.len) {
        ws.logger.error('malformed utf8 payload, payload len: (${payload.len})')
        ws.send_error_event('Received malformed utf8.')
        ws.close(1007, 'malformed utf8 payload')!
        return error('malformed utf8 payload')
    }
}

// read_next_message reads 1 to n frames to compose a message
pub fn (mut ws Client) read_next_message() !Message {
    for {
        frame := ws.parse_frame_header()!
        ws.validate_frame(&frame)!
        frame_payload := ws.read_payload(&frame)!
        if is_control_frame(frame.opcode) {
            // Control frames can interject other frames
            // and need to be returned immediately
            msg := Message{
                opcode:  OPCode(frame.opcode)
                payload: frame_payload.clone()
            }
            unsafe { frame_payload.free() }
            return msg
        }
        // if the message is fragmented we just put it on fragments
        // a fragment is allowed to have zero size payload
        if !frame.fin {
            ws.fragments << &Fragment{
                data:   frame_payload.clone()
                opcode: frame.opcode
            }
            unsafe { frame_payload.free() }
            continue
        }
        if ws.fragments.len == 0 {
            ws.validate_utf_8(frame.opcode, frame_payload) or {
                ws.logger.error('UTF8 validation error: ${err}, len of payload(${frame_payload.len})')
                ws.send_error_event('UTF8 validation error: ${err}, len of payload(${frame_payload.len})')
                return err
            }
            msg := Message{
                opcode:  OPCode(frame.opcode)
                payload: frame_payload.clone()
            }
            unsafe { frame_payload.free() }
            return msg
        }
        defer {
            ws.fragments = []
        }
        if is_data_frame(frame.opcode) {
            ws.close(0, '')!
            return error('Unexpected frame opcode')
        }
        payload := ws.payload_from_fragments(frame_payload)!
        opcode := ws.opcode_from_fragments()
        ws.validate_utf_8(opcode, payload)!
        msg := Message{
            opcode:  opcode
            payload: payload.clone()
        }
        unsafe {
            frame_payload.free()
            payload.free()
        }
        return msg
    }
    return error('none')
}

// payload_from_fragments returns the whole payload from fragmented message
fn (ws &Client) payload_from_fragments(fin_payload []u8) ![]u8 {
    mut total_size := 0
    for f in ws.fragments {
        if f.data.len > 0 {
            total_size += f.data.len
        }
    }
    total_size += fin_payload.len
    if total_size == 0 {
        return []u8{}
    }
    mut total_buffer := []u8{cap: total_size}
    for f in ws.fragments {
        if f.data.len > 0 {
            total_buffer << f.data
        }
    }
    total_buffer << fin_payload
    return total_buffer
}

// opcode_from_fragments returns the opcode for message from the first fragment sent
fn (ws &Client) opcode_from_fragments() OPCode {
    return OPCode(ws.fragments[0].opcode)
}

// parse_frame_header parses next message by decoding the incoming frames
pub fn (mut ws Client) parse_frame_header() !Frame {
    mut buffer := [256]u8{}
    mut bytes_read := 0
    mut frame := Frame{}
    mut rbuff := [1]u8{}
    mut mask_end_byte := 0
    for ws.get_state() == .open {
        read_bytes := ws.socket_read_ptr(&rbuff[0], 1)!
        if read_bytes == 0 {
            return error('websocket peer closed connection')
        }
        buffer[bytes_read] = rbuff[0]
        bytes_read++
        // parses the first two header bytes to get basic frame information
        if bytes_read == header_len_offset {
            frame.fin = (buffer[0] & 0x80) == 0x80
            frame.rsv1 = (buffer[0] & 0x40) == 0x40
            frame.rsv2 = (buffer[0] & 0x20) == 0x20
            frame.rsv3 = (buffer[0] & 0x10) == 0x10
            frame.opcode = unsafe { OPCode(int(buffer[0] & 0x7F)) }
            frame.has_mask = (buffer[1] & 0x80) == 0x80
            frame.payload_len = buffer[1] & 0x7F
            // if the frame has a mask, set the byte position where the mask ends
            if frame.has_mask {
                mask_end_byte = if frame.payload_len < 126 {
                    header_len_offset + 4
                } else if frame.payload_len == 126 {
                    header_len_offset + 6
                } else if frame.payload_len == 127 {
                    header_len_offset + 12
                } else {
                    0
                } // impossible
            }
            frame.payload_len = frame.payload_len
            frame.frame_size = frame.header_len + frame.payload_len
            if !frame.has_mask && frame.payload_len < 126 {
                break
            }
        }
        if frame.payload_len == 126 && bytes_read == extended_payload16_end_byte {
            frame.header_len += 2
            frame.payload_len = 0
            frame.payload_len |= int(u32(buffer[2]) << 8)
            frame.payload_len |= int(buffer[3])
            frame.frame_size = frame.header_len + frame.payload_len
            if !frame.has_mask {
                break
            }
        }
        if frame.payload_len == 127 && bytes_read == extended_payload64_end_byte {
            frame.header_len += 8
            // these shift operators needs 64 bit on clang with -prod flag
            mut payload_len := u64(0)
            payload_len |= u64(buffer[2]) << 56
            payload_len |= u64(buffer[3]) << 48
            payload_len |= u64(buffer[4]) << 40
            payload_len |= u64(buffer[5]) << 32
            payload_len |= u64(buffer[6]) << 24
            payload_len |= u64(buffer[7]) << 16
            payload_len |= u64(buffer[8]) << 8
            payload_len |= u64(buffer[9])
            frame.payload_len = int(payload_len)
            if !frame.has_mask {
                break
            }
        }
        if frame.has_mask && bytes_read == mask_end_byte {
            frame.masking_key[0] = buffer[mask_end_byte - 4]
            frame.masking_key[1] = buffer[mask_end_byte - 3]
            frame.masking_key[2] = buffer[mask_end_byte - 2]
            frame.masking_key[3] = buffer[mask_end_byte - 1]
            break
        }
    }
    return frame
}

// unmask_sequence unmask any given sequence
fn (f &Frame) unmask_sequence(mut buffer []u8) {
    for i in 0 .. buffer.len {
        buffer[i] ^= f.masking_key[i % 4] & 0xff
    }
}