taskserver_protocol 0.1.1

use itertools::Itertools;
use std::{borrow::BorrowMut, collections::HashMap};

/// All Response codes as defined in
/// <https://taskwarrior.org/docs/design/protocol.html>
pub enum ResponseCode {
    Success,
    NoChange,
    Redirect,
    MalformedData,
    UnsupportedEncoding,
    TemporarilyUnavailable,
    ShuttingDown,
    AccessDenied,
    AccountSuspended,
    AccountTerminated,
    SyntaxErrorInRequest,
    SyntaxErrorIllegalParameters,
    NotImplemented,
    CommandParameterNotImplemented,
    RequestTooBig,
}

impl ResponseCode {
    fn get_code(&self) -> u16 {
        match self {
            ResponseCode::Success => 200,
            ResponseCode::NoChange => 201,
            ResponseCode::Redirect => 301,
            ResponseCode::MalformedData => 400,
            ResponseCode::UnsupportedEncoding => 401,
            ResponseCode::TemporarilyUnavailable => 420,
            ResponseCode::ShuttingDown => 421,
            ResponseCode::AccessDenied => 430,
            ResponseCode::AccountSuspended => 431,
            ResponseCode::AccountTerminated => 432,
            ResponseCode::SyntaxErrorInRequest => 500,
            ResponseCode::SyntaxErrorIllegalParameters => 501,
            ResponseCode::NotImplemented => 502,
            ResponseCode::CommandParameterNotImplemented => 503,
            ResponseCode::RequestTooBig => 504,
        }
    }
    fn get_status_text(&self) -> &str {
        match self {
            ResponseCode::Success => "ok",
            ResponseCode::NoChange => "no change",
            ResponseCode::Redirect => "redirect to specified port",
            ResponseCode::MalformedData => "malformed data",
            ResponseCode::UnsupportedEncoding => "unsupported encoding",
            ResponseCode::TemporarilyUnavailable => "max number of concurrent connections reached",
            ResponseCode::ShuttingDown => "Server shutting down on operator request",
            ResponseCode::AccessDenied => "access denied",
            ResponseCode::AccountSuspended => "account suspended",
            ResponseCode::AccountTerminated => "account terminated",
            ResponseCode::SyntaxErrorInRequest => "syntax error in request",
            ResponseCode::SyntaxErrorIllegalParameters => "illegal request parameters:",
            ResponseCode::NotImplemented => "not implemented",
            ResponseCode::CommandParameterNotImplemented => "command parameter not implemented",
            ResponseCode::RequestTooBig => "request too big",
        }
    }
}

struct SerializedResponse {
    serialized_header: String,
    payload: Vec<String>,
}

