arcium-macros 0.2.0

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

// Semantic encryption components extracted from any structure
#[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,
        }
    }

    fn extract_from(&mut self, value: &Value) {
        match value {
            Value::PublicKey { .. } => {
                self.has_public_key = true;
            }
            Value::Scalar { size_in_bits: 128 } => {
                self.has_nonce = true;
            }
            Value::Ciphertext { .. } => {
                self.ciphertext_count += 1;
            }
            Value::Array(values) if !values.is_empty() => {
                // Skip empty arrays (phantom data)
                for v in values {
                    self.extract_from(v);
                }
            }
            Value::Struct(values) => {
                for v in values {
                    self.extract_from(v);
                }
            }
            Value::Tuple(values) => {
                for v in values {
                    self.extract_from(v);
                }
            }
            // Other primitive types don't contribute to encryption semantics
            _ => {}
        }
    }
}

/// Generates a struct definition for the outputs of a given circuit interface
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 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),*
        }
    };
    x
}

/// Maps Value to a Vec of Rust types as TokenStreams, mapping encrypted types to [u8; 32]
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 } => 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 scalar size: {}", 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: {}", size_in_bits),
        },
        Value::Bool => vec![quote::quote!(bool)],
        Value::PublicKey { .. } => 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() {
        // Pass the full expected struct name as the prefix for top-level output structs
        let prefix = format!("{}OutputStruct{}", base_prefix, i);
        collect_structs(val, &mut seen, &mut structs, &prefix);
    }
    structs
}

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, append a sequential number
            let struct_name = if prefix.contains("OutputStruct") {
                prefix.to_string()
            } else {
                format!("{}{}", prefix, seen.len())
            };
            if seen.insert(struct_name.clone()) {
                let ident = syn::Ident::new(&struct_name, proc_macro2::Span::call_site());
                let field_types: Vec<_> = inner
                    .iter()
                    .map(|v| value_to_type_for_structs(v, prefix, seen))
                    .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),*
                    }
                });
                // Recurse into inner fields
                for v in inner {
                    collect_structs(v, seen, structs, &struct_name);
                }
            }
        }
        Value::Array(inner) => {
            if !inner.is_empty() {
                collect_structs(&inner[0], seen, structs, prefix);
            }
        }
        Value::Tuple(inner) => {
            // Generate a unique name for tuple struct
            // Handle both top-level and nested tuple structs
            let struct_name = if prefix.contains("OutputStruct") {
                // For top-level output that's a tuple, replace OutputStruct with TupleStruct
                prefix.replace("OutputStruct", "TupleStruct")
            } else {
                format!("{}TupleStruct{}", prefix, seen.len())
            };
            if seen.insert(struct_name.clone()) {
                let ident = syn::Ident::new(&struct_name, proc_macro2::Span::call_site());
                let field_types: Vec<_> = inner
                    .iter()
                    .map(|v| value_to_type_for_structs(v, prefix, seen))
                    .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),*
                    }
                });
                // Recurse into inner fields
                for v in inner {
                    collect_structs(v, seen, structs, &struct_name);
                }
            }
        }
        _ => {}
    }
}

// Helper for struct field types
fn value_to_type_for_structs(
    val: &Value,
    prefix: &str,
    seen: &mut HashSet<String>,
) -> 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 => {
                    let struct_name = format!("{}{}", prefix, seen.len());
                    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(inner) => {
            let tys = inner
                .iter()
                .map(|v| value_to_type_for_structs(v, prefix, seen));
            quote::quote!((#(#tys),*))
        }
        _ => 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!("{}TupleStruct{}", 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::PublicKey { size_in_bits: 255 },
                Value::Scalar { size_in_bits: 128 },
            ]),
            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 }]),
            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 }, 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 }, Value::Bool]),
                Value::Tuple(vec![
                    Value::Scalar { size_in_bits: 64 },
                    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_value_to_type_for_output() {
        // Test regular struct
        let val = Value::Struct(vec![Value::Scalar { size_in_bits: 32 }, 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 },
            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::PublicKey { size_in_bits: 255 },
                Value::Scalar { size_in_bits: 128 },
            ]),
            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_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 }, Value::Bool]),
                Value::Array(vec![Value::Tuple(vec![
                    Value::Scalar { size_in_bits: 64 },
                    Value::Float { size_in_bits: 32 },
                ])]),
            ])],
        };

        let structs = gen_all_custom_structs(&iface);
        // Should generate exactly 2 structs:
        // 1. NestedTestOutputStruct0 - the outer struct
        // 2. 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);
    }
}