use crate::config::QrConfig;
use crate::error::{MemvidError, Result};
use base64::{Engine as _, engine::general_purpose};
use flate2::Compression;
use flate2::write::GzEncoder;
use image::{DynamicImage, Luma};
use qrcode::{EcLevel, QrCode};
use std::io::Write;
pub struct QrEncoder {
config: QrConfig,
}
#[derive(Debug, Clone)]
pub struct QrFrame {
pub image: DynamicImage,
pub original_text: String,
pub encoded_size: usize,
pub compressed: bool,
}
impl Default for QrEncoder {
fn default() -> Self {
Self::new(QrConfig::default())
}
}
impl QrEncoder {
pub fn new(config: QrConfig) -> Self {
Self { config }
}
pub fn encode_text(&self, text: &str) -> Result<QrFrame> {
let original_text = text.to_string();
let (data_to_encode, compressed) = self.prepare_data(text)?;
let qr_code = self.create_qr_code(&data_to_encode)?;
let image = self.qr_to_image(&qr_code)?;
Ok(QrFrame {
image,
original_text,
encoded_size: data_to_encode.len(),
compressed,
})
}
pub fn encode_chunks(&self, texts: &[String]) -> Result<Vec<QrFrame>> {
texts.iter().map(|text| self.encode_text(text)).collect()
}
fn prepare_data(&self, text: &str) -> Result<(String, bool)> {
if !self.config.enable_compression || text.len() < self.config.compression_threshold {
return Ok((text.to_string(), false));
}
let compressed_data = self.compress_text(text)?;
let compressed_with_prefix = format!("GZ:{}", compressed_data);
if compressed_with_prefix.len() < text.len() {
Ok((compressed_with_prefix, true))
} else {
Ok((text.to_string(), false))
}
}
fn compress_text(&self, text: &str) -> Result<String> {
let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
encoder
.write_all(text.as_bytes())
.map_err(|e| MemvidError::QrCode(format!("Compression failed: {}", e)))?;
let compressed_data = encoder
.finish()
.map_err(|e| MemvidError::QrCode(format!("Compression finalization failed: {}", e)))?;
Ok(general_purpose::STANDARD.encode(&compressed_data))
}
fn create_qr_code(&self, data: &str) -> Result<QrCode> {
let ec_level: EcLevel = self.config.error_correction.clone().into();
let qr_builder = QrCode::with_error_correction_level(data, ec_level);
if let Some(version) = self.config.version {
if !(1..=40).contains(&version) {
return Err(MemvidError::QrCode(format!(
"Invalid QR version: {}",
version
)));
}
}
qr_builder.map_err(|e| MemvidError::QrCode(format!("QR code creation failed: {}", e)))
}
fn qr_to_image(&self, qr_code: &QrCode) -> Result<DynamicImage> {
let qr_image = qr_code
.render::<Luma<u8>>()
.quiet_zone(true) .module_dimensions(self.config.box_size, self.config.box_size) .dark_color(Luma([0u8])) .light_color(Luma([255u8])) .build();
let dynamic_image = DynamicImage::ImageLuma8(qr_image);
Ok(dynamic_image)
}
pub fn estimate_capacity(&self, text: &str) -> Result<QrCapacityInfo> {
let mut capacities = Vec::new();
for &ec_level in &[EcLevel::L, EcLevel::M, EcLevel::Q, EcLevel::H] {
match QrCode::with_error_correction_level(text, ec_level) {
Ok(qr) => {
let version = match qr.version() {
qrcode::Version::Normal(v) => v,
qrcode::Version::Micro(v) => v,
};
capacities.push(QrLevelCapacity {
error_correction: ec_level,
version,
fits: true,
});
}
Err(_) => {
capacities.push(QrLevelCapacity {
error_correction: ec_level,
version: 0,
fits: false,
});
}
}
}
Ok(QrCapacityInfo {
text_length: text.len(),
capacities,
})
}
pub fn get_max_capacity(version: i16, ec_level: EcLevel) -> usize {
match (version, ec_level) {
(1, EcLevel::L) => 25,
(1, EcLevel::M) => 20,
(1, EcLevel::Q) => 16,
(1, EcLevel::H) => 10,
(10, EcLevel::L) => 174,
(10, EcLevel::M) => 136,
(10, EcLevel::Q) => 100,
(10, EcLevel::H) => 74,
(20, EcLevel::L) => 370,
(20, EcLevel::M) => 290,
(20, EcLevel::Q) => 216,
(20, EcLevel::H) => 158,
(40, EcLevel::L) => 852,
(40, EcLevel::M) => 666,
(40, EcLevel::Q) => 496,
(40, EcLevel::H) => 364,
_ => {
let base_capacity = match ec_level {
EcLevel::L => 20,
EcLevel::M => 15,
EcLevel::Q => 12,
EcLevel::H => 8,
};
(base_capacity * version as usize).min(1000)
}
}
}
}
#[derive(Debug, Clone)]
pub struct QrCapacityInfo {
pub text_length: usize,
pub capacities: Vec<QrLevelCapacity>,
}
#[derive(Debug, Clone)]
pub struct QrLevelCapacity {
pub error_correction: EcLevel,
pub version: i16,
pub fits: bool,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_qr_encoding() {
let encoder = QrEncoder::default();
let frame = encoder.encode_text("Hello, World!").unwrap();
assert_eq!(frame.original_text, "Hello, World!");
assert!(!frame.compressed); }
#[test]
fn test_compression() {
let mut config = QrConfig::default();
config.compression_threshold = 10;
let encoder = QrEncoder::new(config);
let long_text = "This is a longer text that should be compressed when encoding into QR codes for efficient storage in video frames.";
let frame = encoder.encode_text(long_text).unwrap();
assert_eq!(frame.original_text, long_text);
}
#[test]
fn test_capacity_estimation() {
let encoder = QrEncoder::default();
let capacity_info = encoder.estimate_capacity("Test").unwrap();
assert_eq!(capacity_info.text_length, 4);
assert_eq!(capacity_info.capacities.len(), 4);
assert!(capacity_info.capacities.iter().any(|c| c.fits));
}
#[test]
fn test_max_capacity() {
assert!(QrEncoder::get_max_capacity(1, EcLevel::L) > 0);
assert!(
QrEncoder::get_max_capacity(40, EcLevel::L)
> QrEncoder::get_max_capacity(1, EcLevel::L)
);
assert!(
QrEncoder::get_max_capacity(20, EcLevel::L)
> QrEncoder::get_max_capacity(20, EcLevel::H)
);
}
#[test]
fn test_empty_text() {
let encoder = QrEncoder::default();
let frame = encoder.encode_text("").unwrap();
assert_eq!(frame.original_text, "");
assert!(!frame.compressed);
}
}