#[derive(Clone)]
enum RespIterValue {
    MessageSizeBuf([u8; 4]),
    BorrowedSlice(&'static [u8]),
    String(String),
}

impl AsRef<[u8]> for RespIterValue {
    fn as_ref(&self) -> &[u8] {
        match self {
            RespIterValue::MessageSizeBuf(buf) => buf,
            RespIterValue::BorrowedSlice(slice) => *slice,
            RespIterValue::String(string) => string.as_bytes(),
        }
    }
}

impl SerializedResponse {
    fn into_buf(self) -> impl bytes::Buf {
        let header_size: usize = self.serialized_header.len();
        let payload_size: usize =
            self.payload.len() + self.payload.iter().map(|s| s.len()).sum::<usize>();
        let total_size: u32 = (header_size + payload_size + 4) as u32;

        let data_iter = std::iter::once(self.serialized_header)
            .chain(self.payload.into_iter())
            .map(RespIterValue::String);
        let newlines_iter = std::iter::repeat(RespIterValue::BorrowedSlice("\n".as_bytes()));
        let data_iter = data_iter.interleave_shortest(newlines_iter);

        ResponseBuf {
            remaining_bytes: total_size as usize,
            current_index: 0,
            current_buf: RespIterValue::MessageSizeBuf(total_size.to_be_bytes()),
            remaining_entries: data_iter,
        }
    }
}

/// Taskserver Response
///
/// see [ResponseBuildExt](crate::response::ResponseBuildExt) for details on modifying it
pub struct Response {
    /// the Response code the server is responding with
    pub code: ResponseCode,
    /// contains all headers. the client, code, status and protocol version headers are added when into_buf is called
    pub extra_headers: HashMap<String, String>,
    /// The payload the server is going to respond with
    pub payload: Vec<String>,
}

struct ResponseBuf<I, T> {
    remaining_bytes: usize,
    current_index: usize,
    current_buf: T,
    remaining_entries: I,
}

impl<T: AsRef<[u8]>, I: Iterator<Item = T>> bytes::Buf for ResponseBuf<I, T> {
    fn remaining(&self) -> usize {
        self.remaining_bytes
    }
    fn bytes(&self) -> &[u8] {
        &self.current_buf.as_ref()
    }
    fn advance(&mut self, mut cnt: usize) {
        if cnt > self.remaining_bytes {
            panic!("trying to advance beyond the buffer size");
        }
        self.remaining_bytes -= cnt;
        while cnt > 0 {
            let current_buf_len_remaining = self.current_buf.as_ref().len() - self.current_index;
            if current_buf_len_remaining > cnt {
                self.current_index += cnt;
                return;
            }
            cnt -= current_buf_len_remaining;
            self.current_index = 0;
            self.current_buf = self
                .remaining_entries
                .next()
                .expect("the underlying iterator did not provide enough bytes to fill the buffer");
        }
    }
}

impl Response {
    /// Creates new Response, with success Response code and empty payload
    pub fn new() -> Response {
        let mut headers = HashMap::new();
        headers.insert(
            "client".into(),
            concat!("taskd_rs ", env!("CARGO_PKG_VERSION")).into(),
        );
        Response {
            code: ResponseCode::Success,
            extra_headers: headers,
            payload: Vec::new(),
        }
    }

    /// creates a buffer adapter to allow for writing out the response in one call
    pub fn into_buf(self) -> impl bytes::Buf {
        self.serialize().into_buf()
    }

    fn serialize(mut self) -> SerializedResponse {
        self.extra_headers
            .insert("code".into(), format!("{}", self.code.get_code()));
        self.extra_headers
            .insert("status".into(), self.code.get_status_text().into());
        self.extra_headers.insert("type".into(), "response".into());
        self.extra_headers.insert("protocol".into(), "v1".into());
        let len = self
            .extra_headers
            .iter()
            .map(|(k, v)| k.len() + v.len() + 2 + 1) // length of key and value including ': ' and \n
            .sum::<usize>();
        let mut serialized_header = String::with_capacity(len);
        for (name, value) in self.extra_headers {
            serialized_header.push_str(&name);
            serialized_header.push_str(": ");
            serialized_header.push_str(&value);
            serialized_header.push('\n');
        }

        SerializedResponse {
            serialized_header,
            payload: self.payload,
        }
    }
}

/// Trait to allow modifying the response by either passing it by value or by mutable reference
///
/// ```
/// use taskserver_protocol::response::*;
/// let mut response = Response::new();
/// let mut borrowed = &mut response;
/// borrowed.code(ResponseCode::NoChange);
///
/// //we can still access the value (it was passed by ref)
/// response.header("test", "foo");
/// ```
///
pub trait ResponseBuildExt {
    /// sets the code
    fn code(self, code: ResponseCode) -> Self;
    /// sets the given header to the provided name
    fn header<I1: Into<String>, I2: Into<String>>(self, name: I1, value: I2) -> Self;
    /// sets the response payload
    fn payload(self, payload: Vec<String>) -> Self;
}

impl<T: BorrowMut<Response>> ResponseBuildExt for T {
    fn code(mut self, code: ResponseCode) -> Self {
        self.borrow_mut().code = code;
        self
    }
    fn header<I1: Into<String>, I2: Into<String>>(mut self, name: I1, value: I2) -> Self {
        self.borrow_mut()
            .extra_headers
            .insert(name.into(), value.into());
        self
    }
    fn payload(mut self, payload: Vec<String>) -> Self {
        self.borrow_mut().payload = payload;
        self
    }
}