promocrypt_core/

generator.rs

//! Code generation using HMAC-SHA256.
//!
//! This module provides functions for generating promotional codes
//! from counter values using HMAC-SHA256 and the Damm check digit algorithm.

use ring::hmac;
use serde::{Deserialize, Serialize};

use crate::alphabet::Alphabet;
use crate::damm::DammTable;

/// Index-based check digit position.
///
/// The check digit can be placed at any position in the code using positive
/// or negative indexing:
/// - `0`: Start (first position)
/// - `-1`: End (last position, default)
/// - Positive N: Position N from start
/// - Negative -N: Position N from end
///
/// # Examples
/// ```
/// use promocrypt_core::CheckPosition;
///
/// let pos = CheckPosition::End;
/// assert_eq!(pos.to_index(10), 9);
///
/// let pos = CheckPosition::Start;
/// assert_eq!(pos.to_index(10), 0);
///
/// let pos = CheckPosition::Index(4);
/// assert_eq!(pos.to_index(10), 4);
///
/// let pos = CheckPosition::Index(-3);
/// assert_eq!(pos.to_index(10), 7);
/// ```
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum CheckPosition {
    /// Check digit at the start of the code (index 0)
    #[serde(rename = "start")]
    Start,
    /// Check digit at the end of the code (index -1)
    #[serde(rename = "end")]
    #[default]
    End,
    /// Check digit at specific index (positive or negative)
    Index(i8),
}

impl CheckPosition {
    /// Create from raw index value.
    pub fn new(index: i8) -> Self {
        match index {
            0 => CheckPosition::Start,
            -1 => CheckPosition::End,
            n => CheckPosition::Index(n),
        }
    }

    /// Get actual index given total code length (including check digit).
    pub fn to_index(&self, total_length: usize) -> usize {
        match self {
            CheckPosition::Start => 0,
            CheckPosition::End => total_length.saturating_sub(1),
            CheckPosition::Index(idx) => {
                if *idx >= 0 {
                    (*idx as usize).min(total_length.saturating_sub(1))
                } else {
                    let from_end = (-*idx) as usize;
                    total_length.saturating_sub(from_end)
                }
            }
        }
    }

    /// Get raw index value.
    pub fn raw(&self) -> i8 {
        match self {
            CheckPosition::Start => 0,
            CheckPosition::End => -1,
            CheckPosition::Index(n) => *n,
        }
    }

    /// Parse from string (for CLI and config).
    pub fn parse_str(s: &str) -> Option<Self> {
        match s.to_lowercase().as_str() {
            "start" | "beginning" | "front" | "0" => Some(CheckPosition::Start),
            "end" | "back" | "tail" | "-1" => Some(CheckPosition::End),
            _ => {
                // Try parsing as integer
                s.parse::<i8>().ok().map(CheckPosition::new)
            }
        }
    }
}

impl std::str::FromStr for CheckPosition {
    type Err = String;

    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        Self::parse_str(s).ok_or_else(|| format!("Invalid check position: {}", s))
    }
}

impl std::fmt::Display for CheckPosition {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            CheckPosition::Start => write!(f, "start"),
            CheckPosition::End => write!(f, "end"),
            CheckPosition::Index(n) => write!(f, "{}", n),
        }
    }
}

/// Code format options for prefix, suffix, and separators.
///
/// # Examples
/// ```
/// use promocrypt_core::CodeFormat;
///
/// let format = CodeFormat::new()
///     .with_prefix("PROMO-")
///     .with_suffix("-2024")
///     .with_separator('-', vec![4, 8]);
///
/// assert_eq!(format.format("A3KF7NP2XM"), "PROMO-A3KF-7NP2-XM-2024");
/// ```
#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct CodeFormat {
    /// Optional prefix prepended to code
    #[serde(skip_serializing_if = "Option::is_none")]
    pub prefix: Option<String>,
    /// Optional suffix appended to code
    #[serde(skip_serializing_if = "Option::is_none")]
    pub suffix: Option<String>,
    /// Separator character
    #[serde(skip_serializing_if = "Option::is_none")]
    pub separator: Option<char>,
    /// Positions where separator is inserted (0-based, before the character at that position)
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub separator_positions: Vec<usize>,
}

impl CodeFormat {
    /// Create a new empty code format.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set prefix.
    pub fn with_prefix(mut self, prefix: &str) -> Self {
        self.prefix = Some(prefix.to_string());
        self
    }

    /// Set suffix.
    pub fn with_suffix(mut self, suffix: &str) -> Self {
        self.suffix = Some(suffix.to_string());
        self
    }

