penis 0.1.1

#![allow(clippy::needless_range_loop)]
use core::panic;

use crate::{PenisBlock, PrivateBlock, State};
use alloc::vec::Vec;

/// Partially Executed Nth-round Intermediate State
///
/// This struct represents the intermediate state of a SHA-256 hash computation that has been
/// partially executed up to a specified round n. It contains a sequence of blocks that are
/// either public (raw data) or private (intermediate state with partial schedule array).
///
/// The Penis structure enables verifiable hashing while maintaining confidentiality of
/// selected portions of the input data. It achieves this by storing:
/// - Public blocks: Contains raw data that can be freely shared
/// - Private blocks: Contains intermediate state and partial schedule array for confidential data
///
/// # Security Considerations
///
/// - The security of the protocol depends on the chosen round number (n)
/// - Private blocks should be properly encrypted during network transmission
/// - Timing attacks should be considered in security-critical applications
///
/// # Examples
///
/// ```ignore
/// use penis::{Penis, PenisBlock};
///
/// // Create a Penis instance from blocks
/// let penis = Penis(vec![
///     PenisBlock::Public([0u8; 64]),    // Public block
///     PenisBlock::Private(private_block) // Private block
/// ]);
///
/// // Encode for transmission
/// let compact = penis.encode_sparse();
///
/// // Decode received data
/// let decoded = Penis::decode_sparse(&compact);
/// ```
#[cfg_attr(test, derive(Debug, PartialEq))]
pub struct Penis(pub Vec<PenisBlock>);

impl Penis {
    /// Encodes the Penis structure into a compact byte format suitable for network transmission.
    ///
    /// The encoding format is:
    /// - For each block:
    ///   - 1 byte: block type (0 = public, 1 = private)
    ///   - For private blocks:
    ///     - 64 bytes: partial schedule array (16 u32 values)
    ///     - 32 bytes: intermediate state
    ///   - For public blocks:
    ///     - 64 bytes: raw block data
    ///
    /// # Returns
    ///
    /// A vector of bytes containing the encoded data
    pub fn encode_compact(&self) -> Vec<u8> {
        // Estimate the capacity needed for the compact_bytes vector
        let estimated_capacity = self.0.iter().map(|block| block.encoded_size() + 1).sum();

        let mut encoded_bytes = Vec::with_capacity(estimated_capacity);

        for block in &self.0 {
            match block {
                PenisBlock::Private(private_block) => {
                    encoded_bytes.push(1);
                    for &value in &private_block.psa {
                        encoded_bytes.extend(&value.to_be_bytes());
                    }
                    encoded_bytes.extend(private_block.state.finalize().as_ref());
                }
                PenisBlock::Public(raw_block) => {
                    encoded_bytes.push(0);
                    encoded_bytes.extend(raw_block);
                }
            }
        }
        encoded_bytes
    }

    /// Decodes a Penis structure from its compact byte representation.
    ///
    /// # Arguments
    ///
    /// * `compact_bytes` - The encoded byte array to decode
    ///
    /// # Returns
    ///
    /// A new Penis instance containing the decoded blocks
    ///
    /// # Panics
    ///
    /// Panics if the input bytes are not in the correct format or if an unknown block type is encountered
    pub fn decode_compact(compact_bytes: &[u8]) -> Self {
        let mut penis_blocks = Vec::new();
        let mut i = 0;

        while i < compact_bytes.len() {
            let block_type = compact_bytes[i];
            i += 1;

            match block_type {
                1 => {
                    // Decode PrivateBlock
                    let mut psa = [0u32; 16];
                    for j in 0..16 {
                        let base_ptr = i + j * 4;
                        psa[j] = u32::from_be_bytes([
                            compact_bytes[base_ptr],
                            compact_bytes[base_ptr + 1],
                            compact_bytes[base_ptr + 2],
                            compact_bytes[base_ptr + 3],
                        ]);
                    }
                    i += 64;
                    let state = State::from_bytes(&compact_bytes[i..i + 32]);
                    penis_blocks.push(PenisBlock::Private(PrivateBlock { psa, state }));
                    i += 32;
                }
                0 => {
                    // Decode PublicBlock
                    let mut raw_block = [0u8; 64];
                    raw_block.copy_from_slice(&compact_bytes[i..i + 64]);
                    penis_blocks.push(PenisBlock::Public(raw_block));
                    i += 64;
                }

                _ => panic!("Unknown block type"),
            }
        }

        Self(penis_blocks)
    }

