arcium-macros 0.5.3

//! Automatic generation of Rust output types from circuit interfaces.
//!
//! This module reads circuit interface files (`.idarc`) and generates corresponding Rust structs
//! for computation outputs. The key challenge: Anchor doesn't support tuples, so all outputs are
//! converted to structs with `field_N` naming.
//!
//! **Encryption Pattern Detection**: The system semantically analyzes the interface to detect
//! encryption patterns (PublicKey + Nonce + Ciphertexts) and generates specialized types like
//! `SharedEncryptedStruct<N>` and `MXEEncryptedStruct<N>` instead of expanding them as raw fields.
//!
//! See `gen_callback_output_struct` documentation for comprehensive naming and type generation
//! rules.

use crate::utils::read_conf_ix_interface;
use arcis_interface::{CircuitInterface, ScalarKind, Value};
use convert_case::{Case, Casing};
use std::collections::HashSet;

/// Size in bits for nonce values used in Arcium's encryption schemes.
/// 128-bit unsigned scalars are recognized as nonces during encryption pattern detection.
const NONCE_SIZE_BITS: usize = 128;

/// Calculates the size in bytes for a given Value type.
///
/// All types in circuit interfaces are constant-sized, enabling compile-time size calculation.
fn value_size_in_bytes(val: &Value) -> usize {
    match val {
        Value::Ciphertext { .. } | Value::MScalar { .. } | Value::MFloat { .. } | Value::MBool => {
            32
        }
        Value::Scalar { size_in_bits, .. } => *size_in_bits / 8,
        Value::Float { size_in_bits } => *size_in_bits / 8,
        Value::Bool => 1,
        Value::ArcisX25519Pubkey => 32,
        Value::Point => 32,
        Value::Array(inner) => {
            if inner.is_empty() {
                0
            } else {
                value_size_in_bytes(&inner[0]) * inner.len()
            }
        }
        Value::Tuple(inner) => inner.iter().map(value_size_in_bytes).sum(),
        Value::Struct(inner) => {
            // Check for encryption patterns first
            match EncryptionComponents::from_value(&Value::Struct(inner.clone())).get_type() {
                // SharedEncryptedStruct: 32 (pubkey) + 16 (nonce) + N * 32 (ciphertexts)
                Some((EncryptionType::Shared, len)) => 32 + 16 + len * 32,
                // MXEEncryptedStruct: 16 (nonce) + N * 32 (ciphertexts)
                Some((EncryptionType::Mxe, len)) => 16 + len * 32,
                // EncDataStruct: N * 32 (ciphertexts)
                Some((EncryptionType::EncData, len)) => len * 32,
                None => inner.iter().map(value_size_in_bytes).sum(),
            }
        }
    }
}

/// Semantic encryption components extracted from circuit interface structures.
///
/// Used to detect patterns like `{PublicKey, Nonce, [Ciphertext]}` and generate specialized types.
#[derive(Debug)]
struct EncryptionComponents {
    has_public_key: bool,
    has_nonce: bool,
    ciphertext_count: usize,
}

impl EncryptionComponents {
    fn new() -> Self {
        Self {
            has_public_key: false,
            has_nonce: false,
            ciphertext_count: 0,
        }
    }

    /// Extract encryption components from a Value and return a new EncryptionComponents instance
    fn from_value(value: &Value) -> Self {
        let mut components = Self::new();
        components.extract_from(value);
        components
    }

    /// Get the encryption type if this represents a valid encryption pattern
    fn get_type(&self) -> Option<(EncryptionType, usize)> {
        match (
            self.has_public_key,
            self.has_nonce,
            self.ciphertext_count > 0,
        ) {
            (true, true, true) => Some((EncryptionType::Shared, self.ciphertext_count)),
            (false, true, true) => Some((EncryptionType::Mxe, self.ciphertext_count)),
            (false, false, true) => Some((EncryptionType::EncData, self.ciphertext_count)),
            _ => None,
        }
    }

    /// Extracts encryption components from a value by traversing its structure.
    ///
    /// # Encryption Component Detection
    /// - `PublicKey` → Sets `has_public_key = true`
    /// - 128-bit unsigned `Scalar` → Sets `has_nonce = true` (nonces are always 128-bit)
    /// - `Ciphertext` → Increments `ciphertext_count`
    /// - Composite types (Array, Struct, Tuple) → Recursively searched
    /// - Empty arrays are skipped (used as phantom data markers)
    fn extract_from(&mut self, value: &Value) {
        match value {
            Value::ArcisX25519Pubkey => {
                self.has_public_key = true;
            }
            // 128-bit unsigned scalars are used as nonces in Arcium's encryption scheme
            Value::Scalar {
                size_in_bits: NONCE_SIZE_BITS,
                ..
            } => {
                self.has_nonce = true;
            }
            Value::Ciphertext { .. } => {
                self.ciphertext_count += 1;
            }
            Value::Array(values) if !values.is_empty() => {
                // Recursively search arrays; empty arrays are phantom data and skipped
                for v in values {
                    self.extract_from(v);
                }
            }
            Value::Struct(values) | Value::Tuple(values) => {
                // Recursively search composite types
                for v in values {
                    self.extract_from(v);
                }
            }
            // Primitive types (Bool, Float, etc.) don't contribute to encryption semantics
            _ => {}
        }
    }
}

