herodb/src/protocol.rs

use core::fmt;

use crate::error::DBError;

#[derive(Debug, Clone)]
pub enum Protocol {
    SimpleString(String),
    BulkString(String),
    Null,
    Array(Vec<Protocol>),
    Error(String), // NEW
}

impl fmt::Display for Protocol {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.decode().as_str())
    }
}

impl Protocol {
    pub fn from(protocol: &str) -> Result<(Self, &str), DBError> {
        if protocol.is_empty() {
            // Incomplete frame; caller should read more bytes
            return Err(DBError("[incomplete] empty".to_string()));
        }
        let ret = match protocol.chars().nth(0) {
            Some('+') => Self::parse_simple_string_sfx(&protocol[1..]),
            Some('$') => Self::parse_bulk_string_sfx(&protocol[1..]),
            Some('*') => Self::parse_array_sfx(&protocol[1..]),
            _ => Err(DBError(format!(
                "[from] unsupported protocol: {:?}",
                protocol
            ))),
        };
        ret
    }

    pub fn from_vec(array: Vec<&str>) -> Self {
        let array = array
            .into_iter()
            .map(|x| Protocol::BulkString(x.to_string()))
            .collect();
        Protocol::Array(array)
    }

    #[inline]
    pub fn ok() -> Self {
        Protocol::SimpleString("ok".to_string())
    }

    #[inline]
    pub fn err(msg: &str) -> Self {
        Protocol::Error(msg.to_string())
    }

    #[inline]
    pub fn write_on_slave_err() -> Self {
        Self::err("DISALLOW WRITE ON SLAVE")
    }

    #[inline]
    pub fn psync_on_slave_err() -> Self {
        Self::err("PSYNC ON SLAVE IS NOT ALLOWED")
    }

    #[inline]
    pub fn none() -> Self {
        Self::SimpleString("none".to_string())
    }

    pub fn decode(&self) -> String {
        match self {
            Protocol::SimpleString(s) => s.to_string(),
            Protocol::BulkString(s) => s.to_string(),
            Protocol::Null => "".to_string(),
            Protocol::Array(s) => s.iter().map(|x| x.decode()).collect::<Vec<_>>().join(" "),
            Protocol::Error(s) => s.to_string(),
        }
    }

    pub fn encode(&self) -> String {
        match self {
            Protocol::SimpleString(s) => format!("+{}\r\n", s),
            Protocol::BulkString(s)   => format!("${}\r\n{}\r\n", s.len(), s),
            Protocol::Array(ss)       => {
                format!("*{}\r\n", ss.len()) + &ss.iter().map(|x| x.encode()).collect::<String>()
            }
            Protocol::Null            => "$-1\r\n".to_string(),
            Protocol::Error(s)        => format!("-{}\r\n", s), // proper RESP error
        }
    }

    fn parse_simple_string_sfx(protocol: &str) -> Result<(Self, &str), DBError> {
        match protocol.find("\r\n") {
            Some(x) => Ok((Self::SimpleString(protocol[..x].to_string()), &protocol[x + 2..])),
            _ => Err(DBError(format!(
                "[new simple string] unsupported protocol: {:?}",
                protocol
            ))),
        }
    }

    fn parse_bulk_string_sfx(protocol: &str) -> Result<(Self, &str), DBError> {
        if let Some(len_end) = protocol.find("\r\n") {
            let size = Self::parse_usize(&protocol[..len_end])?;
            let data_start = len_end + 2;
            let data_end = data_start + size;

            // If we don't yet have the full bulk payload + trailing CRLF, signal INCOMPLETE
            if protocol.len() < data_end + 2 {
                return Err(DBError("[incomplete] bulk body".to_string()));
            }
            if &protocol[data_end..data_end + 2] != "\r\n" {
                return Err(DBError("[incomplete] bulk terminator".to_string()));
            }

            let s = Self::parse_string(&protocol[data_start..data_end])?;
            Ok((Protocol::BulkString(s), &protocol[data_end + 2..]))
        } else {
            // No CRLF after bulk length header yet
            Err(DBError("[incomplete] bulk header".to_string()))
        }
    }

    fn parse_array_sfx(s: &str) -> Result<(Self, &str), DBError> {
        if let Some(len_end) = s.find("\r\n") {
            let array_len = s[..len_end].parse::<usize>()?;
            let mut remaining = &s[len_end + 2..];
            let mut vec = vec![];
            for _ in 0..array_len {
                match Protocol::from(remaining) {
                    Ok((p, rem)) => {
                        vec.push(p);
                        remaining = rem;
                    }
                    Err(e) => {
                        // Propagate incomplete so caller can read more bytes
                        if e.0.starts_with("[incomplete]") {
                            return Err(e);
                        } else {
                            return Err(e);
                        }
                    }
                }
            }
            Ok((Protocol::Array(vec), remaining))
        } else {
            // No CRLF after array header yet
            Err(DBError("[incomplete] array header".to_string()))
        }
    }

    fn parse_usize(protocol: &str) -> Result<usize, DBError> {
        if protocol.is_empty() {
            Err(DBError("Cannot parse usize from empty string".to_string()))
        } else {
            protocol
                .parse::<usize>()
                .map_err(|_| DBError(format!("Failed to parse usize from: {}", protocol)))
        }
    }

    fn parse_string(protocol: &str) -> Result<String, DBError> {
        if protocol.is_empty() {
            // Allow empty strings, but handle appropriately
            Ok("".to_string())
        } else {
            Ok(protocol.to_string())
        }
    }
}