    /// Set separator and positions.
    pub fn with_separator(mut self, sep: char, positions: Vec<usize>) -> Self {
        self.separator = Some(sep);
        self.separator_positions = positions;
        self
    }

    /// Apply formatting to base code.
    pub fn format(&self, base_code: &str) -> String {
        let chars: Vec<char> = base_code.chars().collect();
        let sep_count = self
            .separator_positions
            .iter()
            .filter(|&&p| p > 0 && p < chars.len())
            .count();
        let prefix_len = self.prefix.as_ref().map(|p| p.len()).unwrap_or(0);
        let suffix_len = self.suffix.as_ref().map(|s| s.len()).unwrap_or(0);

        let mut result = String::with_capacity(prefix_len + chars.len() + sep_count + suffix_len);

        // Add prefix
        if let Some(ref prefix) = self.prefix {
            result.push_str(prefix);
        }

        // Add base code with separators
        if let Some(sep) = self.separator {
            for (i, c) in chars.iter().enumerate() {
                if i > 0 && self.separator_positions.contains(&i) {
                    result.push(sep);
                }
                result.push(*c);
            }
        } else {
            result.push_str(base_code);
        }

        // Add suffix
        if let Some(ref suffix) = self.suffix {
            result.push_str(suffix);
        }

        result
    }

    /// Strip formatting from code to get base code.
    pub fn strip(&self, formatted_code: &str) -> Option<String> {
        let mut code = formatted_code.to_string();

        // Strip prefix
        if let Some(ref prefix) = self.prefix {
            code = code.strip_prefix(prefix)?.to_string();
        }

        // Strip suffix
        if let Some(ref suffix) = self.suffix {
            code = code.strip_suffix(suffix)?.to_string();
        }

        // Strip separators
        if let Some(sep) = self.separator {
            code = code.chars().filter(|&c| c != sep).collect();
        }

        Some(code)
    }

    /// Calculate expected total length of formatted code.
    pub fn total_length(&self, base_length: usize) -> usize {
        let prefix_len = self.prefix.as_ref().map(|p| p.len()).unwrap_or(0);
        let suffix_len = self.suffix.as_ref().map(|s| s.len()).unwrap_or(0);
        let sep_count = self
            .separator_positions
            .iter()
            .filter(|&&p| p > 0 && p < base_length)
            .count();

        prefix_len + base_length + sep_count + suffix_len
    }

    /// Check if any formatting is applied.
    pub fn has_formatting(&self) -> bool {
        self.prefix.is_some() || self.suffix.is_some() || self.separator.is_some()
    }
}

/// Generate a single promotional code from a counter value.
///
/// # Arguments
/// * `secret_key` - 32-byte secret key for HMAC
/// * `counter` - Counter value (unique per code)
/// * `alphabet` - Character set for the code
/// * `code_length` - Number of random characters (check digit added separately)
/// * `check_position` - Where to place the check digit (index-based)
/// * `damm_table` - Damm table for check digit calculation
///
/// # Returns
/// Generated code string with length = code_length + 1 (including check digit)
///
/// # Example
///
/// ```
/// use promocrypt_core::{generate_code, Alphabet, DammTable, CheckPosition};
///
/// let secret = [0u8; 32];
/// let alphabet = Alphabet::default_alphabet();
/// let damm = DammTable::new(alphabet.len());
///
/// let code = generate_code(&secret, 0, &alphabet, 9, CheckPosition::End, &damm);
/// assert_eq!(code.len(), 10);
/// ```
pub fn generate_code(
    secret_key: &[u8; 32],
    counter: u64,
    alphabet: &Alphabet,
    code_length: usize,
    check_position: CheckPosition,
    damm_table: &DammTable,
) -> String {
    // 1. HMAC-SHA256 of counter
    let key = hmac::Key::new(hmac::HMAC_SHA256, secret_key);
    let signature = hmac::sign(&key, &counter.to_le_bytes());

    // 2. Convert first 8 bytes to u64
    let hash_bytes: [u8; 8] = signature.as_ref()[0..8].try_into().unwrap();
    let mut value = u64::from_le_bytes(hash_bytes);

    // 3. Convert to base-N characters
    let base = alphabet.len() as u64;
    let mut chars = Vec::with_capacity(code_length);

    for _ in 0..code_length {
        let index = (value % base) as usize;
        chars.push(alphabet.char_at(index));
        value /= base;
    }

    // 4. Get check position index
    let total_length = code_length + 1;
    let check_index = check_position.to_index(total_length);

    // 5. Build code with check digit placeholder to calculate check
    // For any position, we need to calculate check so that Damm validation passes
    let check = if check_index == 0 {
        // Check at start: find X such that processing [X, chars...] = 0
        damm_table.calculate_for_start(&chars, alphabet)
    } else if check_index >= code_length {
        // Check at end: process chars and get final check digit
        damm_table.calculate(&chars, alphabet)
    } else {
        // Check in middle: need to find check digit for middle position
        // We calculate check digit such that Damm validation passes
        damm_table.calculate_for_position(&chars, alphabet, check_index)
    };

    // 6. Build final code with check digit at specified position
    let mut result = String::with_capacity(total_length);

    for i in 0..total_length {
        if i == check_index {
            result.push(check);
        } else {
            // Adjust index for chars array based on check position
            let char_idx = if i < check_index { i } else { i - 1 };
            if char_idx < chars.len() {
                result.push(chars[char_idx]);
            }
        }
    }

    result
}