    /// Encodes the Penis structure into a u32 array format, suitable for local processing.
    ///
    /// The encoding format is similar to compact encoding but uses u32 values directly:
    /// - For each block:
    ///   - 1 u32: block type (0 = public, 1 = private)
    ///   - For private blocks:
    ///     - 16 u32: partial schedule array
    ///     - 8 u32: intermediate state
    ///   - For public blocks:
    ///     - 16 u32: raw block data (converted from bytes)
    ///
    /// # Returns
    ///
    /// A vector of u32 values containing the encoded data
    pub fn encode_u32(&self) -> Vec<u32> {
        // Estimate the capacity needed for the compact_bytes vector
        let estimated_capacity = self
            .0
            .iter()
            .map(|block| block.encoded_size() / 4 + 1)
            .sum();
        let mut encoded_bytes = Vec::with_capacity(estimated_capacity);

        for block in &self.0 {
            match block {
                PenisBlock::Private(private_block) => {
                    encoded_bytes.push(1);
                    encoded_bytes.extend(private_block.psa);
                    encoded_bytes.extend(private_block.state.finalize_u32().as_ref());
                }
                PenisBlock::Public(raw_block) => {
                    encoded_bytes.push(0);
                    encoded_bytes.extend(raw_block.iter().map(|&x| x as u32));
                }
            }
        }
        encoded_bytes
    }

    /// Decodes a Penis structure from its u32 array representation.
    ///
    /// # Arguments
    ///
    /// * `u32_data` - The encoded u32 array to decode
    ///
    /// # Returns
    ///
    /// A new Penis instance containing the decoded blocks
    ///
    /// # Panics
    ///
    /// Panics if the input data is not in the correct format or if an unknown block type is encountered
    pub fn decode_u32(u32_data: &[u32]) -> Self {
        let mut penis_blocks = Vec::new();
        let mut i = 0;

        while i < u32_data.len() {
            let block_type = u32_data[i];
            i += 1;

            match block_type {
                1 => {
                    // Decode PrivateBlock
                    let mut psa = [0u32; 16];
                    psa.copy_from_slice(&u32_data[i..i + 16]);

                    // Create state from the next 8 u32 values
                    let state_data = &u32_data[i + 16..i + 24];
                    let state = State {
                        a: state_data[0],
                        b: state_data[1],
                        c: state_data[2],
                        d: state_data[3],
                        e: state_data[4],
                        f: state_data[5],
                        g: state_data[6],
                        h: state_data[7],
                    };

                    penis_blocks.push(PenisBlock::Private(PrivateBlock { psa, state }));
                    i += 24;
                }
                0 => {
                    // Decode PublicBlock
                    let mut raw_block = [0u8; 64];

                    // Convert 16 u32 values to 64 bytes
                    for j in 0..16 {
                        let bytes = u32_data[i + j].to_be_bytes();
                        raw_block[j * 4..j * 4 + 4].copy_from_slice(&bytes);
                    }

                    penis_blocks.push(PenisBlock::Public(raw_block));
                    i += 16;
                }
                _ => panic!("Unknown block type"),
            }
        }

        Self(penis_blocks)
    }

