use crate::{Package, PackageManifest, Resource, ResourceType};
use oxiarc_archive::zip::{ZipCompressionLevel, ZipWriter};
use std::fs::File;
use std::io::Write;
use std::path::{Path, PathBuf};
use torsh_core::error::{Result, TorshError};
#[derive(Debug, Clone)]
pub struct ExportConfig {
pub compression: ZipCompressionLevel,
pub include_source: bool,
pub include_debug_info: bool,
pub sign_package: bool,
pub verbose: bool,
pub max_size: u64,
}
impl Default for ExportConfig {
fn default() -> Self {
Self {
compression: ZipCompressionLevel::Normal,
include_source: false,
include_debug_info: false,
sign_package: false,
verbose: false,
max_size: 0,
}
}
}
pub struct PackageExporter {
config: ExportConfig,
}
impl PackageExporter {
pub fn new(config: ExportConfig) -> Self {
Self { config }
}
pub fn export_package<P: AsRef<Path>>(&self, package: &Package, path: P) -> Result<()> {
let path = path.as_ref();
package
.manifest
.validate()
.map_err(|e| TorshError::InvalidArgument(format!("Invalid manifest: {}", e)))?;
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let file = File::create(path)?;
let mut zip = ZipWriter::new(file);
zip.set_compression(self.config.compression);
if self.config.verbose {
println!("Writing manifest...");
}
self.write_manifest(&mut zip, &package.manifest)?;
let mut total_size = 0u64;
for (name, resource) in &package.resources {
if self.config.verbose {
println!("Writing resource: {}", name);
}
if self.config.max_size > 0 {
total_size += resource.size() as u64;
if total_size > self.config.max_size {
return Err(TorshError::InvalidArgument(format!(
"Package size ({} bytes) exceeds maximum allowed size ({} bytes)",
total_size, self.config.max_size
)));
}
}
if resource.resource_type == ResourceType::Source && !self.config.include_source {
continue;
}
self.write_resource(&mut zip, name, resource)?;
}
zip.finish()
.map_err(|e| TorshError::IoError(e.to_string()))?;
if self.config.verbose {
println!("Package exported successfully to: {:?}", path);
}
Ok(())
}
fn write_manifest<W: Write>(
&self,
zip: &mut ZipWriter<W>,
manifest: &PackageManifest,
) -> Result<()> {
let json = serde_json::to_string_pretty(manifest)
.map_err(|e| TorshError::SerializationError(e.to_string()))?;
zip.add_file("MANIFEST.json", json.as_bytes())
.map_err(|e| TorshError::IoError(e.to_string()))?;
Ok(())
}
fn write_resource<W: Write>(
&self,
zip: &mut ZipWriter<W>,
name: &str,
resource: &Resource,
) -> Result<()> {
let archive_path = match resource.resource_type {
ResourceType::Model => format!("models/{}", name),
ResourceType::Source => format!("src/{}", name),
ResourceType::Data => format!("data/{}", name),
ResourceType::Config => format!("config/{}", name),
ResourceType::Documentation => format!("docs/{}", name),
_ => format!("resources/{}", name),
};
zip.add_file(&archive_path, &resource.data)
.map_err(|e| TorshError::IoError(e.to_string()))?;
if !resource.metadata.is_empty() {
let metadata_path = format!("{}.metadata", archive_path);
let metadata_json = serde_json::to_string(&resource.metadata)
.map_err(|e| TorshError::SerializationError(e.to_string()))?;
zip.add_file(&metadata_path, metadata_json.as_bytes())
.map_err(|e| TorshError::IoError(e.to_string()))?;
}
Ok(())
}
}
pub struct ExportBuilder {
package: Package,
config: ExportConfig,
output_path: Option<PathBuf>,
}
impl ExportBuilder {
pub fn new(name: String, version: String) -> Self {
Self {
package: Package::new(name, version),
config: ExportConfig::default(),
output_path: None,
}
}
pub fn with_config(mut self, config: ExportConfig) -> Self {
self.config = config;
self
}
pub fn output_path<P: AsRef<Path>>(mut self, path: P) -> Self {
self.output_path = Some(path.as_ref().to_path_buf());
self
}
pub fn author(mut self, author: String) -> Self {
self.package.manifest.author = Some(author);
self
}
pub fn description(mut self, description: String) -> Self {
self.package.manifest.description = Some(description);
self
}
pub fn license(mut self, license: String) -> Self {
self.package.manifest.license = Some(license);
self
}
#[cfg(feature = "with-nn")]
pub fn add_module<M: torsh_nn::Module>(mut self, name: &str, module: &M) -> Result<Self> {
self.package
.add_module(name, module, self.config.include_source)?;
Ok(self)
}
pub fn add_data_file<P: AsRef<Path>>(mut self, name: &str, path: P) -> Result<Self> {
self.package.add_data_file(name, path)?;
Ok(self)
}
pub fn add_resource(mut self, resource: Resource) -> Self {
self.package
.resources
.insert(resource.name.clone(), resource);
self
}
pub fn add_metadata(mut self, key: String, value: String) -> Self {
self.package.manifest.metadata.insert(key, value);
self
}
pub fn export(self) -> Result<PathBuf> {
let output_path = self.output_path.unwrap_or_else(|| {
PathBuf::from(format!(
"{}-{}.torshpkg",
self.package.manifest.name, self.package.manifest.version
))
});
let exporter = PackageExporter::new(self.config);
exporter.export_package(&self.package, &output_path)?;
Ok(output_path)
}
}
#[cfg(feature = "with-nn")]
pub fn export_single_module<M: torsh_nn::Module, P: AsRef<Path>>(
module: &M,
name: &str,
version: &str,
output_path: P,
) -> Result<()> {
ExportBuilder::new(name.to_string(), version.to_string())
.add_module("main", module)?
.output_path(output_path)
.export()?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[test]
fn test_export_config() {
let config = ExportConfig::default();
assert_eq!(config.compression, ZipCompressionLevel::Normal);
assert!(!config.include_source);
}
#[test]
fn test_package_export() {
let temp_dir = TempDir::new().expect("Failed to create temp directory for test");
let output_path = temp_dir.path().join("test.torshpkg");
let mut package = Package::new("test_package".to_string(), "1.0.0".to_string());
let resource = Resource::new(
"test.txt".to_string(),
ResourceType::Text,
b"Hello, World!".to_vec(),
);
package.resources.insert(resource.name.clone(), resource);
let exporter = PackageExporter::new(ExportConfig::default());
exporter
.export_package(&package, &output_path)
.expect("Failed to export package in test");
assert!(output_path.exists());
}
#[test]
fn test_export_builder() {
let temp_dir = TempDir::new().expect("Failed to create temp directory for test");
let output_path = temp_dir.path().join("test.torshpkg");
let path = ExportBuilder::new("test".to_string(), "1.0.0".to_string())
.author("Test Author".to_string())
.description("Test package".to_string())
.license("MIT".to_string())
.output_path(&output_path)
.export()
.expect("Failed to export using builder in test");
assert_eq!(path, output_path);
assert!(path.exists());
}
}