use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{self, BufWriter, Write};
use std::path::{Path, PathBuf};
use undoc::detect::FormatType;
use undoc::model::{Document, Section};
use undoc::render::{
clean_text, refine, render_section_to_string, to_json, to_markdown, to_text, JsonFormat,
RenderOptions,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OutputFormat {
Markdown,
Text,
Json,
}
impl OutputFormat {
pub fn parse(s: &str) -> Option<Self> {
match s.trim().to_ascii_lowercase().as_str() {
"md" | "markdown" => Some(Self::Markdown),
"txt" | "text" => Some(Self::Text),
"json" => Some(Self::Json),
_ => None,
}
}
}
#[derive(Default)]
pub struct WriteSummary {
pub md_path: Option<PathBuf>,
pub txt_path: Option<PathBuf>,
pub json_path: Option<PathBuf>,
pub word_count: usize,
pub section_count: usize,
}
pub struct MultiFormatWriter<'a> {
out_dir: &'a Path,
formats: &'a [OutputFormat],
render_opts: &'a RenderOptions,
}
impl<'a> MultiFormatWriter<'a> {
pub fn new(
out_dir: &'a Path,
formats: &'a [OutputFormat],
render_opts: &'a RenderOptions,
) -> Self {
Self {
out_dir,
formats,
render_opts,
}
}
pub fn write_document(&self, doc: &Document) -> io::Result<WriteSummary> {
let want_md = self.formats.contains(&OutputFormat::Markdown);
let want_txt = self.formats.contains(&OutputFormat::Text);
let want_json = self.formats.contains(&OutputFormat::Json);
let mut summary = WriteSummary::default();
summary.section_count = doc.sections.len();
if want_md {
let markdown = to_markdown(doc, self.render_opts)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?;
let path = self.out_dir.join("extract.md");
fs::write(&path, &markdown)?;
summary.word_count = markdown.split_whitespace().count();
summary.md_path = Some(path);
}
if want_txt {
let text = to_text(doc, self.render_opts)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?;
let path = self.out_dir.join("extract.txt");
fs::write(&path, text)?;
summary.txt_path = Some(path);
}
if want_json {
let json = to_json(doc, JsonFormat::Pretty)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?;
let path = self.out_dir.join("content.json");
fs::write(&path, json)?;
summary.json_path = Some(path);
}
Ok(summary)
}
pub fn open_streaming(
&self,
doc_format: FormatType,
image_map: HashMap<String, String>,
) -> io::Result<StreamingWriter> {
StreamingWriter::new(
self.out_dir,
self.formats,
self.render_opts,
doc_format,
image_map,
)
}
}
pub struct StreamingWriter {
md: Option<BufWriter<File>>,
md_path: Option<PathBuf>,
txt: Option<BufWriter<File>>,
txt_path: Option<PathBuf>,
json: Option<BufWriter<File>>,
json_path: Option<PathBuf>,
json_first_section: bool,
render_opts: RenderOptions,
doc_format: FormatType,
image_map: HashMap<String, String>,
word_count: usize,
section_index: usize,
}
impl StreamingWriter {
fn new(
out_dir: &Path,
formats: &[OutputFormat],
render_opts: &RenderOptions,
doc_format: FormatType,
image_map: HashMap<String, String>,
) -> io::Result<Self> {
let want_md = formats.contains(&OutputFormat::Markdown);
let want_txt = formats.contains(&OutputFormat::Text);
let want_json = formats.contains(&OutputFormat::Json);
let (md, md_path) = if want_md {
let p = out_dir.join("extract.md");
let f = File::create(&p)?;
(Some(BufWriter::new(f)), Some(p))
} else {
(None, None)
};
let (txt, txt_path) = if want_txt {
let p = out_dir.join("extract.txt");
let f = File::create(&p)?;
(Some(BufWriter::new(f)), Some(p))
} else {
(None, None)
};
let (json, json_path) = if want_json {
let p = out_dir.join("content.json");
let f = File::create(&p)?;
(Some(BufWriter::new(f)), Some(p))
} else {
(None, None)
};
Ok(Self {
md,
md_path,
txt,
txt_path,
json,
json_path,
json_first_section: true,
render_opts: render_opts.clone(),
doc_format,
image_map,
word_count: 0,
section_index: 0,
})
}
pub fn write_section(&mut self, section: &Section) -> io::Result<()> {
let idx = self.section_index;
self.section_index += 1;
if let Some(ref mut md) = self.md {
let rendered = render_section_to_string(
section,
idx,
self.doc_format,
&self.render_opts,
&self.image_map,
);
self.word_count += rendered.split_whitespace().count();
md.write_all(rendered.as_bytes())?;
md.write_all(b"\n")?;
}
if let Some(ref mut txt) = self.txt {
for block in §ion.content {
use undoc::model::Block;
match block {
Block::Paragraph(p) => {
let line = p.plain_text();
if !line.is_empty() {
writeln!(txt, "{}", line)?;
}
}
Block::Table(t) => {
for row in &t.rows {
for cell in &row.cells {
let text = cell.plain_text();
if !text.is_empty() {
writeln!(txt, "{}", text)?;
}
}
}
}
_ => {}
}
}
}
if let Some(ref mut json) = self.json {
if !self.json_first_section {
write!(json, ",")?;
}
let section_json = serde_json::to_string(section)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
write!(json, "{}", section_json)?;
self.json_first_section = false;
}
Ok(())
}
pub fn write_json_start(&mut self) -> io::Result<()> {
if let Some(ref mut json) = self.json {
write!(json, "{{\"sections\":[")?;
}
Ok(())
}
pub fn finish(mut self) -> io::Result<WriteSummary> {
let mut summary = WriteSummary::default();
if let (Some(mut md), Some(md_path)) = (self.md.take(), self.md_path.take()) {
md.flush()?;
drop(md);
if self.render_opts.cleanup.is_some() || self.render_opts.refine.is_some() {
let raw = std::fs::read_to_string(&md_path)?;
let cleaned = match self.render_opts.cleanup {
Some(ref cleanup) => clean_text(&raw, cleanup),
None => raw,
};
let refined = match self.render_opts.refine {
Some(ref refine_options) => refine(&cleaned, refine_options),
None => cleaned,
};
std::fs::write(&md_path, refined)?;
}
summary.md_path = Some(md_path);
}
if let (Some(mut txt), Some(txt_path)) = (self.txt.take(), self.txt_path.take()) {
txt.flush()?;
summary.txt_path = Some(txt_path);
}
if let (Some(mut json), Some(json_path)) = (self.json.take(), self.json_path.take()) {
write!(json, "]}}")?;
json.flush()?;
summary.json_path = Some(json_path);
}
summary.word_count = self.word_count;
summary.section_count = self.section_index;
Ok(summary)
}
}