    /// Encodes the Penis structure into a sparse format, suitable when there are few private blocks
    ///
    /// The consecutive same-type private blocks are encoded without block flag. Instead of using block flag in each block, we use a length prefix to indicate how many blocks of the same type follow.
    ///
    /// - Consecutive group (repeat)
    ///   - type flag (1 bit)
    ///     - 0: This consecutive group is public blocks
    ///     - 1: This consecutive group is private blocks
    ///   - length prefix
    ///     - extended length flag (1 bit)
    ///      - 0: short form
    ///        - length field (6 bits)
    ///      - 1: extended form
    ///        - reserved for future use (6 bits) (all-zero for now)
    ///        - length field (64 bits)
    ///   - blocks (repeat)
    ///     - If type is public: 64 bytes of public block is repeated
    ///     - If type is private: 96 bytes of private block is repeated
    ///
    /// # Returns
    ///
    /// A vector of bytes containing the encoded data
    pub fn encode_sparse(&self) -> Vec<u8> {
        let mut encoded_bytes = Vec::new();
        let mut last_proc = 0; // Index of the last block that was processed
        let mut current_type = self.0[0].type_flag(); // Type flag of the current block type being processed

        while last_proc < self.0.len() {
            // counting phase
            let mut consec = 1;
            while last_proc + consec < self.0.len()
                && self.0[last_proc + consec].type_flag() == current_type
            {
                consec += 1;
            }

            // length encoding phase
            if consec < 64 {
                // If consecutive blocks are less than 64, encode the length as short form
                encoded_bytes.push((current_type << 7) | (consec as u8));
            } else {
                // Encode the length in extended form
                encoded_bytes.push((current_type << 7) | (1 << 6));
                let length_field: u64 = consec as u64;
                encoded_bytes.extend(&length_field.to_be_bytes());
            }

            // data encoding phase
            for i in 0..consec {
                match &self.0[last_proc + i] {
                    PenisBlock::Private(private_block) => {
                        for &value in &private_block.psa {
                            encoded_bytes.extend(&value.to_be_bytes());
                        }
                        encoded_bytes.extend(private_block.state.finalize().as_ref());
                    }
                    PenisBlock::Public(raw_block) => {
                        encoded_bytes.extend(raw_block);
                    }
                }
            }

            last_proc += consec;
            if last_proc >= self.0.len() {
                break;
            }
            current_type = self.0[last_proc].type_flag();
        }
        encoded_bytes
    }

    /// Decodes a Penis structure from its sparse representation.
    ///
    /// # Arguments
    ///
    /// * `sparse_bytes` - The encoded byte array to decode
    ///
    /// # Returns
    ///
    /// A new Penis instance containing the decoded blocks
    ///
    /// # Panics
    ///
    /// Panics if the input bytes are not in the correct format or if an unknown block type is encountered
    pub fn decode_sparse(sparse_bytes: &[u8]) -> Self {
        let mut penis_blocks = Vec::new();
        let mut i = 0;

        while i < sparse_bytes.len() {
            let header = sparse_bytes[i];
            i += 1;
            let block_type = (header & 0x80) >> 7;
            let length_flag = (header & 0x40) >> 6; // Extract the second highest bit

            let consecutive_count: usize = if length_flag == 0 {
                (header & 0b111111).into()
            } else {
                if (header & 0b111111) != 0 {
                    panic!("Invalid format: RFU field must be zero");
                }
                let mut extended_length_bytes = [0u8; 8];
                extended_length_bytes[0..8].copy_from_slice(&sparse_bytes[i..i + 8]);
                i += 8;
                let len = u64::from_be_bytes(extended_length_bytes) as usize;
                if len < 64 {
                    panic!("Invalid format: length when in extended length must be at least 64");
                }
                len
            };

            if consecutive_count == 0 {
                panic!("Invalid format: consecutive_count cannot be zero");
            }

            for _ in 0..consecutive_count {
                match block_type {
                    1 => {
                        // Decode PrivateBlock
                        let mut psa = [0u32; 16];
                        for j in 0..16 {
                            psa[j] = u32::from_be_bytes([
                                sparse_bytes[i],
                                sparse_bytes[i + 1],
                                sparse_bytes[i + 2],
                                sparse_bytes[i + 3],
                            ]);
                            i += 4;
                        }
                        let state = State::from_bytes(&sparse_bytes[i..i + 32]);
                        i += 32;
                        penis_blocks.push(PenisBlock::Private(PrivateBlock { psa, state }));
                    }
                    0 => {
                        // Decode PublicBlock
                        let mut raw_block = [0u8; 64];
                        raw_block.copy_from_slice(&sparse_bytes[i..i + 64]);
                        i += 64;
                        penis_blocks.push(PenisBlock::Public(raw_block));
                    }
                    _ => unreachable!("Unknown block type"),
                }
            }
        }

        Self(penis_blocks)
    }
}

