fmi-export-derive 0.3.0

use attribute_derive::FromAttr;
use proc_macro_error2::emit_error;

mod field_attr;
pub use field_attr::{
    AliasAttribute, ChildAttribute, FieldAttribute, FieldAttributeOuter, TerminalAttribute,
};

/// Helper function to extract docstring from a syn::Attribute
/// Follows DRY principles by centralizing doc attribute parsing logic
fn parse_doc_attribute(attr: &syn::Attribute) -> Option<String> {
    if attr.meta.path().is_ident("doc") {
        attr.meta.require_name_value().ok().and_then(|name_value| {
            if let syn::Expr::Lit(syn::ExprLit {
                lit: syn::Lit::Str(lit_str),
                ..
            }) = &name_value.value
            {
                let doc_line = lit_str.value();
                Some(doc_line.strip_prefix(' ').unwrap_or(&doc_line).to_string())
            } else {
                None
            }
        })
    } else {
        None
    }
}

/// StructAttribute represents the attributes that can be applied to the model struct
#[derive(Debug, attribute_derive::FromAttr, PartialEq, Clone, Default)]
#[attribute(ident = model)]
pub struct StructAttr {
    /// Optional model description (defaults to the struct docstring)
    #[attribute(optional)]
    pub description: Option<String>,

    /// Enable Model Exchange interface
    #[attribute(optional)]
    pub model_exchange: Option<bool>,

    /// Enable Co-Simulation interface
    #[attribute(optional)]
    pub co_simulation: Option<bool>,

    /// Enable Scheduled Execution interface
    #[attribute(optional)]
    pub scheduled_execution: Option<bool>,

    /// Whether to auto-generate a UserModel impl (default: true)
    #[attribute(optional)]
    pub user_model: Option<bool>,
}

#[derive(Debug, PartialEq, Clone)]
pub enum StructAttrOuter {
    Docstring(String),
    Model(StructAttr),
    Terminal(TerminalStructAttr),
}

/// Struct-level attributes for terminal generation
#[derive(Debug, attribute_derive::FromAttr, PartialEq, Clone, Default)]
#[attribute(ident = terminal)]
pub struct TerminalStructAttr {
    /// Optional terminal name override (defaults to struct name)
    #[attribute(optional)]
    pub name: Option<String>,
    /// Optional matching rule override (defaults to "bus")
    #[attribute(optional)]
    pub matching_rule: Option<String>,
    /// Optional terminal kind override
    #[attribute(optional)]
    pub terminal_kind: Option<String>,
}

/// Representation of an FmuModel field with it's parsed attributes
#[derive(Debug, PartialEq, Clone)]
pub struct Field {
    pub ident: syn::Ident,
    pub rust_type: syn::Type,
    pub attrs: Vec<FieldAttributeOuter>,
}

#[derive(Debug, PartialEq, Clone)]
pub struct Model {
    pub ident: syn::Ident,
    pub attrs: Vec<StructAttrOuter>,
    pub fields: Vec<Field>,
}

impl Field {
    /// Extract the description from field docstrings
    pub fn fold_description(&self) -> String {
        self.attrs
            .iter()
            .find_map(|attr| {
                if let FieldAttributeOuter::Docstring(doc) = attr {
                    Some(doc.clone())
                } else {
                    None
                }
            })
            .unwrap_or_else(|| "".to_string())
    }
}

impl Model {
    /// Returns an iterator over all continuous state fields.
    /// A field is considered a continuous state if another field has it as a derivative.
    pub fn iter_continuous_states(&self) -> impl Iterator<Item = &Field> {
        self.fields.iter().filter(move |field| {
            let field_name = field.ident.to_string();
            self.is_continuous_state(&field_name)
        })
    }

