armour-rpc 0.1.6

use crate::error::RpcError;
use std::ops::Bound;
use strum::FromRepr;

#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, FromRepr)]
pub enum OpCode {
    Get = 0x01,
    Contains = 0x02,
    First = 0x03,
    Last = 0x04,
    Range = 0x05,
    RangeKeys = 0x06,
    Count = 0x07,
    Upsert = 0x10,
    Remove = 0x11,
    Take = 0x12,
    ApplyBatch = 0x13,
    ListCollections = 0x20,
    GetSchema = 0x21,
}

#[derive(Debug, Clone)]
pub enum UpsertKey {
    Sequence,
    Provided(Vec<u8>),
}

#[derive(Debug)]
pub struct Request {
    pub op: OpCode,
    pub hashname: u64,
    pub payload: RequestPayload,
}

#[derive(Debug)]
pub enum RequestPayload {
    Key(Vec<u8>),
    Empty,
    Range {
        start: Bound<Vec<u8>>,
        end: Bound<Vec<u8>>,
    },
    Upsert {
        key: UpsertKey,
        flag: Option<bool>,
        value: Vec<u8>,
    },
    Remove {
        key: Vec<u8>,
        soft: bool,
    },
    Take {
        key: Vec<u8>,
        soft: bool,
    },
    Count {
        exact: bool,
    },
    Batch(Vec<(Vec<u8>, Option<Vec<u8>>)>),
    ListCollections,
}

/// Response decoded from the wire.
/// `Ok` contains the raw payload bytes; each client method decodes them.
#[derive(Debug)]
pub enum Response {
    Ok(Vec<u8>),
    Err { code: u16, message: String },
}

// --- Encoding helpers ---

pub fn write_bound(buf: &mut Vec<u8>, bound: &Bound<Vec<u8>>) {
    match bound {
        Bound::Unbounded => buf.push(0x00),
        Bound::Included(bytes) => {
            buf.push(0x01);
            buf.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
            buf.extend_from_slice(bytes);
        }
        Bound::Excluded(bytes) => {
            buf.push(0x02);
            buf.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
            buf.extend_from_slice(bytes);
        }
    }
}

pub fn read_bound(buf: &[u8], pos: &mut usize) -> Result<Bound<Vec<u8>>, RpcError> {
    let tag = read_u8(buf, pos)?;
    match tag {
        0x00 => Ok(Bound::Unbounded),
        0x01 => {
            let bytes = read_bytes(buf, pos)?;
            Ok(Bound::Included(bytes))
        }
        0x02 => {
            let bytes = read_bytes(buf, pos)?;
            Ok(Bound::Excluded(bytes))
        }
        _ => Err(RpcError::Protocol("invalid bound tag".to_string())),
    }
}

pub fn read_upsert_key(buf: &[u8], pos: &mut usize) -> Result<UpsertKey, RpcError> {
    let tag = read_u8(buf, pos)?;
    match tag {
        0x00 => Ok(UpsertKey::Sequence),
        0x01 => {
            let key = read_bytes(buf, pos)?;
            Ok(UpsertKey::Provided(key))
        }
        _ => Err(RpcError::Protocol("invalid upsert key tag".to_string())),
    }
}

pub fn write_upsert_key(buf: &mut Vec<u8>, key: &UpsertKey) {
    match key {
        UpsertKey::Sequence => buf.push(0x00),
        UpsertKey::Provided(bytes) => {
            buf.push(0x01);
            buf.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
            buf.extend_from_slice(bytes);
        }
    }
}

// --- Primitive read/write helpers ---

pub fn read_u8(buf: &[u8], pos: &mut usize) -> Result<u8, RpcError> {
    if *pos >= buf.len() {
        return Err(RpcError::UnexpectedEof);
    }
    let v = buf[*pos];
    *pos += 1;
    Ok(v)
}

pub fn read_u16_be(buf: &[u8], pos: &mut usize) -> Result<u16, RpcError> {
    if *pos + 2 > buf.len() {
        return Err(RpcError::UnexpectedEof);
    }
    let v = u16::from_be_bytes([buf[*pos], buf[*pos + 1]]);
    *pos += 2;
    Ok(v)
}

pub fn read_u32_be(buf: &[u8], pos: &mut usize) -> Result<u32, RpcError> {
    if *pos + 4 > buf.len() {
        return Err(RpcError::UnexpectedEof);
    }
    let v = u32::from_be_bytes([buf[*pos], buf[*pos + 1], buf[*pos + 2], buf[*pos + 3]]);
    *pos += 4;
    Ok(v)
}

pub fn read_u64_be(buf: &[u8], pos: &mut usize) -> Result<u64, RpcError> {
    if *pos + 8 > buf.len() {
        return Err(RpcError::UnexpectedEof);
    }
    let v = u64::from_be_bytes([
        buf[*pos],
        buf[*pos + 1],
        buf[*pos + 2],
        buf[*pos + 3],
        buf[*pos + 4],
        buf[*pos + 5],
        buf[*pos + 6],
        buf[*pos + 7],
    ]);
    *pos += 8;
    Ok(v)
}

pub fn read_bytes(buf: &[u8], pos: &mut usize) -> Result<Vec<u8>, RpcError> {
    let len = read_u32_be(buf, pos)? as usize;
    if *pos + len > buf.len() {
        return Err(RpcError::UnexpectedEof);
    }
    let v = buf[*pos..*pos + len].to_vec();
    *pos += len;
    Ok(v)
}

pub fn write_bytes(buf: &mut Vec<u8>, bytes: &[u8]) {
    buf.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
    buf.extend_from_slice(bytes);
}

/// JSON-serialised schema response for [`OpCode::GetSchema`].
///
/// Available under the `schema` feature. Clients typically parse it out of
/// [`Response::Ok`] bytes via `serde_json::from_slice`.
#[cfg(feature = "schema")]
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct SchemaResponse {
    pub name: String,
    pub version: u16,
    pub typ: armour_core::Typ,
    pub key_scheme: armour_core::KeyScheme,
}