//! Implements writer that sends data to UE2 game server.

use std::cmp::{max, min};
use std::collections::VecDeque;
use std::convert::TryFrom;
use std::iter::Extend;

//  Defines how many bytes is used to encode "LENGTH" field in the chunk sent to ue-server
const CHUNK_LENGTH_FIELD: usize = 2;
//  Maximum amount of bytes ue-server is able to receive at once
const UE_INPUT_BUFFER: usize = 4095;
//  Minimal payload size (in bytes) to send, unless there is not enough data left
const MIN_PAYLOAD_SIZE: usize = 50;

/// For converting byte stream that is expected from the ue-server into actual messages.
/// Conversion process has two steps:
/// 1.  Every string message is converted into it's utf8 representation and is pre-pended with
/// it's own length in format:
///
/// | Data | Length |
/// |---------|---------|
/// | Message `LENGTH` | 4 bytes: u32 BE |
/// | UTF8-encoded string | `LENGTH` bytes|
///
/// Resulting byte sequences from all the messages then concatenated, in order, into
///  a single data stream.
///
/// 2.  Resulting data stream is then separated into "chunks" that can be accepted by
/// the ue-server (each no longer than `UE_INPUT_BUFFER` in total) and are sent in a format:
///
/// | Data | Length |
/// |---------|---------|
/// | Chunk `LENGTH` | 2 bytes: u16 BE |
/// | UTF8-encoded string | `LENGTH` bytes|
///
/// Use `push()` to input string messages and `try_pop()` to retrieve next chunk, if ue-server
/// can accept it.
/// NOTE: `try_pop()` can return `None` even if not all message data has been transferred,
/// in case ue-server's buffer does not have enough space.
///
/// Call `update_ue_received_bytes()` to update `MessageWriter`'s information about
/// how many bytes ue-server has received so far.
///
/// Use `is_empty()` call to check for whether `MessageWriter` has depleted all it's data.
pub struct MessageWriter {
    sent_bytes: u64,
    ue_received_bytes: u64,
    pending_data: VecDeque<u8>,
}

impl MessageWriter {
    pub fn new() -> MessageWriter {
        MessageWriter {
            sent_bytes: 0,
            ue_received_bytes: 0,
            //  This value should be more than enough for typical use:
            //  it will take at least one second to send this much data to a 30 tick rate ue-server.
            pending_data: VecDeque::with_capacity(30 * UE_INPUT_BUFFER),
        }
    }

    pub fn push(&mut self, message: &str) {
        let message_as_utf8 = message.as_bytes();
        let message_as_utf8 = [
            &(message_as_utf8.len() as u32).to_be_bytes(),
            message_as_utf8,
        ]
        .concat();
        self.pending_data.extend(message_as_utf8.into_iter());
    }

    /// This method will always return chunk if all remaining data will fit inside it, otherwise it
    /// will wait until ue-server's buffer has enough space for at least `MIN_PAYLOAD_SIZE` bytes.
    pub fn try_pop(&mut self) -> Option<Vec<u8>> {
        if self.is_empty() {
            return None;
        }
        let required_payload_size = min(self.pending_data.len(), MIN_PAYLOAD_SIZE);
        let available_payload_space = self
            .available_ue_buffer_capacity()
            .checked_sub(CHUNK_LENGTH_FIELD)
            .unwrap_or_default();
        if required_payload_size > available_payload_space {
            return None;
        }
        let payload_size = min(available_payload_space, self.pending_data.len());
        let mut bytes_to_send = Vec::with_capacity(CHUNK_LENGTH_FIELD + payload_size);
        bytes_to_send.extend((payload_size as u16).to_be_bytes().iter());
        for next_byte in self.pending_data.drain(..payload_size) {
            bytes_to_send.push(next_byte);
        }
        self.sent_bytes += bytes_to_send.len() as u64;
        Some(bytes_to_send)
    }

    pub fn is_empty(&self) -> bool {
        self.pending_data.is_empty()
    }

    pub fn update_ue_received_bytes(&mut self, ue_received_bytes: u64) {
        self.ue_received_bytes = max(ue_received_bytes, self.ue_received_bytes);
    }

    fn available_ue_buffer_capacity(&self) -> usize {
        match usize::try_from(self.sent_bytes - self.ue_received_bytes).ok() {
            Some(ue_buffered_bytes) => max(0, UE_INPUT_BUFFER - ue_buffered_bytes),
            _ => 0,
        }
    }
}

#[test]
fn writer_contents_after_creation() {
    let writer = MessageWriter::new();
    assert_eq!(writer.is_empty(), true);
    assert_eq!(writer.available_ue_buffer_capacity(), UE_INPUT_BUFFER);
}

#[test]
fn writer_content_after_push() {
    let mut writer = MessageWriter::new();
    writer.push("Hello, world!");
    assert_eq!(writer.is_empty(), false);
    assert_eq!(writer.available_ue_buffer_capacity(), UE_INPUT_BUFFER);
}