    /// Checks if a field with the given name is a continuous state variable.
    /// A field is a continuous state if another field has it as a derivative.
    pub fn is_continuous_state(&self, field_name: &str) -> bool {
        self.fields.iter().any(|other_field| {
            other_field.attrs.iter().any(|attr| {
                let derivative_ref = match attr {
                    FieldAttributeOuter::Variable(var_attr) => &var_attr.derivative,
                    _ => return false,
                };

                derivative_ref.as_ref().map(|d| d.to_string()) == Some(field_name.to_string())
            })
        })
    }

    /// Iterator over all fields that are derivatives.
    /// A field is considered a derivative if it has a derivative attribute.
    pub fn iter_derivatives(&self) -> impl Iterator<Item = &Field> {
        self.fields.iter().filter(|field| self.is_derivative(field))
    }

    /// Checks if a field is a derivative variable.
    /// A field is a derivative if it has a derivative attribute.
    pub fn is_derivative(&self, field: &Field) -> bool {
        field.attrs.iter().any(|attr| match attr {
            FieldAttributeOuter::Variable(var_attr) => var_attr.derivative.is_some(),
            _ => false,
        })
    }

    /// Get the parsed model attribute, if present
    pub fn get_model_attr(&self) -> Option<&StructAttr> {
        self.attrs.iter().find_map(|attr| match attr {
            StructAttrOuter::Model(model_attr) => Some(model_attr),
            _ => None,
        })
    }

    /// Get the parsed terminal attribute, if present
    pub fn get_terminal_attr(&self) -> Option<&TerminalStructAttr> {
        self.attrs.iter().find_map(|attr| match attr {
            StructAttrOuter::Terminal(terminal_attr) => Some(terminal_attr),
            _ => None,
        })
    }

    /// Check if Model Exchange is supported
    pub fn supports_model_exchange(&self) -> bool {
        self.get_model_attr()
            .and_then(|attr| attr.model_exchange)
            .unwrap_or(true) // Default to true if not specified
    }

    /// Check if Co-Simulation is supported
    pub fn supports_co_simulation(&self) -> bool {
        self.get_model_attr()
            .and_then(|attr| attr.co_simulation)
            .unwrap_or(false)
    }

    /// Check if Scheduled Execution is supported
    pub fn supports_scheduled_execution(&self) -> bool {
        self.get_model_attr()
            .and_then(|attr| attr.scheduled_execution)
            .unwrap_or(false)
    }

    /// Check if UserModel impl should be auto-generated
    pub fn auto_user_model(&self) -> bool {
        self.get_model_attr()
            .and_then(|attr| attr.user_model)
            .unwrap_or(true)
    }
}

impl TryFrom<syn::Field> for Field {
    type Error = String;
    fn try_from(field: syn::Field) -> Result<Self, String> {
        use attribute_derive::Attribute;
        let attrs = field
            .attrs
            .iter()
            .filter_map(|attr| match attr.meta.path().get_ident() {
                Some(ident) if ident == "doc" => {
                    parse_doc_attribute(attr).map(FieldAttributeOuter::Docstring)
                }

                Some(ident) if ident == "variable" => {
                    match FieldAttribute::from_attribute(attr).map(FieldAttributeOuter::Variable) {
                        Ok(attr) => Some(attr),
                        Err(e) => {
                            emit_error!(attr, format!("{e}"));
                            None
                        }
                    }
                }

                Some(ident) if ident == "alias" => {
                    match AliasAttribute::from_attribute(attr).map(FieldAttributeOuter::Alias) {
                        Ok(attr) => Some(attr),
                        Err(e) => {
                            emit_error!(attr, format!("{e}"));
                            None
                        }
                    }
                }

                Some(ident) if ident == "child" => {
                    match ChildAttribute::from_attribute(attr).map(FieldAttributeOuter::Child) {
                        Ok(attr) => Some(attr),
                        Err(e) => {
                            emit_error!(attr, format!("{e}"));
                            None
                        }
                    }
                }

                Some(ident) if ident == "terminal" => {
                    match TerminalAttribute::from_attribute(attr).map(FieldAttributeOuter::Terminal)
                    {
                        Ok(attr) => Some(attr),
                        Err(e) => {
                            emit_error!(attr, format!("{e}"));
                            None
                        }
                    }
                }

                _ => None,
            })
            .collect();

        Ok(Self {
            ident: field.ident.expect("Expected named field"),
            rust_type: field.ty,
            attrs,
        })
    }
}