#[cfg(test)]
mod tests {
    use alloc::vec;

    use super::*;

    #[test]
    fn test_sparse_encoding() {
        // Create test data with mixed public and private blocks
        let private1 = PrivateBlock {
            psa: [1u32; 16],
            state: State {
                a: 0x1111,
                b: 0x2222,
                c: 0x3333,
                d: 0x4444,
                e: 0x5555,
                f: 0x6666,
                g: 0x7777,
                h: 0x8888,
            },
        };
        let private2 = PrivateBlock {
            psa: [2u32; 16],
            state: State {
                a: 0x1111,
                b: 0x2222,
                c: 0x3333,
                d: 0x4444,
                e: 0x5555,
                f: 0x6666,
                g: 0x7777,
                h: 0x8888,
            },
        };
        let private3 = PrivateBlock {
            psa: [3u32; 16],
            state: State {
                a: 0x1111,
                b: 0x2222,
                c: 0x3333,
                d: 0x4444,
                e: 0x5555,
                f: 0x6666,
                g: 0x7777,
                h: 0x8888,
            },
        };

        let public1 = [1u8; 64];
        let public2 = [2u8; 64];
        let public3 = [3u8; 64];
        let public4 = [4u8; 64];

        let penis = Penis(vec![
            PenisBlock::Public(public1),
            PenisBlock::Public(public2),
            PenisBlock::Private(private1),
            PenisBlock::Private(private2),
            PenisBlock::Public(public3),
            PenisBlock::Private(private3),
            PenisBlock::Public(public4),
        ]);

        // Test encoding and decoding
        let encoded = penis.encode_sparse();
        let decoded = Penis::decode_sparse(&encoded);

        assert_eq!(penis, decoded);

        let expected_length = 5 + (4 * 64) + (3 * 96);
        assert_eq!(encoded.len(), expected_length);
    }

    #[test]
    fn test_sparse_encoding_long_sequence() {
        // Create a long sequence of 100 public blocks
        let public_blocks: Vec<PenisBlock> = (0..100)
            .map(|i| PenisBlock::Public([i as u8; 64]))
            .collect();

        // Create a long sequence of 70 private blocks
        let private_blocks: Vec<PenisBlock> = (0..70)
            .map(|i| {
                PenisBlock::Private(PrivateBlock {
                    psa: [i as u32; 16],
                    state: State {
                        a: 0x1111,
                        b: 0x2222,
                        c: 0x3333,
                        d: 0x4444,
                        e: 0x5555,
                        f: 0x6666,
                        g: 0x7777,
                        h: 0x8888,
                    },
                })
            })
            .collect();

        // Combine blocks with mixed long sequences
        let mut all_blocks = Vec::new();
        all_blocks.extend(public_blocks);
        all_blocks.extend(private_blocks);

        let penis = Penis(all_blocks);

        // Test encoding and decoding
        let encoded = penis.encode_sparse();
        let decoded = Penis::decode_sparse(&encoded);

        assert_eq!(penis, decoded);

        // Verify structure of encoded data
        // Header for public blocks (extended form): 2 bytes (header + RFU) + 8 bytes length
        // Header for private blocks (extended form): 2 bytes (header + RFU) + 8 bytes length
        // Data: 100 public blocks * 64 bytes + 70 private blocks * 96 bytes
        let expected_length = (1 + 8) + (100 * 64) + (1 + 8) + (70 * 96);
        assert_eq!(encoded.len(), expected_length);

        // Check that the first byte has extended length flag set
        assert_eq!(encoded[0] & 0x40, 0x40);
    }

