QRGen/core/src/qr.rs

use crate::ecc::reed_solomon;
use crate::encoder::bitstream::build_codewords;
use crate::encoder::segment::{segment_bit_length, segment_text};
use crate::matrix::grid::Matrix;
use crate::matrix::mask::best_mask;
use crate::matrix::patterns::{
    encode_format_info, encode_version_info, place_alignment_patterns, place_finder_patterns,
    place_format_info, place_timing_patterns, place_version_info, reserve_format_areas,
    reserve_version_areas,
};
use crate::matrix::placement::place_data;
use crate::version::{get_data_capacity, EcLevel, Version};

/// 版本选择模式
#[derive(Debug, Clone)]
pub enum VersionMode {
    Auto,
    Fixed(u8),
}

/// QR 码配置
#[derive(Debug, Clone)]
pub struct QrConfig {
    pub level: EcLevel,
    pub version: VersionMode,
    pub margin: u8,
}

impl Default for QrConfig {
    fn default() -> Self {
        QrConfig {
            level: EcLevel::M,
            version: VersionMode::Auto,
            margin: 4,
        }
    }
}

/// 生成的 QR 码
pub struct QrCode {
    pub version: Version,
    pub level: EcLevel,
    pub mask: u8,
    matrix: Matrix,
    pub margin: u8,
}

impl QrCode {
    /// 编码字符串生成 QR 码
    pub fn encode(text: &str, config: QrConfig) -> Result<Self, String> {
        // 1. 分段
        let segments = segment_text(text);
        if segments.is_empty() {
            return Err("输入为空".into());
        }

        // 2. 确定版本
        let version = match config.version {
            VersionMode::Fixed(v) => Version::new(v).ok_or("无效版本号 (1-40)")?,
            VersionMode::Auto => {
                let mut selected = None;
                for v in 1..=40 {
                    let ver = Version(v);
                    let total_bits: u16 = segments.iter().map(|s| segment_bit_length(s, v)).sum();
                    let cap = get_data_capacity(ver, config.level) as u32 * 8;
                    if cap >= total_bits as u32 {
                        selected = Some(ver);
                        break;
                    }
                }
                selected.ok_or("数据过长，超出 QR 码最大容量".to_string())?
            }
        };

        // 3. 构建数据码字
        let data = build_codewords(text, version, config.level);

        // 4. 纠错编码 — 分组并计算 RS 纠错码
        let ec_info = version.ec_info(config.level);
        let mut blocks: Vec<Vec<u8>> = Vec::new();
        let mut pos = 0usize;
        for binfo in &ec_info.blocks {
            for _ in 0..binfo.count {
                let end = pos + binfo.data_codewords as usize;
                if end > data.len() {
                    return Err("内部错误: 数据码字不足".into());
                }
                blocks.push(data[pos..end].to_vec());
                pos = end;
            }
        }
        let final_codewords = reed_solomon::interleave(&blocks, ec_info.ec_per_block);

        // 5. 构建矩阵 + 放置功能图案
        let mut matrix = Matrix::new(version.size());
        place_finder_patterns(&mut matrix);
        place_timing_patterns(&mut matrix);
        place_alignment_patterns(&mut matrix, version.alignment_positions());
        reserve_format_areas(&mut matrix);
        if version.0 >= 7 {
            reserve_version_areas(&mut matrix);
        }

        // 6. 蛇形放置数据
        place_data(&mut matrix, &final_codewords);

        // 7. 掩码评分 → 选最佳
        let (best_idx, best_matrix) = best_mask(&matrix);

        // 8. 写入格式信息
        let format = encode_format_info(config.level.indicator_bits(), best_idx);
        let mut final_matrix = best_matrix;
        place_format_info(&mut final_matrix, format);

        // 9. 写入版本信息（版本 ≥ 7）
        if version.0 >= 7 {
            let ver_info = encode_version_info(version.0);
            place_version_info(&mut final_matrix, ver_info);
        }

        Ok(QrCode {
            version,
            level: config.level,
            mask: best_idx,
            matrix: final_matrix,
            margin: config.margin,
        })
    }

    pub fn modules(&self) -> &[Vec<bool>] {
        &self.matrix.modules
    }

    pub fn size(&self) -> u8 {
        self.matrix.size
    }

    pub fn to_svg(&self) -> String {
        crate::render::svg::render_svg(self)
    }

    pub fn to_ascii(&self, invert: bool) -> String {
        crate::render::ascii::render_ascii(self, invert)
    }

    pub fn to_png_bytes(&self, module_size: u8) -> Vec<u8> {
        crate::render::png::render_png(self, module_size)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_encode_simple() {
        let qr = QrCode::encode("HELLO", QrConfig::default()).unwrap();
        assert_eq!(qr.version.0, 1);
        assert_eq!(qr.size(), 21);
    }

    #[test]
    fn test_encode_numeric() {
        let qr = QrCode::encode("1234567890", QrConfig::default()).unwrap();
        assert_eq!(qr.size(), 21);
    }

    #[test]
    fn test_encode_empty_fails() {
        assert!(QrCode::encode("", QrConfig::default()).is_err());
    }

    #[test]
    fn test_fixed_version() {
        let config = QrConfig {
            version: VersionMode::Fixed(3),
            ..Default::default()
        };
        let qr = QrCode::encode("FIXED VERSION", config).unwrap();
        assert_eq!(qr.version.0, 3);
    }

    #[test]
    fn test_all_levels() {
        for level in [EcLevel::L, EcLevel::M, EcLevel::Q, EcLevel::H] {
            let config = QrConfig {
                level,
                ..Default::default()
            };
            let qr = QrCode::encode("TEST", config).unwrap();
            assert!(qr.size() >= 21);
        }
    }
}