/// Generate a batch of promotional codes.
///
/// # Arguments
/// * `secret_key` - 32-byte secret key for HMAC
/// * `start_counter` - Starting counter value
/// * `count` - Number of codes to generate
/// * `alphabet` - Character set for the codes
/// * `code_length` - Number of random characters per code
/// * `check_position` - Where to place the check digit
/// * `damm_table` - Damm table for check digit calculation
///
/// # Returns
/// Vector of generated codes
///
/// # Example
///
/// ```
/// use promocrypt_core::{generate_batch, Alphabet, DammTable, CheckPosition};
///
/// let secret = [0u8; 32];
/// let alphabet = Alphabet::default_alphabet();
/// let damm = DammTable::new(alphabet.len());
///
/// let codes = generate_batch(&secret, 0, 100, &alphabet, 9, CheckPosition::End, &damm);
/// assert_eq!(codes.len(), 100);
/// ```
pub fn generate_batch(
    secret_key: &[u8; 32],
    start_counter: u64,
    count: usize,
    alphabet: &Alphabet,
    code_length: usize,
    check_position: CheckPosition,
    damm_table: &DammTable,
) -> Vec<String> {
    (0..count)
        .map(|i| {
            generate_code(
                secret_key,
                start_counter + i as u64,
                alphabet,
                code_length,
                check_position,
                damm_table,
            )
        })
        .collect()
}

/// Generate codes into a pre-allocated vector for better performance.
///
/// This is useful when generating very large batches where allocation
/// overhead matters.
#[allow(clippy::too_many_arguments)]
pub fn generate_batch_into(
    secret_key: &[u8; 32],
    start_counter: u64,
    count: usize,
    alphabet: &Alphabet,
    code_length: usize,
    check_position: CheckPosition,
    damm_table: &DammTable,
    output: &mut Vec<String>,
) {
    output.clear();
    output.reserve(count);

    for i in 0..count {
        output.push(generate_code(
            secret_key,
            start_counter + i as u64,
            alphabet,
            code_length,
            check_position,
            damm_table,
        ));
    }
}

/// Iterator for generating codes lazily.
///
/// Useful for streaming large batches without holding all codes in memory.
pub struct CodeGenerator<'a> {
    secret_key: &'a [u8; 32],
    alphabet: &'a Alphabet,
    code_length: usize,
    check_position: CheckPosition,
    damm_table: &'a DammTable,
    current_counter: u64,
    end_counter: u64,
}

impl<'a> CodeGenerator<'a> {
    /// Create a new code generator iterator.
    pub fn new(
        secret_key: &'a [u8; 32],
        start_counter: u64,
        count: usize,
        alphabet: &'a Alphabet,
        code_length: usize,
        check_position: CheckPosition,
        damm_table: &'a DammTable,
    ) -> Self {
        Self {
            secret_key,
            alphabet,
            code_length,
            check_position,
            damm_table,
            current_counter: start_counter,
            end_counter: start_counter + count as u64,
        }
    }
}

impl Iterator for CodeGenerator<'_> {
    type Item = String;

    fn next(&mut self) -> Option<Self::Item> {
        if self.current_counter >= self.end_counter {
            return None;
        }

        let code = generate_code(
            self.secret_key,
            self.current_counter,
            self.alphabet,
            self.code_length,
            self.check_position,
            self.damm_table,
        );

        self.current_counter += 1;
        Some(code)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let remaining = (self.end_counter - self.current_counter) as usize;
        (remaining, Some(remaining))
    }
}