/// Generates a struct definition for the outputs of a given circuit interface
///
/// ## Output Type Naming Specification
///
/// This function generates output types for encrypted computations with comprehensive naming
/// conventions that integrate with the Arcium callback system. The naming follows predictable
/// patterns tested with real circuit examples.
///
/// ### Encrypted Interface Integration
///
/// The generated output structs work with encrypted computation callbacks:
///
/// ```rust
/// // For encrypted interface named "complex_example"
/// #[callback_accounts("complex_example", payer)]
/// pub struct ComplexExampleCallback<'info> {
///     // ... account definitions
/// }
///
/// #[arcium_callback(encrypted_ix = "complex_example")]
/// pub fn complex_example_callback(
///     ctx: Context<ComplexExampleCallback>,
///     output: SignedComputationOutputs<ComplexExampleOutput>,  // <- Generated by this function
/// ) {
///     // Handle computation results
/// }
/// ```
///
/// ### Main Output Struct Naming
/// - **Pattern**: `{CircuitName}Output`
/// - **Example**: For circuit "complex_example" → `ComplexExampleOutput`
/// - **Conversion**: Uses PascalCase conversion from the circuit interface name
/// - **Integration**: Must match the callback function's second parameter type
///
/// ### Field Naming in Output Structs
/// - **Pattern**: `field_{index}`
/// - **Example**: `field_0`, `field_1`, `field_2`, etc.
/// - **Rationale**: Anchor doesn't support tuples/tuple structs, so we use numbered fields
/// - **Usage**: Access results via `output.field_0`, `output.field_1`, etc.
///
/// ### Important: Tuple Output Behavior
/// **When a circuit returns a tuple, the entire tuple becomes a single `field_0`:**
///
/// ```rust
/// // Circuit returning: (UserData, Enc<Shared, u32>, (u64, f32), Enc<Mxe, bool>)
/// // Generates:
/// pub struct ComplexExampleOutput {
///     pub field_0: ComplexExampleOutputStruct0,  // The entire tuple as one field
/// }
///
/// // The ComplexExampleOutputStruct0 expands to:
/// pub struct ComplexExampleOutputStruct0 {
///     pub field_0: ComplexExampleOutputStruct00,  // UserData struct
///     pub field_1: SharedEncryptedStruct<1>,      // Enc<Shared, u32>
///     pub field_2: ComplexExampleOutputStruct02,  // (u64, f32) tuple
///     pub field_3: MXEEncryptedStruct<1>,         // Enc<Mxe, bool>
/// }
/// ```
///
/// ### Custom Struct Types Generated
///
/// #### Output Structs
/// - **Pattern**: `{CircuitName}OutputStruct{index}`
/// - **Example**: For circuit "vote" with output at index 0 → `VoteOutputStruct0`
/// - **Note**: Both single values and tuples use this naming since outputs are always single
///   elements or tuples
/// - **Fields**: Each field follows `field_{index}` pattern
///
/// #### Nested Structs
/// - **Pattern**: `{ParentName}OutputStruct{counter}`
/// - **Example**: Nested struct in `VoteOutputStruct0` → `VoteOutputStruct00`, `VoteOutputStruct01`
/// - **Counter**: Incremented based on seen structs to ensure uniqueness
///
/// ### Special Encryption Types
///
/// The system automatically recognizes encryption patterns and generates specialized types:
///
/// #### SharedEncryptedStruct<N>
/// - **Pattern**: Has `PublicKey` + `Scalar(128-bit nonce)` + N `Ciphertext` values
/// - **Generated Type**: `SharedEncryptedStruct<N>` where N = ciphertext count
/// - **Actual Struct**:
/// ```rust
/// pub struct SharedEncryptedStruct<const LEN: usize> {
///   pub encryption_key: [u8; 32],  // PublicKey (255-bit)
///   pub nonce: u128,               // 128-bit
///   pub ciphertexts: [[u8; 32]; LEN], // Array of ciphertexts
/// }
/// ```
/// - **Use Case**: Shared encryption where multiple parties can decrypt with the same key
///
/// #### MXEEncryptedStruct<N>
/// - **Pattern**: Has `Scalar(128-bit nonce)` + N `Ciphertext` values (no PublicKey)
/// - **Generated Type**: `MXEEncryptedStruct<N>` where N = ciphertext count
/// - **Actual Struct**:
/// ```rust
/// pub struct MXEEncryptedStruct<const LEN: usize> {
///   pub nonce: u128,  // 128-bit scalar
///   pub ciphertexts: [[u8; 32]; LEN], // Array of ciphertexts
/// }
/// ```
/// - **Use Case**: Multi-party exchange encryption without shared public key
///
/// #### EncDataStruct<N>
/// - **Pattern**: Has only N `Ciphertext` values (no PublicKey or nonce)
/// - **Generated Type**: `EncDataStruct<N>` where N = ciphertext count
/// - **Note**: This type is referenced but not fully implemented in current codebase
/// - **Use Case**: Simple encrypted data without key exchange metadata
///
/// ### Pattern Recognition Logic
/// The system uses semantic analysis to detect encryption patterns:
/// 1. **PublicKey Detection**: `Value::PublicKey { size_in_bits: 255 }`
/// 2. **Nonce Detection**: `Value::Scalar { size_in_bits: 128, kind: ScalarKind::Unsigned }`
/// 3. **Ciphertext Counting**: Count of `Value::Ciphertext { .. }` values
/// 4. **Type Determination**:
///    - `(PublicKey=true, Nonce=true, Ciphertexts>0)` → SharedEncryptedStruct
///    - `(PublicKey=false, Nonce=true, Ciphertexts>0)` → MXEEncryptedStruct
///    - `(PublicKey=false, Nonce=false, Ciphertexts>0)` → EncDataStruct
///
/// ### Real Example from Testing
///
/// For a circuit `complex_example` returning:
/// ```rust
/// (UserData, Enc<Shared, u32>, (u64, f32), Enc<Mxe, bool>)
/// ```
///
/// **Generated Rust Types:**
/// ```rust
/// // Main output struct - tuple becomes single field_0
/// pub struct ComplexExampleOutput {
///     pub field_0: ComplexExampleOutputStruct0,
/// }
///
/// // The output struct containing all return values
/// pub struct ComplexExampleOutputStruct0 {
///     pub field_0: ComplexExampleOutputStruct00,  // UserData
///     pub field_1: SharedEncryptedStruct<1>,      // Enc<Shared, u32>
///     pub field_2: ComplexExampleOutputStruct02,  // (u64, f32) tuple
///     pub field_3: MXEEncryptedStruct<1>,         // Enc<Mxe, bool>
/// }
///
/// // UserData struct
/// pub struct ComplexExampleOutputStruct00 {
///     pub field_0: u32,   // id
///     pub field_1: bool,  // active
/// }
///
/// // Inner tuple (u64, f32)
/// pub struct ComplexExampleOutputStruct02 {
///     pub field_0: u64,   // timestamp
///     pub field_1: f32,
/// }
/// ```
///
/// **Callback Function Usage:**
/// ```ignore
/// #[arcium_callback(encrypted_ix = "complex_example")]
/// pub fn complex_example_callback(
///     ctx: Context<ComplexExampleCallback>,
///     output: SignedComputationOutputs<ComplexExampleOutput>,
/// ) -> Result<()> {
///     // Use verify_output to verify the BLS signature and deserialize
///     let result = output.verify_output(
///         &ctx.accounts.cluster_account,
///         &ctx.accounts.computation_account
///     )?;
///
///     // Destructure the result
///     let ComplexExampleOutputStruct0 {
///         field_0: ComplexExampleOutputStruct00 { field_0: user_id, field_1: is_active },
///         field_1: shared_encrypted,
///         field_2: ComplexExampleOutputStruct02 { field_0: timestamp, field_1: value },
///         field_3: mxe_encrypted,
///     } = result.field_0;
///
///     // Use the data
///     msg!("User ID: {}, Active: {}", user_id, is_active);
///     msg!("Shared nonce: {}", shared_encrypted.nonce);
///     msg!("Tuple values: {}, {}", timestamp, value);
///     msg!("MXE nonce: {}", mxe_encrypted.nonce);
///
///     Ok(())
/// }
/// ```
pub fn gen_callback_output_struct(conf_ix_name: &str) -> proc_macro2::TokenStream {
    let iface: CircuitInterface = read_conf_ix_interface(conf_ix_name);
    let struct_name = syn::Ident::new(
        &format!("{}Output", iface.name.to_case(Case::Pascal)),
        proc_macro2::Span::call_site(),
    );

    // Generate all custom structs first
    let custom_structs = gen_all_custom_structs(&iface);

    // Generate the main output struct fields
    let fields = iface.outputs.iter().enumerate().map(|(i, val)| {
        let field_name = syn::Ident::new(&format!("field_{}", i), proc_macro2::Span::call_site());
        let ty = value_to_type_for_output(val, &iface.name.to_case(Case::Pascal), i);
        quote::quote! { pub #field_name: #ty }
    });

    let total_size: usize = iface.outputs.iter().map(value_size_in_bytes).sum();

    let x = quote::quote! {
        #(#custom_structs)*

        /// The output of the callback instruction. Provided as a struct with ordered fields
        /// as anchor does not support tuples and tuple structs yet.
        #[derive(AnchorSerialize, AnchorDeserialize)]
        pub struct #struct_name {
            #(#fields),*
        }

        impl #struct_name {
            /// The size of this struct in bytes.
            pub const SIZE: usize = #total_size;
        }

        impl ::arcium_anchor::HasSize for #struct_name {
            const SIZE: usize = #total_size;
        }
    };
    x
}