    #[test]
    fn test_compact_encoding() {
        // Create test data
        let private = PrivateBlock {
            psa: [0xDEADBEEF; 16],
            state: State {
                a: 0x67452301,
                b: 0xEFCDAB89,
                c: 0x98BADCFE,
                d: 0x10325476,
                e: 0xC3D2E1F0,
                f: 0x1F83D9AB,
                g: 0x2B87E4CD,
                h: 0x3C9BE0F1,
            },
        };

        let public = [0xAA; 64];

        let penis = Penis(vec![
            PenisBlock::Public(public),
            PenisBlock::Private(private),
            PenisBlock::Public(public),
        ]);

        // Test encoding and decoding
        let encoded = penis.encode_compact();
        let decoded = Penis::decode_compact(&encoded);

        assert_eq!(penis, decoded);

        // Verify expected length
        // Each block: 1 byte type + data (64 bytes public or 96 bytes private)
        let expected_length = (1 + 64) + (1 + 96) + (1 + 64);
        assert_eq!(encoded.len(), expected_length);
    }

    #[test]
    #[should_panic(expected = "Invalid format: RFU field must be zero")]
    fn test_sparse_decode_invalid_rfu() {
        // Create a header with extended length flag set and non-zero RFU bits
        let mut encoded = vec![0b01100000];
        encoded.extend([0u8; 8]); // Extended length bytes
        Penis::decode_sparse(&encoded);
    }

    #[test]
    #[should_panic(expected = "Invalid format: consecutive_count cannot be zero")]
    fn test_sparse_decode_zero_count() {
        // Create a header with count = 0
        let encoded = vec![0b00000000]; // Type=0, Extended=0, Count=0
        Penis::decode_sparse(&encoded);
    }
    #[test]
    #[should_panic(expected = "Invalid format: length when in extended length must be at least 64")]
    fn test_sparse_decode_zero_count_long() {
        // Create a header with extended length flag set and non-zero RFU bits
        let mut encoded = vec![0b11000000];
        encoded.extend([0u8; 8]); // Extended length bytes
        Penis::decode_sparse(&encoded);
    }
    #[test]
    #[should_panic(expected = "length when in extended length must be at least 64")]
    fn test_sparse_decode_invalid_len_long() {
        // Create a header with extended length flag set and non-zero RFU bits
        let mut encoded = vec![0b11000000];
        encoded.extend((1u64).to_be_bytes()); // Extended length bytes
        Penis::decode_sparse(&encoded);
    }

    #[test]
    #[should_panic(expected = "index out of bounds: the len is 1 but the index is 1")]
    fn test_sparse_decode_invalid_format() {
        // Create a header with extended length flag set and non-zero RFU bits
        let encoded = vec![0b10000001];
        Penis::decode_sparse(&encoded);
    }
    #[test]
    #[should_panic(expected = "index out of bounds: the len is 9 but the index is 9")]
    fn test_sparse_decode_invalid_format_long() {
        // Create a header with extended length flag set and non-zero RFU bits
        let mut encoded = vec![0b11000000];
        encoded.extend((64u64).to_be_bytes()); // Extended length bytes
        Penis::decode_sparse(&encoded);
    }

    #[test]
    #[should_panic]
    fn test_sparse_decode_truncated_data() {
        // Create a valid header but truncate the data
        let mut encoded = vec![];
        encoded.push(0b00000010); // Type=0, Count=2
        encoded.extend([1u8; 64]); // Only one block worth of data
        Penis::decode_sparse(&encoded);
    }
}