impl ExactSizeIterator for CodeGenerator<'_> {}

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

    fn setup() -> (Alphabet, DammTable, [u8; 32]) {
        let alphabet = Alphabet::default_alphabet();
        let damm = DammTable::new(alphabet.len());
        let secret = [42u8; 32];
        (alphabet, damm, secret)
    }

    #[test]
    fn test_generate_code_length() {
        let (alphabet, damm, secret) = setup();

        let code = generate_code(&secret, 0, &alphabet, 9, CheckPosition::End, &damm);
        assert_eq!(code.len(), 10); // 9 random + 1 check
    }

    #[test]
    fn test_generate_code_deterministic() {
        let (alphabet, damm, secret) = setup();

        let code1 = generate_code(&secret, 12345, &alphabet, 9, CheckPosition::End, &damm);
        let code2 = generate_code(&secret, 12345, &alphabet, 9, CheckPosition::End, &damm);

        assert_eq!(code1, code2);
    }

    #[test]
    fn test_different_counters_different_codes() {
        let (alphabet, damm, secret) = setup();

        let code1 = generate_code(&secret, 0, &alphabet, 9, CheckPosition::End, &damm);
        let code2 = generate_code(&secret, 1, &alphabet, 9, CheckPosition::End, &damm);

        assert_ne!(code1, code2);
    }

    #[test]
    fn test_different_secrets_different_codes() {
        let (alphabet, damm, _) = setup();

        let secret1 = [1u8; 32];
        let secret2 = [2u8; 32];

        let code1 = generate_code(&secret1, 0, &alphabet, 9, CheckPosition::End, &damm);
        let code2 = generate_code(&secret2, 0, &alphabet, 9, CheckPosition::End, &damm);

        assert_ne!(code1, code2);
    }

    #[test]
    fn test_check_position_start() {
        let (alphabet, damm, secret) = setup();

        let code = generate_code(&secret, 0, &alphabet, 9, CheckPosition::Start, &damm);
        assert_eq!(code.len(), 10);

        // Verify it validates correctly
        assert!(damm.validate(&code, &alphabet));
    }

    #[test]
    fn test_check_position_end() {
        let (alphabet, damm, secret) = setup();

        let code = generate_code(&secret, 0, &alphabet, 9, CheckPosition::End, &damm);
        assert_eq!(code.len(), 10);

        // Verify it validates correctly
        assert!(damm.validate(&code, &alphabet));
    }

    #[test]
    fn test_all_characters_in_alphabet() {
        let (alphabet, damm, secret) = setup();

        let code = generate_code(&secret, 0, &alphabet, 9, CheckPosition::End, &damm);

        for c in code.chars() {
            assert!(alphabet.contains(c), "Character '{}' not in alphabet", c);
        }
    }

    #[test]
    fn test_generate_batch() {
        let (alphabet, damm, secret) = setup();

        let codes = generate_batch(&secret, 0, 100, &alphabet, 9, CheckPosition::End, &damm);

        assert_eq!(codes.len(), 100);

        // All codes should be unique
        let unique: std::collections::HashSet<_> = codes.iter().collect();
        assert_eq!(unique.len(), 100);

        // All codes should be valid
        for code in &codes {
            assert!(damm.validate(code, &alphabet));
        }
    }

    #[test]
    fn test_code_generator_iterator() {
        let (alphabet, damm, secret) = setup();

        let code_gen = CodeGenerator::new(&secret, 0, 10, &alphabet, 9, CheckPosition::End, &damm);

        let codes: Vec<_> = code_gen.collect();
        assert_eq!(codes.len(), 10);

        // Should match batch generation
        let batch = generate_batch(&secret, 0, 10, &alphabet, 9, CheckPosition::End, &damm);
        assert_eq!(codes, batch);
    }

    #[test]
    fn test_code_generator_exact_size() {
        let (alphabet, damm, secret) = setup();

        let code_gen = CodeGenerator::new(&secret, 0, 50, &alphabet, 9, CheckPosition::End, &damm);
        assert_eq!(code_gen.len(), 50);
    }

    #[test]
    fn test_various_code_lengths() {
        let (alphabet, damm, secret) = setup();

        for length in [4, 6, 8, 9, 12, 16, 20] {
            let code = generate_code(&secret, 0, &alphabet, length, CheckPosition::End, &damm);
            assert_eq!(code.len(), length + 1);
            assert!(damm.validate(&code, &alphabet));
        }
    }

    #[test]
    fn test_check_position_from_str() {
        assert_eq!(
            CheckPosition::parse_str("start"),
            Some(CheckPosition::Start)
        );
        assert_eq!(CheckPosition::parse_str("end"), Some(CheckPosition::End));
        assert_eq!(
            CheckPosition::parse_str("beginning"),
            Some(CheckPosition::Start)
        );
        assert_eq!(CheckPosition::parse_str("invalid"), None);
    }
}