/// Converts a circuit interface Value to Rust type TokenStreams for code generation.
///
/// # Type Mappings
/// - Encrypted types (`Ciphertext`, `MScalar`, `MFloat`, `MBool`) → `[u8; 32]`
/// - Scalars → Rust integer types (`u8`, `u16`, `u32`, `u64`, `u128`, `i8`, `i16`, `i32`, `i64`,
///   `i128`)
/// - Floats → `f32` or `f64`
/// - Bools → `bool`
/// - PublicKey/Point → `[u8; 32]`
/// - Structs → May be recognized as encryption patterns (`SharedEncryptedStruct<N>`,
///   `MXEEncryptedStruct<N>`) or flattened into component types
/// - Arrays → `[T; N]` where T is the element type
/// - Tuples → `(T1, T2, ...)`
///
/// # Panics
/// Panics on unsupported scalar or float sizes.
fn value_to_type(val: &Value) -> Vec<proc_macro2::TokenStream> {
    match val {
        Value::Ciphertext { .. } | Value::MScalar { .. } | Value::MFloat { .. } | Value::MBool => {
            vec![quote::quote!([u8; 32])]
        }
        Value::Scalar { size_in_bits, kind } => match kind {
            ScalarKind::Unsigned => match size_in_bits {
                8 => vec![quote::quote!(u8)],
                16 => vec![quote::quote!(u16)],
                32 => vec![quote::quote!(u32)],
                64 => vec![quote::quote!(u64)],
                128 => vec![quote::quote!(u128)],
                _ => panic!(
                    "Unsupported unsigned integer size: {} bits. Supported sizes are: 8, 16, 32, 64, 128",
                    size_in_bits
                ),
            },
            ScalarKind::Signed => match size_in_bits {
                8 => vec![quote::quote!(i8)],
                16 => vec![quote::quote!(i16)],
                32 => vec![quote::quote!(i32)],
                64 => vec![quote::quote!(i64)],
                128 => vec![quote::quote!(i128)],
                _ => panic!(
                    "Unsupported signed integer size: {} bits. Supported sizes are: 8, 16, 32, 64, 128",
                    size_in_bits
                ),
            },
        },
        Value::Float { size_in_bits } => match size_in_bits {
            32 => vec![quote::quote!(f32)],
            64 => vec![quote::quote!(f64)],
            _ => panic!(
                "Unsupported float size: {} bits. Supported sizes are: 32 (f32), 64 (f64)",
                size_in_bits
            ),
        },
        Value::Bool => vec![quote::quote!(bool)],
        Value::ArcisX25519Pubkey => vec![quote::quote!([u8; 32])],
        Value::Point => vec![quote::quote!([u8; 32])],
        Value::Array(inner) => {
            let len = inner.len();
            if len == 0 {
                return vec![];
            }
            let tys = value_to_type(&inner[0]);
            if tys.is_empty() {
                vec![]
            } else {
                vec![quote::quote!([#(tys[0]); #len])]
            }
        }
        Value::Tuple(inner) => {
            let tys = inner.iter().flat_map(value_to_type);
            vec![quote::quote!((#(#tys),*))]
        }
        Value::Struct(inner) => {
            // Special case for encryption struct array pattern
            match extract_and_get_encryption_type(&Value::Struct(inner.clone())) {
                Some((EncryptionType::Shared, len)) => vec![quote::quote! {
                    SharedEncryptedStruct<#len>
                }],
                Some((EncryptionType::Mxe, len)) => vec![quote::quote! {
                    MXEEncryptedStruct<#len>
                }],
                Some((EncryptionType::EncData, len)) => vec![quote::quote! {
                    EncDataStruct<#len>
                }],
                None => {
                    // Flatten regular struct fields
                    inner.iter().flat_map(value_to_type).collect()
                }
            }
        }
    }
}

/// Extract encryption components from a Value and return its type if it's an encryption pattern
fn extract_and_get_encryption_type(value: &Value) -> Option<(EncryptionType, usize)> {
    EncryptionComponents::from_value(value).get_type()
}

#[derive(Debug, PartialEq)]
enum EncryptionType {
    Shared,
    Mxe,
    EncData,
}

/// Recursively collects and generates all custom struct definitions from a CircuitInterface
pub fn gen_all_custom_structs(iface: &CircuitInterface) -> Vec<proc_macro2::TokenStream> {
    let mut seen = HashSet::new();
    let mut structs = Vec::new();
    let base_prefix = iface.name.to_case(Case::Pascal);

    for (i, val) in iface.outputs.iter().enumerate() {
        // All outputs use OutputStruct naming convention
        let prefix = format!("{}OutputStruct{}", base_prefix, i);
        collect_structs(val, &mut seen, &mut structs, &prefix);
    }
    structs
}

/// Generate a unique struct name for nested structs.
///
/// # Arguments
/// * `prefix` - The base prefix for the struct name
/// * `seen_count` - Number of structs seen so far (used for uniqueness)
///
/// # Returns
/// The generated struct name
fn generate_nested_struct_name(prefix: &str, seen_count: usize) -> String {
    format!("{}{}", prefix, seen_count)
}

fn collect_structs(
    val: &Value,
    seen: &mut HashSet<String>,
    structs: &mut Vec<proc_macro2::TokenStream>,
    prefix: &str,
) {
    match val {
        Value::Struct(inner) => {
            // Skip special encryption structs
            if extract_and_get_encryption_type(&Value::Struct(inner.clone())).is_some() {
                return;
            }
            // For top-level output structs, use the prefix as-is (e.g., VoteOutputStruct0)
            // For nested structs, use the shared naming logic
            let struct_name = if prefix.contains("OutputStruct") && !seen.contains(prefix) {
                prefix.to_string()
            } else {
                generate_nested_struct_name(prefix, 0)
            };
            if seen.insert(struct_name.clone()) {
                let ident = syn::Ident::new(&struct_name, proc_macro2::Span::call_site());
                let field_types: Vec<_> = inner
                    .iter()
                    .enumerate()
                    .map(|(i, v)| value_to_type_for_structs_with_index(v, &struct_name, seen, i))
                    .collect();
                let fields = field_types.iter().enumerate().map(|(i, ty)| {
                    let field_name =
                        syn::Ident::new(&format!("field_{}", i), proc_macro2::Span::call_site());
                    quote::quote! { pub #field_name: #ty }
                });
                let struct_size: usize = inner.iter().map(value_size_in_bytes).sum();
                structs.push(quote::quote! {
                    #[derive(AnchorSerialize, AnchorDeserialize)]
                    pub struct #ident {
                        #(#fields),*
                    }

                    impl #ident {
                        /// The size of this struct in bytes.
                        pub const SIZE: usize = #struct_size;
                    }

                    impl ::arcium_anchor::HasSize for #ident {
                        const SIZE: usize = #struct_size;
                    }
                });
                // Recurse into inner fields
                for (i, v) in inner.iter().enumerate() {
                    collect_structs(v, seen, structs, &format!("{}{}", struct_name, i));
                }
            }
        }
        Value::Array(inner) => {
            if !inner.is_empty() {
                // Only recurse into arrays if they don't contain tuples
                // Tuples inside arrays remain as tuple types, not converted to structs
                if !matches!(&inner[0], Value::Tuple(_)) {
                    collect_structs(&inner[0], seen, structs, prefix);
                }
            }
        }
        Value::Tuple(inner) => {
            // Since outputs are always single elements or tuples, we use OutputStruct naming
            let struct_name = prefix.to_string();
            if seen.insert(struct_name.clone()) {
                let ident = syn::Ident::new(&struct_name, proc_macro2::Span::call_site());
                let struct_size: usize = inner.iter().map(value_size_in_bytes).sum();

                // Handle empty tuples - they should generate empty structs
                if inner.is_empty() {
                    structs.push(quote::quote! {
                        #[derive(AnchorSerialize, AnchorDeserialize)]
                        pub struct #ident {
                            // Empty struct for empty tuple output
                        }

                        impl #ident {
                            /// The size of this struct in bytes.
                            pub const SIZE: usize = #struct_size;
                        }

                        impl ::arcium_anchor::HasSize for #ident {
                            const SIZE: usize = #struct_size;
                        }
                    });
                } else {
                    let field_types: Vec<_> = inner
                        .iter()
                        .enumerate()
                        .map(|(i, v)| {
                            value_to_type_for_structs_with_index(v, &struct_name, seen, i)
                        })
                        .collect();
                    let fields = field_types.iter().enumerate().map(|(i, ty)| {
                        let field_name = syn::Ident::new(
                            &format!("field_{}", i),
                            proc_macro2::Span::call_site(),
                        );
                        quote::quote! { pub #field_name: #ty }
                    });
                    structs.push(quote::quote! {
                        #[derive(AnchorSerialize, AnchorDeserialize)]
                        pub struct #ident {
                            #(#fields),*
                        }

                        impl #ident {
                            /// The size of this struct in bytes.
                            pub const SIZE: usize = #struct_size;
                        }

                        impl ::arcium_anchor::HasSize for #ident {
                            const SIZE: usize = #struct_size;
                        }
                    });
                    // Recurse into inner fields with field-specific naming for tuples
                    for (field_idx, v) in inner.iter().enumerate() {
                        // For tuples, nested structs should be named based on the original output
                        // index and field index using the shared naming logic
                        let nested_prefix = generate_nested_struct_name(&struct_name, field_idx);
                        collect_structs(v, seen, structs, &nested_prefix);
                    }
                }
            }
        }
        _ => {}
    }
}