#[test]
fn writing_single_short_message() {
    let mut writer = MessageWriter::new();
    writer.push("Hello, world!");
    let resulting_bytes = writer.try_pop().unwrap();
    let expected_bytes = [
        0, 17, // Bytes in the chunk = message length (4 bytes) + message (13 bytes)
        0, 0, 0, 13, // Bytes in the message
        b'H', b'e', b'l', b'l', b'o', b',', b' ', b'w', b'o', b'r', b'l', b'd', b'!',
    ];
    assert_eq!(writer.is_empty(), true);
    assert_eq!(
        writer.available_ue_buffer_capacity(),
        UE_INPUT_BUFFER - "Hello, world!".len() - 2 - 4
    );
    assert_eq!(resulting_bytes, expected_bytes);
    assert_eq!(writer.sent_bytes, expected_bytes.len() as u64);
}

#[test]
fn writing_first_chunk_of_single_long_message() {
    let mut writer = MessageWriter::new();
    //  Because we also have to pass lengths, this will go over the sending limit
    let long_message = "Q".repeat(UE_INPUT_BUFFER);
    writer.push(&long_message);
    let resulting_bytes = writer.try_pop().unwrap();
    assert_eq!(writer.is_empty(), false);
    assert_eq!(resulting_bytes.len(), UE_INPUT_BUFFER);
    assert_eq!(writer.available_ue_buffer_capacity(), 0);
    //  Bytes in the chunk = 4095 - 2 = 4093 = 0x0ffd
    assert_eq!(resulting_bytes[0], 0x0f);
    assert_eq!(resulting_bytes[1], 0xfd);
    //  Bytes in message = 4095 = 0x0fff
    assert_eq!(resulting_bytes[2], 0);
    assert_eq!(resulting_bytes[3], 0);
    assert_eq!(resulting_bytes[4], 0x0f);
    assert_eq!(resulting_bytes[5], 0xff);
    for &byte in resulting_bytes[6..].iter() {
        assert_eq!(byte, b'Q');
    }
    assert_eq!(writer.try_pop(), None);
    assert_eq!(writer.is_empty(), false);
}

#[test]
fn writing_second_chunk_of_single_long_message() {
    let mut writer = MessageWriter::new();
    //  Because we also have to pass lengths, this will go over the sending limit
    let long_message = "Q".repeat(UE_INPUT_BUFFER);
    writer.push(&long_message);
    //  This pops all but 6 bytes of `long_message`, that were required to encode lengths of
    //  message and first chunk
    let first_bytes = writer.try_pop().unwrap();
    writer.update_ue_received_bytes(first_bytes.len() as u64);
    let resulting_bytes = writer.try_pop().unwrap();
    assert_eq!(
        writer.available_ue_buffer_capacity(),
        UE_INPUT_BUFFER - resulting_bytes.len()
    );
    assert_eq!(writer.is_empty(), true);
    //  Bytes in the chunk = 6
    assert_eq!(resulting_bytes[0], 0);
    assert_eq!(resulting_bytes[1], 6);
    assert_eq!(resulting_bytes[2..], [b'Q', b'Q', b'Q', b'Q', b'Q', b'Q'])
}

#[test]
fn will_write_small_chunks_if_no_more_data() {
    let mut writer = MessageWriter::new();
    //  Because we also have to pass lengths (of chunk `CHUNK_LENGTH_FIELD` amd of message `4`),
    //  sending this will leave us with exactly 10 free bytes in the buffer
    let long_message = "Q".repeat(UE_INPUT_BUFFER / 2);
    writer.push(&long_message);
    writer.try_pop();
    let short_message = "Hello, world!";
    writer.push(&short_message);
    let expected_bytes = [
        0, 17, // Bytes in the chunk = message length (4 bytes) + message (13 bytes)
        0, 0, 0, 13, // Bytes in the message
        b'H', b'e', b'l', b'l', b'o', b',', b' ', b'w', b'o', b'r', b'l', b'd', b'!',
    ];
    //  There should be enough space in the ue-server buffer to send `short_message`
    let resulting_bytes = writer.try_pop().unwrap();
    assert_eq!(resulting_bytes, expected_bytes);
    assert_eq!(writer.try_pop(), None);
    assert_eq!(writer.is_empty(), true);
}

#[test]
fn will_not_write_small_chunks_if_more_data_remains() {
    let mut writer = MessageWriter::new();
    //  Because we also have to pass lengths (of chunk `CHUNK_LENGTH_FIELD` amd of message `4`),
    //  sending this will leave us with exactly 10 free bytes in the buffer
    let long_message = "Q".repeat(UE_INPUT_BUFFER - CHUNK_LENGTH_FIELD - 4 - 10);
    writer.push(&long_message);
    writer.try_pop();
    let short_message = "Hello, world!";
    writer.push(&short_message);
    //  `MessageWriter` can neither send full message, nor a chunk of size 10
    //  (because it is too short)
    assert_eq!(writer.try_pop(), None);
    assert_eq!(writer.is_empty(), false);
}