llmoxide 0.1.0

use crate::error::{Error, Result};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SseFrame {
    pub event: Option<String>,
    pub data: String,
}

pub fn drain_sse_frames(buf: &mut String, out: &mut Vec<SseFrame>) -> Result<()> {
    loop {
        let lf_idx = buf.find("\n\n");
        let crlf_idx = buf.find("\r\n\r\n");

        let (idx, drain_len) = match (lf_idx, crlf_idx) {
            (None, None) => break,
            (Some(i), None) => (i, 2),
            (None, Some(i)) => (i, 4),
            (Some(a), Some(b)) => {
                if a <= b {
                    (a, 2)
                } else {
                    (b, 4)
                }
            }
        };

        let frame = buf[..idx].to_string();
        buf.drain(..idx + drain_len);

        let mut event: Option<String> = None;
        let mut data = String::new();

        for raw_line in frame.lines() {
            let line = raw_line.trim_end_matches('\r');
            if let Some(v) = line.strip_prefix("event:") {
                event = Some(v.trim().to_string());
            } else if let Some(v) = line.strip_prefix("data:") {
                if !data.is_empty() {
                    data.push('\n');
                }
                data.push_str(v.trim());
            }
        }

        if data.is_empty() {
            continue;
        }

        if data.len() > 8 * 1024 * 1024 {
            return Err(Error::InvalidInput("sse frame too large".into()));
        }

        out.push(SseFrame { event, data });
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn drains_single_frame_with_event_and_data() {
        let mut buf = "event: foo\ndata: {\"a\":1}\n\n".to_string();
        let mut frames = Vec::new();
        drain_sse_frames(&mut buf, &mut frames).unwrap();
        assert_eq!(buf, "");
        assert_eq!(
            frames,
            vec![SseFrame {
                event: Some("foo".to_string()),
                data: "{\"a\":1}".to_string()
            }]
        );
    }

    #[test]
    fn drains_multiline_data_and_crlf() {
        let mut buf = "event: e\r\ndata: one\r\ndata: two\r\n\r\n".to_string();
        let mut frames = Vec::new();
        drain_sse_frames(&mut buf, &mut frames).unwrap();
        assert_eq!(frames.len(), 1);
        assert_eq!(frames[0].event.as_deref(), Some("e"));
        assert_eq!(frames[0].data, "one\ntwo");
    }
}