// Helper for struct field types
fn value_to_type_for_structs(
    val: &Value,
    prefix: &str,
    seen: &mut HashSet<String>,
) -> proc_macro2::TokenStream {
    value_to_type_for_structs_with_index(val, prefix, seen, 0)
}

fn value_to_type_for_structs_with_index(
    val: &Value,
    prefix: &str,
    seen: &mut HashSet<String>,
    index: usize,
) -> proc_macro2::TokenStream {
    match val {
        Value::Struct(inner) => {
            match extract_and_get_encryption_type(&Value::Struct(inner.clone())) {
                Some((EncryptionType::Shared, len)) => {
                    quote::quote! { SharedEncryptedStruct<#len> }
                }
                Some((EncryptionType::Mxe, len)) => quote::quote! { MXEEncryptedStruct<#len> },
                Some((EncryptionType::EncData, len)) => quote::quote! { EncDataStruct<#len> },
                None => {
                    // Use shared naming logic to ensure consistency with collect_structs
                    let struct_name = format!("{}{}", prefix, index);
                    let ident = syn::Ident::new(&struct_name, proc_macro2::Span::call_site());
                    quote::quote! { #ident }
                }
            }
        }
        Value::Array(inner) => {
            if inner.is_empty() {
                quote::quote!([(); 0])
            } else {
                let ty = value_to_type_for_structs(&inner[0], prefix, seen);
                let len = inner.len();
                quote::quote!([#ty; #len])
            }
        }
        Value::Tuple(_) => {
            panic!("Tuple in circuit return type is not supported.");
        }
        _ => value_to_type(val)
            .into_iter()
            .next()
            .unwrap_or_else(|| quote::quote!(())),
    }
}

// Helper for output struct field types
fn value_to_type_for_output(val: &Value, prefix: &str, i: usize) -> proc_macro2::TokenStream {
    match val {
        Value::Struct(inner) => {
            // Special case for encryption struct patterns
            match extract_and_get_encryption_type(&Value::Struct(inner.clone())) {
                Some((EncryptionType::Shared, len)) => quote::quote! {
                    SharedEncryptedStruct<#len>
                },
                Some((EncryptionType::Mxe, len)) => quote::quote! {
                    MXEEncryptedStruct<#len>
                },
                Some((EncryptionType::EncData, len)) => quote::quote! {
                    EncDataStruct<#len>
                },
                None => {
                    // Use the custom struct name
                    let struct_name = format!("{}OutputStruct{}", prefix, i);
                    let ident = syn::Ident::new(&struct_name, proc_macro2::Span::call_site());
                    quote::quote! { #ident }
                }
            }
        }
        Value::Array(inner) => {
            let len = inner.len();
            if len == 0 {
                quote::quote!([(); 0])
            } else {
                let ty = value_to_type_for_output(&inner[0], prefix, i);
                quote::quote!([#ty; #len])
            }
        }
        // Convert tuple to struct by creating a custom struct name
        // Note: `collect_structs` will generate the actual struct definition
        Value::Tuple(_inner) => {
            let struct_name = format!("{}OutputStruct{}", prefix, i);
            let ident = syn::Ident::new(&struct_name, proc_macro2::Span::call_site());
            quote::quote! { #ident }
        }
        _ => value_to_type(val)
            .into_iter()
            .next()
            .unwrap_or_else(|| quote::quote!(())),
    }
}

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

    #[test]
    fn test_semantic_encryption_detection() {
        // Test shared encryption
        let shared_v = Value::Struct(vec![
            Value::Struct(vec![
                Value::ArcisX25519Pubkey,
                Value::Scalar {
                    size_in_bits: 128,
                    kind: ScalarKind::Unsigned,
                },
            ]),
            Value::Struct(vec![
                Value::Array(vec![Value::Ciphertext { size_in_bits: 255 }]),
                Value::Array(vec![]),
            ]),
        ]);
        assert_eq!(
            extract_and_get_encryption_type(&shared_v),
            Some((EncryptionType::Shared, 1))
        );

        // Test MXE encryption
        let mxe_v = Value::Struct(vec![
            Value::Struct(vec![Value::Scalar {
                size_in_bits: 128,
                kind: ScalarKind::Unsigned,
            }]),
            Value::Struct(vec![
                Value::Array(vec![
                    Value::Ciphertext { size_in_bits: 255 },
                    Value::Ciphertext { size_in_bits: 255 },
                ]),
                Value::Array(vec![]),
            ]),
        ]);
        assert_eq!(
            extract_and_get_encryption_type(&mxe_v),
            Some((EncryptionType::Mxe, 2))
        );

        // Test EncData (only ciphertext, no public key or nonce)
        let enc_data_v = Value::Struct(vec![Value::Struct(vec![
            Value::Array(vec![Value::Ciphertext { size_in_bits: 255 }]),
            Value::Array(vec![]),
        ])]);
        assert_eq!(
            extract_and_get_encryption_type(&enc_data_v),
            Some((EncryptionType::EncData, 1))
        );

        // Test non-encryption struct
        let normal_v = Value::Struct(vec![
            Value::Scalar {
                size_in_bits: 32,
                kind: ScalarKind::Unsigned,
            },
            Value::Bool,
        ]);
        assert_eq!(extract_and_get_encryption_type(&normal_v), None);
    }

    #[test]
    fn test_gen_all_custom_structs() {
        let iface = CircuitInterface {
            name: "TestInterface".to_string(),
            inputs: vec![],
            outputs: vec![
                Value::Struct(vec![
                    Value::Scalar {
                        size_in_bits: 32,
                        kind: ScalarKind::Unsigned,
                    },
                    Value::Bool,
                ]),
                Value::Tuple(vec![
                    Value::Scalar {
                        size_in_bits: 64,
                        kind: ScalarKind::Unsigned,
                    },
                    Value::Float { size_in_bits: 32 },
                ]),
            ],
        };

        let structs = gen_all_custom_structs(&iface);
        assert_eq!(structs.len(), 2); // One for struct, one for tuple
    }

    #[test]
    fn test_empty_tuple_handling() {
        // Test the manticore_auc case - empty tuple output
        let iface = CircuitInterface {
            name: "ManticoreAuc".to_string(),
            inputs: vec![],
            outputs: vec![Value::Tuple(vec![])], // Empty tuple like manticore_auc
        };

        let structs = gen_all_custom_structs(&iface);
        assert_eq!(structs.len(), 1); // Should generate one empty struct

        // Check that the generated struct has the correct name and is empty
        let struct_code = structs[0].to_string();
        assert!(struct_code.contains("pub struct ManticoreAucOutputStruct0"));
        assert!(struct_code.contains("# [derive (AnchorSerialize , AnchorDeserialize)]"));
        // Should not contain any field definitions
        assert!(!struct_code.contains("pub field_"));

        // Test value_to_type_for_output also handles empty tuples correctly
        let output_type = value_to_type_for_output(&Value::Tuple(vec![]), "ManticoreAuc", 0);
        assert_eq!(output_type.to_string(), "ManticoreAucOutputStruct0");
    }

    #[test]
    fn test_naming_consistency_for_nested_structs() {
        // Test the dark-pool insert_order case - tuple with nested struct
        let iface = CircuitInterface {
            name: "InsertOrder".to_string(),
            inputs: vec![],
            outputs: vec![Value::Tuple(vec![
                // field_0: MXEEncryptedStruct (should be recognized as encryption type)
                Value::Struct(vec![
                    Value::Struct(vec![Value::Scalar {
                        size_in_bits: 128,
                        kind: ScalarKind::Unsigned,
                    }]),
                    Value::Struct(vec![
                        Value::Array(vec![
                            Value::Ciphertext { size_in_bits: 255 },
                            Value::Ciphertext { size_in_bits: 255 },
                        ]),
                        Value::Array(vec![]),
                    ]),
                ]),
                // field_1: Regular struct with 6 u64 fields (should generate custom struct)
                Value::Struct(vec![
                    Value::Scalar {
                        size_in_bits: 64,
                        kind: ScalarKind::Unsigned,
                    }, // u64
                    Value::Scalar {
                        size_in_bits: 64,
                        kind: ScalarKind::Unsigned,
                    }, // u64
                    Value::Scalar {
                        size_in_bits: 64,
                        kind: ScalarKind::Unsigned,
                    }, // u64
                    Value::Scalar {
                        size_in_bits: 64,
                        kind: ScalarKind::Unsigned,
                    }, // u64
                    Value::Scalar {
                        size_in_bits: 64,
                        kind: ScalarKind::Unsigned,
                    }, // u64
                    Value::Scalar {
                        size_in_bits: 64,
                        kind: ScalarKind::Unsigned,
                    }, // u64
                ]),
            ])],
        };

        let structs = gen_all_custom_structs(&iface);

        // Should generate 2 structs:
        // 1. InsertOrderOutputStruct0 (the main output struct)
        // 2. InsertOrderOutputStruct01 (the nested struct for field_1)
        assert_eq!(structs.len(), 2);

        let struct_strings: Vec<String> = structs.iter().map(|s| s.to_string()).collect();

        // Verify main output struct exists
        assert!(struct_strings
            .iter()
            .any(|s| s.contains("pub struct InsertOrderOutputStruct0")));

        // Verify nested struct has correct name
        assert!(struct_strings
            .iter()
            .any(|s| s.contains("pub struct InsertOrderOutputStruct01")));

        // Verify that the main struct references the nested struct with the correct name
        let main_struct = struct_strings
            .iter()
            .find(|s| s.contains("pub struct InsertOrderOutputStruct0"))
            .expect("Main output struct should exist");

        // The field_1 should reference InsertOrderOutputStruct01
        assert!(main_struct.contains("InsertOrderOutputStruct01"));
    }

    #[test]
    fn test_value_to_type_for_output() {
        // Test regular struct
        let val = Value::Struct(vec![
            Value::Scalar {
                size_in_bits: 32,
                kind: ScalarKind::Unsigned,
            },
            Value::Bool,
        ]);
        let ty = value_to_type_for_output(&val, "Test", 0);
        assert!(!ty.to_string().is_empty());

        // Test tuple
        let val = Value::Tuple(vec![
            Value::Scalar {
                size_in_bits: 64,
                kind: ScalarKind::Unsigned,
            },
            Value::Float { size_in_bits: 32 },
        ]);
        let ty = value_to_type_for_output(&val, "Test", 0);
        assert!(!ty.to_string().is_empty());

        // Test special encryption struct
        let val = Value::Struct(vec![
            Value::Struct(vec![
                Value::ArcisX25519Pubkey,
                Value::Scalar {
                    size_in_bits: 128,
                    kind: ScalarKind::Unsigned,
                },
            ]),
            Value::Struct(vec![
                Value::Array(vec![Value::Ciphertext { size_in_bits: 255 }]),
                Value::Array(vec![]),
            ]),
        ]);
        let ty = value_to_type_for_output(&val, "Test", 0);
        assert!(ty.to_string().contains("SharedEncryptedStruct"));

        // Test EncData struct
        let val = Value::Struct(vec![Value::Struct(vec![
            Value::Array(vec![Value::Ciphertext { size_in_bits: 255 }]),
            Value::Array(vec![]),
        ])]);
        let ty = value_to_type_for_output(&val, "Test", 0);
        assert!(ty.to_string().contains("EncDataStruct"));
    }

    #[test]
    fn test_generate_nested_struct_name() {
        // Test regular prefix
        assert_eq!(
            generate_nested_struct_name("ComplexExample", 0),
            "ComplexExample0"
        );
        assert_eq!(
            generate_nested_struct_name("ComplexExample", 5),
            "ComplexExample5"
        );

        // Test OutputStruct prefix
        assert_eq!(
            generate_nested_struct_name("InsertOrderOutputStruct0", 1),
            "InsertOrderOutputStruct01"
        );
        assert_eq!(
            generate_nested_struct_name("ComplexExampleOutputStruct0", 2),
            "ComplexExampleOutputStruct02"
        );
    }

    #[test]
    fn test_struct_generation_with_nested_types() {
        let iface = CircuitInterface {
            name: "NestedTest".to_string(),
            inputs: vec![],
            outputs: vec![Value::Struct(vec![
                Value::Struct(vec![
                    Value::Scalar {
                        size_in_bits: 32,
                        kind: ScalarKind::Unsigned,
                    },
                    Value::Bool,
                ]),
                Value::Array(vec![
                    Value::Scalar {
                        size_in_bits: 64,
                        kind: ScalarKind::Unsigned,
                    },
                    Value::Float { size_in_bits: 32 },
                ]),
            ])],
        };

        let structs = gen_all_custom_structs(&iface);
        // Should generate exactly 2 structs:
        // 1. NestedTestOutputStruct0 - the outer struct
        // 2. NestedTestOutputStruct00 - the inner struct (first field of the outer struct)
        // Note: Tuples inside arrays remain as tuple types, not converted to structs
        assert_eq!(structs.len(), 2);
    }

    #[test]
    fn test_value_size_in_bytes_primitives() {
        assert_eq!(value_size_in_bytes(&Value::Bool), 1);
        assert_eq!(
            value_size_in_bytes(&Value::Scalar {
                size_in_bits: 32,
                kind: ScalarKind::Unsigned
            }),
            4
        );
        assert_eq!(
            value_size_in_bytes(&Value::Scalar {
                size_in_bits: 64,
                kind: ScalarKind::Unsigned
            }),
            8
        );
        assert_eq!(value_size_in_bytes(&Value::Float { size_in_bits: 32 }), 4);
        assert_eq!(
            value_size_in_bytes(&Value::Ciphertext { size_in_bits: 255 }),
            32
        );
    }

    #[test]
    fn test_value_size_in_bytes_composites() {
        // Tuple: u64 + bool + f32 = 8 + 1 + 4 = 13
        assert_eq!(
            value_size_in_bytes(&Value::Tuple(vec![
                Value::Scalar {
                    size_in_bits: 64,
                    kind: ScalarKind::Unsigned
                },
                Value::Bool,
                Value::Float { size_in_bits: 32 },
            ])),
            13
        );

        // Array of 5 u32s = 5 * 4 = 20
        assert_eq!(
            value_size_in_bytes(&Value::Array(vec![
                Value::Scalar {
                    size_in_bits: 32,
                    kind: ScalarKind::Unsigned
                };
                5
            ])),
            20
        );

        // Struct: u32 + bool = 4 + 1 = 5
        assert_eq!(
            value_size_in_bytes(&Value::Struct(vec![
                Value::Scalar {
                    size_in_bits: 32,
                    kind: ScalarKind::Unsigned
                },
                Value::Bool,
            ])),
            5
        );
    }

    #[test]
    fn test_value_size_in_bytes_encryption_types() {
        // SharedEncryptedStruct<1>: 32 (pubkey) + 16 (nonce) + 32 (ciphertext) = 80
        let shared_enc = Value::Struct(vec![
            Value::Struct(vec![
                Value::ArcisX25519Pubkey,
                Value::Scalar {
                    size_in_bits: 128,
                    kind: ScalarKind::Unsigned,
                },
            ]),
            Value::Struct(vec![
                Value::Array(vec![Value::Ciphertext { size_in_bits: 255 }]),
                Value::Array(vec![]),
            ]),
        ]);
        assert_eq!(value_size_in_bytes(&shared_enc), 80);

        // MXEEncryptedStruct<2>: 16 (nonce) + 64 (2 ciphertexts) = 80
        let mxe_enc = Value::Struct(vec![
            Value::Struct(vec![Value::Scalar {
                size_in_bits: 128,
                kind: ScalarKind::Unsigned,
            }]),
            Value::Struct(vec![
                Value::Array(vec![
                    Value::Ciphertext { size_in_bits: 255 },
                    Value::Ciphertext { size_in_bits: 255 },
                ]),
                Value::Array(vec![]),
            ]),
        ]);
        assert_eq!(value_size_in_bytes(&mxe_enc), 80);
    }

    #[test]
    fn test_generated_struct_has_correct_size_constant() {
        let iface = CircuitInterface {
            name: "SizeTest".to_string(),
            inputs: vec![],
            outputs: vec![Value::Tuple(vec![
                Value::Scalar {
                    size_in_bits: 64,
                    kind: ScalarKind::Unsigned,
                },
                Value::Bool,
            ])],
        };

        let structs = gen_all_custom_structs(&iface);
        let struct_code = structs[0].to_string();

        // Should have SIZE = 9 (u64=8 + bool=1)
        assert!(struct_code.contains("const SIZE : usize = 9usize"));
    }
}