use std::collections::HashMap;
use thiserror::Error;
use crate::protocol::{Message, extensions::Batch};
#[derive(Error, Debug)]
pub enum MultilineError {
#[error("Maximum bytes exceeded")]
MaxBytesExceeded,
#[error("Maximum lines exceeded")]
MaxLinesExceeded,
#[error("Invalid target")]
InvalidTarget,
#[error("Invalid multiline batch")]
InvalidBatch,
#[error("Multiline not supported")]
NotSupported,
}
#[derive(Debug, Clone)]
pub struct MultilineCapability {
pub max_bytes: usize,
pub max_lines: Option<usize>,
}
impl MultilineCapability {
pub fn new(max_bytes: usize, max_lines: Option<usize>) -> Self {
Self { max_bytes, max_lines }
}
pub fn to_cap_value(&self) -> String {
let mut value = format!("max-bytes={}", self.max_bytes);
if let Some(lines) = self.max_lines {
value.push_str(&format!(",max-lines={}", lines));
}
value
}
pub fn from_cap_value(value: &str) -> Option<Self> {
let mut max_bytes = None;
let mut max_lines = None;
for param in value.split(',') {
if let Some((key, val)) = param.split_once('=') {
match key {
"max-bytes" => {
max_bytes = val.parse().ok();
}
"max-lines" => {
max_lines = val.parse().ok();
}
_ => {}
}
}
}
max_bytes.map(|bytes| Self::new(bytes, max_lines))
}
}
#[derive(Debug, Clone)]
pub struct MultilineBatch {
pub reference: String,
pub target: String,
pub command: String, pub lines: Vec<MultilinePart>,
pub total_bytes: usize,
pub tags: HashMap<String, Option<String>>,
}
#[derive(Debug, Clone)]
pub struct MultilinePart {
pub content: String,
pub concat: bool, }
impl MultilineBatch {
pub fn new(reference: String, target: String, command: String) -> Self {
Self {
reference,
target,
command,
lines: Vec::new(),
total_bytes: 0,
tags: HashMap::new(),
}
}
pub fn add_line(&mut self, content: String, concat: bool) -> Result<(), MultilineError> {
let part = MultilinePart { content, concat };
self.total_bytes += part.content.len();
self.lines.push(part);
Ok(())
}
pub fn validate(&self, capability: &MultilineCapability) -> Result<(), MultilineError> {
if self.total_bytes > capability.max_bytes {
return Err(MultilineError::MaxBytesExceeded);
}
if let Some(max_lines) = capability.max_lines {
if self.lines.len() > max_lines {
return Err(MultilineError::MaxLinesExceeded);
}
}
if !matches!(self.command.as_str(), "PRIVMSG" | "NOTICE") {
return Err(MultilineError::InvalidBatch);
}
if self.target.is_empty() {
return Err(MultilineError::InvalidTarget);
}
Ok(())
}
pub fn compose_message(&self) -> String {
let mut result = String::new();
for (i, part) in self.lines.iter().enumerate() {
if part.concat {
result.push_str(&part.content);
} else {
if i > 0 {
result.push('\n');
}
result.push_str(&part.content);
}
}
result
}
pub fn split_for_fallback(&self, max_line_length: usize) -> Vec<Message> {
let composed = self.compose_message();
let lines: Vec<&str> = composed.split('\n').collect();
let mut messages = Vec::new();
for (i, line) in lines.iter().enumerate() {
let mut msg = Message::new(&self.command)
.with_params(vec![self.target.clone(), line.to_string()]);
if i == 0 {
for (key, value) in &self.tags {
msg = msg.add_tag(key.clone(), value.clone());
}
}
messages.push(msg);
}
messages
}
}
pub struct MultilineProcessor {
capability: MultilineCapability,
active_batches: HashMap<String, MultilineBatch>,
}
impl MultilineProcessor {
pub fn new(capability: MultilineCapability) -> Self {
Self {
capability,
active_batches: HashMap::new(),
}
}
pub fn start_batch(
&mut self,
reference: String,
target: String,
command: String,
tags: HashMap<String, Option<String>>,
) -> Result<(), MultilineError> {
if target.contains(',') {
return Err(MultilineError::InvalidTarget);
}
let mut batch = MultilineBatch::new(reference.clone(), target, command);
batch.tags = tags;
self.active_batches.insert(reference, batch);
Ok(())
}
pub fn add_batch_line(
&mut self,
reference: &str,
message: Message,
) -> Result<(), MultilineError> {
let batch = self.active_batches.get_mut(reference)
.ok_or(MultilineError::InvalidBatch)?;
let content = message.params.last()
.cloned()
.unwrap_or_default();
let concat = message.tags.contains_key("draft/multiline-concat");
batch.add_line(content, concat)?;
batch.validate(&self.capability)?;
Ok(())
}
pub fn end_batch(&mut self, reference: &str) -> Result<MultilineBatch, MultilineError> {
let batch = self.active_batches.remove(reference)
.ok_or(MultilineError::InvalidBatch)?;
batch.validate(&self.capability)?;
Ok(batch)
}
pub fn create_multiline_message(&self, batch: &MultilineBatch) -> Message {
let content = batch.compose_message();
let mut msg = Message::new(&batch.command)
.with_params(vec![batch.target.clone(), content]);
for (key, value) in &batch.tags {
msg = msg.add_tag(key.clone(), value.clone());
}
msg
}
pub fn get_capability(&self) -> &MultilineCapability {
&self.capability
}
}
pub fn create_multiline_fail_message(
error: MultilineError,
context: &str,
) -> Message {
let error_code = match error {
MultilineError::MaxBytesExceeded => "MULTILINE_MAX_BYTES",
MultilineError::MaxLinesExceeded => "MULTILINE_MAX_LINES",
MultilineError::InvalidTarget => "MULTILINE_INVALID_TARGET",
MultilineError::InvalidBatch => "MULTILINE_INVALID",
MultilineError::NotSupported => "MULTILINE_INVALID",
};
Message::new("FAIL")
.with_params(vec!["BATCH".to_string(), error_code.to_string(), context.to_string(), error.to_string()])
}
pub fn is_multiline_batch(batch_type: &str) -> bool {
batch_type == "draft/multiline"
}
pub fn validate_multiline_message(message: &Message) -> Result<(), MultilineError> {
if !matches!(message.command.as_str(), "PRIVMSG" | "NOTICE") {
return Err(MultilineError::InvalidBatch);
}
if message.params.len() < 2 {
return Err(MultilineError::InvalidBatch);
}
for key in message.tags.keys() {
if !matches!(key.as_str(), "draft/multiline-concat" | "batch") {
return Err(MultilineError::InvalidBatch);
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_multiline_capability_parsing() {
let cap = MultilineCapability::from_cap_value("max-bytes=4096,max-lines=100").unwrap();
assert_eq!(cap.max_bytes, 4096);
assert_eq!(cap.max_lines, Some(100));
let cap_value = cap.to_cap_value();
assert!(cap_value.contains("max-bytes=4096"));
assert!(cap_value.contains("max-lines=100"));
}
#[test]
fn test_multiline_batch_composition() {
let mut batch = MultilineBatch::new(
"ref123".to_string(),
"#channel".to_string(),
"PRIVMSG".to_string(),
);
batch.add_line("Line 1".to_string(), false).unwrap();
batch.add_line("Line 2".to_string(), false).unwrap();
batch.add_line(" continued".to_string(), true).unwrap();
let composed = batch.compose_message();
assert_eq!(composed, "Line 1\nLine 2 continued");
}
#[test]
fn test_multiline_validation() {
let capability = MultilineCapability::new(100, Some(3));
let mut batch = MultilineBatch::new(
"ref123".to_string(),
"#channel".to_string(),
"PRIVMSG".to_string(),
);
batch.validate(&capability).unwrap();
batch.add_line("x".repeat(101), false).unwrap();
assert!(matches!(batch.validate(&capability), Err(MultilineError::MaxBytesExceeded)));
}
}