use crate::llm::types::{DxLlmValue, DxSection};
use std::collections::HashMap;
#[derive(Debug, Clone)]
pub struct TableWrapperConfig {
pub max_width: usize,
pub min_col_width: usize,
pub continuation_indicator: String,
}
impl Default for TableWrapperConfig {
fn default() -> Self {
Self {
max_width: 120,
min_col_width: 5,
continuation_indicator: "↓".to_string(),
}
}
}
impl TableWrapperConfig {
pub fn new() -> Self {
Self::default()
}
pub fn with_max_width(mut self, width: usize) -> Self {
self.max_width = width;
self
}
pub fn with_min_col_width(mut self, width: usize) -> Self {
self.min_col_width = width;
self
}
}
pub struct TableWrapper {
config: TableWrapperConfig,
}
impl TableWrapper {
pub fn new() -> Self {
Self {
config: TableWrapperConfig::default(),
}
}
pub fn with_config(config: TableWrapperConfig) -> Self {
Self { config }
}
pub fn needs_wrapping(&self, col_widths: &[usize]) -> bool {
let total_width = self.calculate_table_width(col_widths);
total_width > self.config.max_width
}
fn calculate_table_width(&self, col_widths: &[usize]) -> usize {
if col_widths.is_empty() {
return 0;
}
1 + col_widths.iter().map(|w| w + 3).sum::<usize>()
}
pub fn calculate_widths(
&self,
_section: &DxSection,
header_widths: &[usize],
cell_widths: &[Vec<usize>],
) -> Vec<usize> {
let mut widths: Vec<usize> = header_widths.to_vec();
for row_widths in cell_widths {
for (i, &w) in row_widths.iter().enumerate() {
if i < widths.len() {
widths[i] = widths[i].max(w);
}
}
}
for w in &mut widths {
*w = (*w).max(self.config.min_col_width);
}
if !self.needs_wrapping(&widths) {
return widths;
}
let current_width = self.calculate_table_width(&widths);
let excess = current_width.saturating_sub(self.config.max_width);
if excess == 0 {
return widths;
}
let shrinkable: Vec<(usize, usize)> = widths
.iter()
.enumerate()
.filter(|&(_, &w)| w > self.config.min_col_width)
.map(|(i, &w)| (i, w - self.config.min_col_width))
.collect();
let total_shrinkable: usize = shrinkable.iter().map(|(_, s)| s).sum();
if total_shrinkable == 0 {
return widths; }
let mut remaining_excess = excess;
for (i, shrink_room) in shrinkable {
let shrink_amount = (shrink_room * excess / total_shrinkable).min(remaining_excess);
widths[i] = widths[i].saturating_sub(shrink_amount);
remaining_excess = remaining_excess.saturating_sub(shrink_amount);
}
widths
}
pub fn wrap_cell(&self, content: &str, max_width: usize) -> Vec<String> {
if content.chars().count() <= max_width {
return vec![content.to_string()];
}
let mut lines = Vec::new();
let mut current_line = String::new();
let mut char_count = 0;
for ch in content.chars() {
if char_count >= max_width {
lines.push(current_line);
current_line = String::new();
char_count = 0;
}
current_line.push(ch);
char_count += 1;
}
if !current_line.is_empty() {
lines.push(current_line);
}
lines
}
pub fn wrap_row(
&self,
row: &[DxLlmValue],
col_widths: &[usize],
refs: &HashMap<String, String>,
format_cell: impl Fn(&DxLlmValue, &HashMap<String, String>) -> String,
) -> Vec<Vec<String>> {
let formatted_cells: Vec<String> = row.iter().map(|v| format_cell(v, refs)).collect();
let wrapped_cells: Vec<Vec<String>> = formatted_cells
.iter()
.enumerate()
.map(|(i, content)| {
let max_width = col_widths
.get(i)
.copied()
.unwrap_or(self.config.min_col_width);
self.wrap_cell(content, max_width)
})
.collect();
let max_lines = wrapped_cells.iter().map(|c| c.len()).max().unwrap_or(1);
let mut result = Vec::new();
for line_idx in 0..max_lines {
let mut line_cells = Vec::new();
for (col_idx, wrapped) in wrapped_cells.iter().enumerate() {
let cell_content = wrapped.get(line_idx).cloned().unwrap_or_default();
let width = col_widths
.get(col_idx)
.copied()
.unwrap_or(self.config.min_col_width);
let padding = width.saturating_sub(cell_content.chars().count());
let padded = format!("{}{}", cell_content, " ".repeat(padding));
line_cells.push(padded);
}
result.push(line_cells);
}
result
}
pub fn config(&self) -> &TableWrapperConfig {
&self.config
}
}
impl Default for TableWrapper {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_needs_wrapping_small_table() {
let wrapper = TableWrapper::new();
let col_widths = vec![5, 10, 8];
assert!(!wrapper.needs_wrapping(&col_widths));
}
#[test]
fn test_needs_wrapping_wide_table() {
let config = TableWrapperConfig::new().with_max_width(50);
let wrapper = TableWrapper::with_config(config);
let col_widths = vec![20, 20, 20];
assert!(wrapper.needs_wrapping(&col_widths));
}
#[test]
fn test_wrap_cell_short_content() {
let wrapper = TableWrapper::new();
let result = wrapper.wrap_cell("hello", 10);
assert_eq!(result.len(), 1);
assert_eq!(result[0], "hello");
}
#[test]
fn test_wrap_cell_long_content() {
let wrapper = TableWrapper::new();
let result = wrapper.wrap_cell("hello world", 5);
assert_eq!(result.len(), 3);
assert_eq!(result[0], "hello");
assert_eq!(result[1], " worl");
assert_eq!(result[2], "d");
}
#[test]
fn test_wrap_cell_exact_width() {
let wrapper = TableWrapper::new();
let result = wrapper.wrap_cell("hello", 5);
assert_eq!(result.len(), 1);
assert_eq!(result[0], "hello");
}
#[test]
fn test_calculate_widths_no_shrink_needed() {
let wrapper = TableWrapper::new();
let section = DxSection::new(vec!["id".to_string(), "name".to_string()]);
let header_widths = vec![2, 4];
let cell_widths = vec![vec![1, 5], vec![2, 6]];
let result = wrapper.calculate_widths(§ion, &header_widths, &cell_widths);
assert_eq!(result[0], 5); assert_eq!(result[1], 6); }
#[test]
fn test_calculate_widths_with_shrink() {
let config = TableWrapperConfig::new()
.with_max_width(50)
.with_min_col_width(3);
let wrapper = TableWrapper::with_config(config);
let section = DxSection::new(vec!["col1".to_string(), "col2".to_string()]);
let header_widths = vec![20, 20];
let cell_widths = vec![vec![20, 20]];
let result = wrapper.calculate_widths(§ion, &header_widths, &cell_widths);
assert!(result[0] >= 3); assert!(result[1] >= 3);
assert_eq!(result[0], 20);
assert_eq!(result[1], 20);
}
#[test]
fn test_wrap_row_no_wrap_needed() {
let wrapper = TableWrapper::new();
let row = vec![DxLlmValue::Num(1.0), DxLlmValue::Str("test".to_string())];
let col_widths = vec![5, 10];
let refs = HashMap::new();
let format_cell = |v: &DxLlmValue, _: &HashMap<String, String>| match v {
DxLlmValue::Num(n) => format!("{}", *n as i64),
DxLlmValue::Str(s) => s.clone(),
_ => String::new(),
};
let result = wrapper.wrap_row(&row, &col_widths, &refs, format_cell);
assert_eq!(result.len(), 1);
assert_eq!(result[0].len(), 2);
}
#[test]
fn test_wrap_row_with_wrap() {
let wrapper = TableWrapper::new();
let row = vec![
DxLlmValue::Num(1.0),
DxLlmValue::Str("this is a long string".to_string()),
];
let col_widths = vec![5, 10];
let refs = HashMap::new();
let format_cell = |v: &DxLlmValue, _: &HashMap<String, String>| match v {
DxLlmValue::Num(n) => format!("{}", *n as i64),
DxLlmValue::Str(s) => s.clone(),
_ => String::new(),
};
let result = wrapper.wrap_row(&row, &col_widths, &refs, format_cell);
assert!(result.len() > 1);
for (i, line) in result.iter().enumerate() {
if i > 0 {
assert!(line[0].trim().is_empty() || line[0].chars().count() <= col_widths[0]);
}
}
}
#[test]
fn test_table_width_calculation() {
let wrapper = TableWrapper::new();
assert_eq!(wrapper.calculate_table_width(&[]), 0);
assert_eq!(wrapper.calculate_table_width(&[5]), 9);
assert_eq!(wrapper.calculate_table_width(&[3, 4]), 14);
}
}