/// Check for variable name conflicts with the built-in "time" variable
/// and for alias/variable name uniqueness.
fn check_time_variable_conflicts(fields: &[Field]) {
    let mut seen_names = std::collections::HashSet::new();

    for field in fields {
        let field_name = field.ident.to_string();

        // Check variable names (including custom name overrides)
        for attr in &field.attrs {
            if let FieldAttributeOuter::Variable(var_attr) = attr {
                let var_name = var_attr.name.as_deref().unwrap_or(&field_name);
                if var_name.to_lowercase() == "time" {
                    emit_error!(field.ident, "'time' is a reserved name.");
                }
                if !seen_names.insert(var_name.to_string()) {
                    emit_error!(
                        field.ident,
                        format!("Duplicate variable or alias name '{var_name}'.")
                    );
                }
            }
        }

        // Check alias names too
        for attr in &field.attrs {
            if let FieldAttributeOuter::Alias(alias_attr) = attr {
                let alias_name = alias_attr.name.as_str();
                if alias_name.to_lowercase() == "time" {
                    emit_error!(field.ident, "'time' is a reserved name.");
                }
                if !seen_names.insert(alias_name.to_string()) {
                    emit_error!(
                        field.ident,
                        format!("Duplicate variable or alias name '{alias_name}'.")
                    );
                }
            }
        }
    }
}

impl From<syn::DeriveInput> for Model {
    fn from(item: syn::DeriveInput) -> Self {
        if let syn::Data::Struct(struct_data) = item.data {
            let attrs = build_attrs(item.attrs);
            let fields = build_fields(struct_data.fields);

            // Check for time variable name conflicts
            check_time_variable_conflicts(&fields);

            Self {
                ident: item.ident,
                attrs,
                fields,
            }
        } else {
            emit_error!(item, "FmuModel can only be derived for structs");
            Self {
                ident: item.ident,
                attrs: vec![],
                fields: vec![],
            }
        }
    }
}

/// Parse struct-level attributes, accepting explicit boolean values
pub fn build_attrs(attrs: Vec<syn::Attribute>) -> Vec<StructAttrOuter> {
    attrs
        .into_iter()
        .filter_map(|attr| match attr.meta.path().get_ident() {
            Some(ident) if ident == "doc" => {
                parse_doc_attribute(&attr).map(StructAttrOuter::Docstring)
            }

            Some(ident) if ident == "model" => {
                // Prefer the derived parser; fall back to a manual tolerant parser for bool flags
                match StructAttr::from_attribute(attr.clone())
                    .or_else(|_e| parse_model_attr_bool(attr.clone()))
                {
                    Ok(attr) => Some(StructAttrOuter::Model(attr)),
                    Err(e) => {
                        emit_error!(attr, format!("{e}"));
                        None
                    }
                }
            }

            Some(ident) if ident == "terminal" => {
                match TerminalStructAttr::from_attribute(attr.clone()) {
                    Ok(attr) => Some(StructAttrOuter::Terminal(attr)),
                    Err(e) => {
                        emit_error!(attr, format!("{e}"));
                        None
                    }
                }
            }

            _ => None,
        })
        .collect()
}

