use super::icons::get_icon;
use crate::error::FormatError;
use crate::format::Format;
use lex_core::lex::ast::trait_helpers::try_as_container;
use lex_core::lex::ast::traits::{AstNode, Container, VisualStructure};
use lex_core::lex::ast::{ContentItem, Document};
use std::collections::HashMap;
fn format_content_item(
item: &ContentItem,
prefix: &str,
child_index: usize,
child_count: usize,
include_all: bool,
show_linum: bool,
) -> String {
let mut output = String::new();
let is_last = child_index == child_count - 1;
let connector = if is_last { "└─" } else { "├─" };
let icon = get_icon(item.node_type());
let linum_prefix = if show_linum {
format!("{:02} ", item.range().start.line + 1)
} else {
String::new()
};
output.push_str(&format!(
"{}{}{} {} {}\n",
linum_prefix,
prefix,
connector,
icon,
item.display_label()
));
let child_prefix = format!("{}{}", prefix, if is_last { " " } else { "│ " });
if include_all {
if item.has_visual_header() {
if let Some(container) = try_as_container(item) {
let header = container.label();
let header_icon = get_icon(item.node_type());
output.push_str(&format!("{child_prefix}├─ {header_icon} {header}\n"));
}
}
match item {
ContentItem::Session(s) => {
for (i, ann) in s.annotations.iter().enumerate() {
let ann_item = ContentItem::Annotation(ann.clone());
output.push_str(&format_content_item(
&ann_item,
&child_prefix,
i + 1,
s.annotations.len() + s.children().len(),
include_all,
show_linum,
));
}
}
ContentItem::ListItem(li) => {
let marker_icon = get_icon("Marker");
output.push_str(&format!(
"{}├─ {} {}\n",
child_prefix,
marker_icon,
li.marker.as_string()
));
for (i, text_part) in li.text.iter().enumerate() {
let text_icon = get_icon("Text");
let connector = if i == li.text.len() - 1 && li.children().is_empty() {
"└─"
} else {
"├─"
};
output.push_str(&format!(
"{}{} {} {}\n",
child_prefix,
connector,
text_icon,
text_part.as_string()
));
}
for ann in &li.annotations {
let ann_item = ContentItem::Annotation(ann.clone());
output.push_str(&format_content_item(
&ann_item,
&child_prefix,
0,
1,
include_all,
show_linum,
));
}
}
ContentItem::Definition(d) => {
for ann in &d.annotations {
let ann_item = ContentItem::Annotation(ann.clone());
output.push_str(&format_content_item(
&ann_item,
&child_prefix,
0,
1,
include_all,
show_linum,
));
}
}
ContentItem::Annotation(a) => {
for param in &a.data.parameters {
let param_icon = get_icon("Parameter");
output.push_str(&format!(
"{}├─ {} {}={}\n",
child_prefix, param_icon, param.key, param.value
));
}
}
_ => {}
}
}
match item {
ContentItem::VerbatimBlock(v) => {
let mut group_output = String::new();
for (idx, group) in v.group().enumerate() {
let group_label = if v.group_len() == 1 {
group.subject.as_string().to_string()
} else {
format!(
"{} (group {} of {})",
group.subject.as_string(),
idx + 1,
v.group_len()
)
};
let group_icon = get_icon("VerbatimGroup");
let is_last_group = idx == v.group_len() - 1;
let group_connector = if is_last_group { "└─" } else { "├─" };
group_output.push_str(&format!(
"{child_prefix}{group_connector} {group_icon} {group_label}\n"
));
let group_child_prefix = format!(
"{}{}",
child_prefix,
if is_last_group { " " } else { "│ " }
);
for (i, child) in group.children.iter().enumerate() {
group_output.push_str(&format_content_item(
child,
&group_child_prefix,
i,
group.children.len(),
include_all,
show_linum,
));
}
}
output + &group_output
}
_ => {
if let Some(container) = try_as_container(item) {
output
+ &format_children(container.children(), &child_prefix, include_all, show_linum)
} else {
output
}
}
}
}
fn format_children(
children: &[ContentItem],
prefix: &str,
include_all: bool,
show_linum: bool,
) -> String {
let mut output = String::new();
let child_count = children.len();
for (i, child) in children.iter().enumerate() {
output.push_str(&format_content_item(
child,
prefix,
i,
child_count,
include_all,
show_linum,
));
}
output
}
pub fn to_treeviz_str(doc: &Document) -> String {
to_treeviz_str_with_params(doc, &HashMap::new())
}
pub fn to_treeviz_str_with_params(doc: &Document, params: &HashMap<String, String>) -> String {
let include_all = params
.get("ast-full")
.map(|v| v.to_lowercase() == "true")
.unwrap_or(false);
let show_linum = params
.get("show-linum")
.map(|v| v != "false")
.unwrap_or(false);
let icon = get_icon("Document");
let mut output = format!(
"{} Document ({} annotations, {} items)\n",
icon,
doc.annotations.len(),
doc.root.children.len()
);
if include_all {
for annotation in &doc.annotations {
let ann_item = ContentItem::Annotation(annotation.clone());
output.push_str(&format_content_item(
&ann_item,
"",
0,
1,
include_all,
show_linum,
));
}
}
let children = &doc.root.children;
output + &format_children(children, "", include_all, show_linum)
}
pub struct TreevizFormat;
impl Format for TreevizFormat {
fn name(&self) -> &str {
"treeviz"
}
fn description(&self) -> &str {
"Visual tree representation with indentation and Unicode icons"
}
fn file_extensions(&self) -> &[&str] {
&["tree", "treeviz"]
}
fn supports_serialization(&self) -> bool {
true
}
fn serialize(&self, doc: &Document) -> Result<String, FormatError> {
Ok(to_treeviz_str(doc))
}
}