use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use clap::ValueEnum;
use flate2::write::GzEncoder;
use flate2::Compression;
use serde::Serialize;
use walkdir::WalkDir;
use zip::write::SimpleFileOptions;
use crate::archive::ArchiveLayout;
use crate::config::Config;
use crate::errors::{Result, SiteforgeError};
use crate::parse::{estimate_tokens, JsonlChunk};
use crate::storage::Storage;
#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
pub enum ExportFormat {
Markdown,
BasicMarkdown,
Jsonl,
Json,
Html,
Warc,
Zip,
Tar,
AgentTar,
All,
}
impl ExportFormat {
pub fn as_str(self) -> &'static str {
match self {
ExportFormat::Markdown => "markdown",
ExportFormat::BasicMarkdown => "basic-markdown",
ExportFormat::Jsonl => "jsonl",
ExportFormat::Json => "json",
ExportFormat::Html => "html",
ExportFormat::Warc => "warc",
ExportFormat::Zip => "zip",
ExportFormat::Tar => "tar",
ExportFormat::AgentTar => "agent-tar",
ExportFormat::All => "all",
}
}
}
pub fn export_archive(
config: &Config,
archive_id: &str,
format: ExportFormat,
output: Option<&Path>,
) -> Result<PathBuf> {
let layout = layout_for(config, archive_id)?;
ensure_archive_exists(&layout)?;
let output = output
.map(Path::to_path_buf)
.unwrap_or_else(|| layout.exports_dir().join(format.as_str()));
fs::create_dir_all(&output)?;
match format {
ExportFormat::Markdown => copy_dir(&layout.readable_markdown_dir(), &output)?,
ExportFormat::BasicMarkdown => copy_dir(&layout.readable_basic_markdown_dir(), &output)?,
ExportFormat::Json => copy_dir(&layout.readable_json_dir(), &output)?,
ExportFormat::Html => copy_dir(&layout.raw_pages_dir(), &output)?,
ExportFormat::Jsonl => {
fs::copy(layout.chunks_jsonl(), output.join("chunks.jsonl"))?;
}
ExportFormat::Warc => {
write_warc_like(&layout, &output.join(format!("{archive_id}.warc.jsonl")))?;
}
ExportFormat::Zip => {
zip_archive(&layout, &output.join(format!("{archive_id}.zip")))?;
}
ExportFormat::Tar => {
tar_archive(&layout, &output.join(format!("{archive_id}.tar")))?;
}
ExportFormat::AgentTar => {
agent_tar_archive(&layout, &output.join(format!("{archive_id}-agent.tar.gz")))?;
}
ExportFormat::All => {
copy_dir(&layout.readable_markdown_dir(), &output.join("markdown"))?;
copy_dir(
&layout.readable_basic_markdown_dir(),
&output.join("basic_markdown"),
)?;
copy_dir(&layout.readable_json_dir(), &output.join("json"))?;
copy_dir(&layout.raw_pages_dir(), &output.join("raw_pages"))?;
fs::copy(layout.chunks_jsonl(), output.join("chunks.jsonl"))?;
write_warc_like(&layout, &output.join(format!("{archive_id}.warc.jsonl")))?;
zip_archive(&layout, &output.join(format!("{archive_id}.zip")))?;
tar_archive(&layout, &output.join(format!("{archive_id}.tar")))?;
agent_tar_archive(&layout, &output.join(format!("{archive_id}-agent.tar.gz")))?;
}
}
Ok(output)
}
pub fn create_ai_packs(
config: &Config,
archive_id: &str,
max_tokens: Option<usize>,
) -> Result<Vec<PathBuf>> {
let layout = layout_for(config, archive_id)?;
ensure_archive_exists(&layout)?;
fs::create_dir_all(layout.packs_dir())?;
for entry in fs::read_dir(layout.packs_dir())? {
let entry = entry?;
if entry.file_name().to_string_lossy().starts_with("pack_") {
let _ = fs::remove_file(entry.path());
}
}
let storage = Storage::open(&layout.database())?;
let chunks_by_page = load_chunks_by_page(&layout)?;
let pages = storage
.list_pages()?
.into_iter()
.filter(|page| page.markdown_path.is_some())
.collect::<Vec<_>>();
let token_budget = max_tokens.unwrap_or(config.ai_pack_token_budget).max(1_000);
let mut packs = Vec::new();
let mut current = String::new();
let mut current_tokens = 0usize;
let mut current_toc: Vec<(String, String)> = Vec::new();
let mut pack_index = 1usize;
for page in pages {
let Some(markdown_rel) = page.markdown_path.as_deref() else {
continue;
};
let markdown_path = layout.base().join(markdown_rel);
let markdown = fs::read_to_string(&markdown_path)?;
let page_tokens = estimate_tokens(&markdown);
let units = if page_tokens > token_budget {
chunks_by_page
.get(&page.page_id)
.map(|chunks| chunk_pack_units(&page.title, &page.source_url, chunks))
.filter(|units| !units.is_empty())
.unwrap_or_else(|| {
vec![PackUnit::new(
page.title.clone(),
page.source_url.clone(),
format!(
"\n\n# {}\n\nSource: {}\n\n{}",
page.title, page.source_url, markdown
),
)]
})
} else {
vec![PackUnit::new(
page.title.clone(),
page.source_url.clone(),
format!(
"\n\n# {}\n\nSource: {}\n\n{}",
page.title, page.source_url, markdown
),
)]
};
for unit in units {
if current_tokens > 0 && current_tokens + unit.tokens > token_budget {
let path = write_pack(&layout, pack_index, ¤t_toc, ¤t)?;
packs.push(path);
pack_index += 1;
current.clear();
current_tokens = 0;
current_toc.clear();
}
current_toc.push((unit.title, unit.source_url));
current.push_str(&unit.body);
current_tokens += unit.tokens;
}
}
if !current.is_empty() {
packs.push(write_pack(&layout, pack_index, ¤t_toc, ¤t)?);
}
Ok(packs)
}
struct PackUnit {
title: String,
source_url: String,
body: String,
tokens: usize,
}
impl PackUnit {
fn new(title: String, source_url: String, body: String) -> Self {
let tokens = estimate_tokens(&body);
Self {
title,
source_url,
body,
tokens,
}
}
}
fn load_chunks_by_page(layout: &ArchiveLayout) -> Result<HashMap<String, Vec<JsonlChunk>>> {
let mut chunks_by_page: HashMap<String, Vec<JsonlChunk>> = HashMap::new();
if !layout.chunks_jsonl().exists() {
return Ok(chunks_by_page);
}
let raw = fs::read_to_string(layout.chunks_jsonl())?;
for line in raw.lines().filter(|line| !line.trim().is_empty()) {
let chunk: JsonlChunk = serde_json::from_str(line)?;
chunks_by_page
.entry(chunk.page_id.clone())
.or_default()
.push(chunk);
}
for chunks in chunks_by_page.values_mut() {
chunks.sort_by_key(|chunk| chunk.order_index);
}
Ok(chunks_by_page)
}
fn chunk_pack_units(title: &str, source_url: &str, chunks: &[JsonlChunk]) -> Vec<PackUnit> {
chunks
.iter()
.map(|chunk| {
let heading = if chunk.heading_path.is_empty() {
title.to_string()
} else {
chunk.heading_path.join(" > ")
};
let body = if chunk.content_type == "code" {
let language = chunk.code_language.as_deref().unwrap_or("");
format!(
"\n\n# {}\n\nSource: {}\n\nHeading: {}\n\n```{}\n{}\n```\n",
title,
source_url,
heading,
language,
chunk.text.trim_end()
)
} else {
format!(
"\n\n# {}\n\nSource: {}\n\nHeading: {}\n\n{}\n",
title,
source_url,
heading,
chunk.text.trim()
)
};
PackUnit::new(
format!("{} / chunk {}", title, chunk.order_index),
source_url.to_string(),
body,
)
})
.collect()
}
fn write_pack(
layout: &ArchiveLayout,
index: usize,
toc: &[(String, String)],
body: &str,
) -> Result<PathBuf> {
let path = layout.packs_dir().join(format!("pack_{index:04}.md"));
let mut content = String::new();
content.push_str("# Siteforge AI Context Pack\n\n");
content.push_str("This pack preserves source attribution and local asset references. OCR-derived content is labeled in page Markdown.\n\n");
content.push_str("## Table of Contents\n\n");
for (title, url) in toc {
content.push_str(&format!("- {} - {}\n", title, url));
}
content.push_str(body);
fs::write(&path, content)?;
Ok(path)
}
fn write_warc_like(layout: &ArchiveLayout, path: &Path) -> Result<()> {
#[derive(Serialize)]
struct Record {
record_type: &'static str,
source_url: String,
final_url: String,
status: i64,
content_hash: Option<String>,
raw_path: Option<String>,
}
let storage = Storage::open(&layout.database())?;
let mut file = File::create(path)?;
for page in storage.list_pages()? {
let record = Record {
record_type: "response",
source_url: page.source_url,
final_url: page.final_url,
status: page.status,
content_hash: page.content_hash,
raw_path: page.raw_path,
};
writeln!(file, "{}", serde_json::to_string(&record)?)?;
}
Ok(())
}
fn zip_archive(layout: &ArchiveLayout, path: &Path) -> Result<()> {
let file = File::create(path)?;
let mut zip = zip::ZipWriter::new(file);
let options = SimpleFileOptions::default().compression_method(zip::CompressionMethod::Deflated);
for entry in WalkDir::new(layout.base())
.into_iter()
.filter_map(std::result::Result::ok)
{
if !entry.file_type().is_file() || should_skip_bundle_path(layout, entry.path(), path) {
continue;
}
let relative = entry
.path()
.strip_prefix(layout.base())
.unwrap_or(entry.path())
.to_string_lossy()
.replace('\\', "/");
zip.start_file(relative, options)?;
let mut source = File::open(entry.path())?;
let mut bytes = Vec::new();
source.read_to_end(&mut bytes)?;
zip.write_all(&bytes)?;
}
zip.finish()?;
Ok(())
}
fn tar_archive(layout: &ArchiveLayout, path: &Path) -> Result<()> {
let file = File::create(path)?;
let mut builder = tar::Builder::new(file);
for entry in WalkDir::new(layout.base())
.into_iter()
.filter_map(std::result::Result::ok)
{
if !entry.file_type().is_file() || should_skip_bundle_path(layout, entry.path(), path) {
continue;
}
let relative = entry
.path()
.strip_prefix(layout.base())
.unwrap_or(entry.path());
builder.append_path_with_name(entry.path(), relative)?;
}
builder.finish()?;
Ok(())
}
fn agent_tar_archive(layout: &ArchiveLayout, path: &Path) -> Result<()> {
let file = File::create(path)?;
let encoder = GzEncoder::new(file, Compression::default());
let mut builder = tar::Builder::new(encoder);
for entry in WalkDir::new(layout.base())
.into_iter()
.filter_map(std::result::Result::ok)
{
if !entry.file_type().is_file() || should_skip_bundle_path(layout, entry.path(), path) {
continue;
}
let relative = entry
.path()
.strip_prefix(layout.base())
.unwrap_or(entry.path());
let relative_text = relative.to_string_lossy().replace('\\', "/");
if !is_agent_bundle_file(&relative_text) {
continue;
}
builder.append_path_with_name(entry.path(), relative)?;
}
builder.finish()?;
let encoder = builder.into_inner()?;
encoder.finish()?;
Ok(())
}
fn should_skip_bundle_path(layout: &ArchiveLayout, entry_path: &Path, output_path: &Path) -> bool {
if entry_path == output_path {
return true;
}
let relative = entry_path
.strip_prefix(layout.base())
.unwrap_or(entry_path)
.to_string_lossy()
.replace('\\', "/");
relative.starts_with("exports/")
|| relative.ends_with(".db-wal")
|| relative.ends_with(".db-shm")
}
fn is_agent_bundle_file(relative: &str) -> bool {
relative == "AGENTS.md"
|| relative == "agent-index.json"
|| relative == "manifest.json"
|| relative == "checksums.json"
|| relative.starts_with("readable/")
|| relative.starts_with("chunks/")
|| relative.starts_with("packs/")
|| relative.starts_with("raw/pages/")
|| relative.starts_with("raw/rendered/")
|| relative.starts_with("raw/assets/")
}
fn copy_dir(source: &Path, destination: &Path) -> Result<()> {
fs::create_dir_all(destination)?;
if !source.exists() {
return Ok(());
}
for entry in WalkDir::new(source)
.into_iter()
.filter_map(std::result::Result::ok)
{
let relative = entry.path().strip_prefix(source).unwrap_or(entry.path());
let target = destination.join(relative);
if entry.file_type().is_dir() {
fs::create_dir_all(target)?;
} else if entry.file_type().is_file() {
if let Some(parent) = target.parent() {
fs::create_dir_all(parent)?;
}
fs::copy(entry.path(), target)?;
}
}
Ok(())
}
fn layout_for(config: &Config, archive_id: &str) -> Result<ArchiveLayout> {
Ok(ArchiveLayout::new(
config.resolved_archive_root()?,
archive_id.to_string(),
))
}
fn ensure_archive_exists(layout: &ArchiveLayout) -> Result<()> {
if !layout.base().exists() {
return Err(SiteforgeError::ArchiveNotFound(layout.archive_id.clone()));
}
Ok(())
}