/// Fallback parser that accepts boolean flags for co_simulation/model_exchange/scheduled_execution
fn parse_model_attr_bool(attr: syn::Attribute) -> Result<StructAttr, String> {
    let mut model_attr = StructAttr::default();

    let list = attr
        .meta
        .require_list()
        .map_err(|_| "expected a model attribute list like #[model(...)]".to_string())?;

    for nested in list
        .parse_args_with(syn::punctuated::Punctuated::<syn::Meta, syn::Token![,]>::parse_terminated)
        .map_err(|_| "failed to parse #[model(...)] arguments".to_string())?
    {
        match nested {
            syn::Meta::NameValue(nv) if nv.path.is_ident("model_exchange") => {
                if let syn::Expr::Lit(syn::ExprLit {
                    lit: syn::Lit::Bool(lit_bool),
                    ..
                }) = nv.value
                {
                    model_attr.model_exchange = Some(lit_bool.value);
                } else {
                    return Err("model_exchange expects a boolean".into());
                }
            }
            syn::Meta::Path(path) if path.is_ident("model_exchange") => {
                return Err(
                    "model_exchange expects a boolean value, e.g. model_exchange = true".into(),
                );
            }
            syn::Meta::List(list) if list.path.is_ident("model_exchange") => {
                let _ = list;
                return Err(
                    "model_exchange expects a boolean value, e.g. model_exchange = true".into(),
                );
            }
            syn::Meta::NameValue(nv) if nv.path.is_ident("co_simulation") => {
                if let syn::Expr::Lit(syn::ExprLit {
                    lit: syn::Lit::Bool(lit_bool),
                    ..
                }) = nv.value
                {
                    model_attr.co_simulation = Some(lit_bool.value);
                } else {
                    return Err("co_simulation expects a boolean".into());
                }
            }
            syn::Meta::Path(path) if path.is_ident("co_simulation") => {
                return Err(
                    "co_simulation expects a boolean value, e.g. co_simulation = true".into(),
                );
            }
            syn::Meta::List(list) if list.path.is_ident("co_simulation") => {
                let _ = list;
                return Err(
                    "co_simulation expects a boolean value, e.g. co_simulation = true".into(),
                );
            }
            syn::Meta::NameValue(nv) if nv.path.is_ident("scheduled_execution") => {
                if let syn::Expr::Lit(syn::ExprLit {
                    lit: syn::Lit::Bool(lit_bool),
                    ..
                }) = nv.value
                {
                    model_attr.scheduled_execution = Some(lit_bool.value);
                } else {
                    return Err("scheduled_execution expects a boolean".into());
                }
            }
            syn::Meta::Path(path) if path.is_ident("scheduled_execution") => {
                return Err(
                    "scheduled_execution expects a boolean value, e.g. scheduled_execution = true"
                        .into(),
                );
            }
            syn::Meta::List(list) if list.path.is_ident("scheduled_execution") => {
                let _ = list;
                return Err(
                    "scheduled_execution expects a boolean value, e.g. scheduled_execution = true"
                        .into(),
                );
            }
            syn::Meta::NameValue(nv) if nv.path.is_ident("user_model") => {
                if let syn::Expr::Lit(syn::ExprLit {
                    lit: syn::Lit::Bool(lit_bool),
                    ..
                }) = nv.value
                {
                    model_attr.user_model = Some(lit_bool.value);
                } else {
                    return Err("user_model expects a boolean".into());
                }
            }
            syn::Meta::Path(path) if path.is_ident("user_model") => {
                return Err("user_model expects a boolean value, e.g. user_model = false".into());
            }
            syn::Meta::List(list) if list.path.is_ident("user_model") => {
                let _ = list;
                return Err("user_model expects a boolean value, e.g. user_model = false".into());
            }
            // Ignore unknown entries here and let the main parser report errors elsewhere
            _ => {}
        }
    }

    Ok(model_attr)
}

/// Check if a field has any FMU-relevant attributes (variable or alias)
fn has_fmu_attributes(field: &syn::Field) -> bool {
    field.attrs.iter().any(|attr| {
        attr.meta
            .path()
            .get_ident()
            .map(|ident| {
                ident == "variable" || ident == "alias" || ident == "child" || ident == "terminal"
            })
            .unwrap_or(false)
    })
}

pub fn build_fields(fields: syn::Fields) -> Vec<Field> {
    match fields {
        syn::Fields::Named(syn::FieldsNamed { named, .. }) => named
            .into_iter()
            .filter(has_fmu_attributes) // Only process fields with FMU attributes
            .filter_map(|ref field| match Field::try_from(field.clone()) {
                Ok(field) => Some(field),
                Err(e) => {
                    emit_error!(field, format!("{e}"));
                    None
                }
            })
            .collect(),
        _ => {
            emit_error!(fields, "Expected named fields in the struct");
            vec![]
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use attribute_derive::FromAttr;
    use fmi::fmi3::schema;

    #[test]
    fn test_attribute() {
        let input: syn::Attribute = syn::parse_quote! {
            #[variable(causality = Parameter, variability = Fixed, start = -9.81)]
        };
        let _attr = FieldAttribute::from_attribute(input).unwrap();

        let input: syn::Attribute = syn::parse_quote! {
            #[variable(causality = Output, start = 0.0)]
        };
        let _attr = FieldAttribute::from_attribute(input).unwrap();
    }

    #[test]
    fn test_fields_and_attributes() {
        let input: syn::ItemStruct = syn::parse_quote! {
            struct TestModel {
                /// Test1
                #[variable(causality = Output, start = 1.0)]
                h: f64,

                /// Test2
                #[variable(causality = Output, start = 0.0)]
                #[alias(name="alias_h", description = "Alias of h")]
                v: f64,
            }
        };
        let fields = build_fields(input.fields);

        assert_eq!(fields.len(), 2, "There should be 2 fields");
        assert_eq!(
            fields[0].attrs,
            vec![
                FieldAttributeOuter::Docstring("Test1".to_string()),
                FieldAttributeOuter::Variable(FieldAttribute {
                    causality: Some(schema::Causality::Output.into()),
                    start: Some(syn::parse_quote!(1.0)),
                    ..Default::default()
                })
            ],
            "First field should have 2 attributes: docstring and variable"
        );
        assert_eq!(
            fields[1].attrs,
            vec![
                FieldAttributeOuter::Docstring("Test2".to_string()),
                FieldAttributeOuter::Variable(FieldAttribute {
                    causality: Some(schema::Causality::Output.into()),
                    start: Some(syn::parse_quote!(0.0)),
                    ..Default::default()
                }),
                FieldAttributeOuter::Alias(AliasAttribute {
                    name: "alias_h".to_string(),
                    description: Some("Alias of h".to_string()),
                    display_unit: None,
                })
            ],
            "Second field should have 3 attributes: docstring, variable, and alias"
        );
    }

    #[test]
    fn test_field_description() {
        let input: syn::Field = syn::parse_quote! {
            /// This is a field description
            #[variable(causality = Output, start = 1.0)]
            height: f64
        };
        let field = Field::try_from(input).unwrap();
        assert_eq!(
            field.fold_description(),
            "This is a field description".to_string(),
            "Field description should match the docstring"
        );

        let input: syn::Field = syn::parse_quote! {
            #[variable(causality = Output, start = 1.0)]
            height: f64
        };
        let field = Field::try_from(input).unwrap();
        assert_eq!(
            field.fold_description(),
            "".to_string(),
            "Field description should be empty when no docstring"
        );
    }

    #[test]
    fn test_fields_without_fmu_attributes_are_ignored() {
        let input: syn::ItemStruct = syn::parse_quote! {
            struct TestModel {
                /// Height above ground (state output) - has FMU attribute
                #[variable(causality = Output, start = 1.0)]
                h: f64,

                /// User variable without FMU attributes - should be ignored
                internal_state: Vec<bool>,

                /// Another user variable - should be ignored
                helper_data: std::collections::HashMap<String, i32>,

                /// Velocity - has FMU attribute
                #[variable(causality = Output, start = 0.0)]
                v: f64,
            }
        };
        let fields = build_fields(input.fields);

        // Only fields with FMU attributes should be included
        assert_eq!(
            fields.len(),
            2,
            "Only fields with FMU attributes should be processed"
        );
        assert_eq!(fields[0].ident.to_string(), "h");
        assert_eq!(fields[1].ident.to_string(), "v");
    }
}