rh_codegen/generators/trait_impl_generator/

mod.rs

//! Trait implementation generation functionality
//!
//! This module handles the generation of trait implementations for FHIR resources.

mod accessor_impl;
mod existence_impl;
mod mutator_impl;

use super::utils::GeneratorUtils;
use crate::fhir_types::StructureDefinition;
use crate::naming::Naming;
use crate::rust_types::RustTraitImpl;
use crate::CodegenResult;

/// Generator for trait implementations
pub struct TraitImplGenerator;

#[allow(dead_code)]
impl TraitImplGenerator {
    /// Create a new trait implementation generator
    pub fn new() -> Self {
        Self
    }

    /// Extract the base resource type from a FHIR baseDefinition URL
    /// For example: "http://hl7.org/fhir/StructureDefinition/Group" -> "Group"
    pub(crate) fn extract_base_resource_type(base_definition: &str) -> Option<String> {
        if base_definition.starts_with("http://hl7.org/fhir/StructureDefinition/") {
            if let Some(last_segment) = base_definition.split('/').next_back() {
                return Some(last_segment.to_string());
            }
        }
        None
    }

    /// Check if a baseDefinition indicates this is a core FHIR resource
    /// Core resources inherit directly from Resource or DomainResource
    pub(crate) fn is_core_resource(base_definition: &str) -> bool {
        matches!(
            base_definition,
            "http://hl7.org/fhir/StructureDefinition/Resource"
                | "http://hl7.org/fhir/StructureDefinition/DomainResource"
        )
    }

    /// Get the core resource type that a profile ultimately inherits from
    /// This handles known profile inheritance chains and can be extended with dynamic loading
    pub(crate) fn resolve_to_core_resource_type(
        base_resource_type: &str,
        _base_definition_url: &str,
    ) -> String {
        match base_resource_type.to_lowercase().as_str() {
            "vitalsigns" => "Observation".to_string(),
            "bodyweight" | "bodyheight" | "bmi" | "bodytemp" | "heartrate" | "resprate"
            | "oxygensat" => "Observation".to_string(),
            _ => base_resource_type.to_string(),
        }
    }

    /// Get the appropriate resource type for a structure definition
    /// For profiles, returns the base resource type; for core resources, returns the struct name
    pub(crate) fn get_resource_type_for_struct(
        struct_name: &str,
        structure_def: &StructureDefinition,
    ) -> String {
        if let Some(base_def) = &structure_def.base_definition {
            if Self::is_core_resource(base_def) {
                return struct_name.to_string();
            }

            if let Some(base_resource_type) = Self::extract_base_resource_type(base_def) {
                let core_resource_type =
                    Self::resolve_to_core_resource_type(&base_resource_type, base_def);

                if GeneratorUtils::is_fhir_resource_type(&core_resource_type) {
                    return core_resource_type;
                }
            }
        }

        struct_name.to_string()
    }

    /// Generate trait implementations for a FHIR resource
    pub fn generate_trait_impls(
        &self,
        structure_def: &StructureDefinition,
    ) -> CodegenResult<Vec<RustTraitImpl>> {
        let mut trait_impls = Vec::new();

        if structure_def.kind != "resource" {
            return Ok(trait_impls);
        }

        let struct_name = Naming::struct_name(structure_def);

        trait_impls.push(self.generate_resource_trait_impl(&struct_name, structure_def));
        trait_impls.push(self.generate_resource_mutators_trait_impl(&struct_name, structure_def));
        trait_impls.push(self.generate_resource_existence_trait_impl(&struct_name, structure_def));

        if let Some(base_def) = &structure_def.base_definition {
            if base_def.contains("DomainResource") {
                trait_impls.push(self.generate_domain_resource_trait_impl(&struct_name));
                trait_impls.push(self.generate_domain_resource_mutators_trait_impl(&struct_name));
                trait_impls.push(self.generate_domain_resource_existence_trait_impl(&struct_name));
            }
        }

        if struct_name != "Resource" {
            let specific_trait_impl =
                self.generate_specific_resource_trait_impl(&struct_name, structure_def);

            if !specific_trait_impl.is_empty() {
                trait_impls.push(specific_trait_impl);
            }

            let specific_mutators_trait_impl =
                self.generate_specific_resource_mutators_trait_impl(&struct_name, structure_def);

            if !specific_mutators_trait_impl.is_empty() {
                trait_impls.push(specific_mutators_trait_impl);
            }

            let specific_existence_trait_impl =
                self.generate_specific_resource_existence_trait_impl(&struct_name, structure_def);

            if !specific_existence_trait_impl.is_empty() {
                trait_impls.push(specific_existence_trait_impl);
            }
        }

        Ok(trait_impls)
    }

    /// Get the base access pattern for resource fields
    pub(crate) fn get_resource_base_access(
        &self,
        struct_name: &str,
        structure_def: &StructureDefinition,
    ) -> (String, bool) {
        if struct_name == "Resource" {
            ("self".to_string(), false)
        } else if struct_name == "DomainResource" {
            ("self.base".to_string(), false)
        } else if let Some(base_def) = &structure_def.base_definition {
            if base_def.contains("DomainResource") {
                ("self.base.base".to_string(), false)
            } else if base_def.contains("Resource") && struct_name != "DomainResource" {
                ("self.base".to_string(), false)
            } else if base_def.starts_with("http://hl7.org/fhir/StructureDefinition/") {
                ("self.base".to_string(), true)
            } else {
                ("self.base.base".to_string(), false)
            }
        } else {
            ("self.base.base".to_string(), false)
        }
    }

    /// Check if we should generate an accessor implementation for this element
    pub(crate) fn should_generate_accessor_impl(
        &self,
        element: &crate::fhir_types::ElementDefinition,
        structure_def: &StructureDefinition,
    ) -> bool {
        let field_path = &element.path;
        let base_name = &structure_def.name;

        if !field_path.starts_with(base_name) {
            return false;
        }

        let path_parts: Vec<&str> = field_path.split('.').collect();
        if path_parts.len() != 2 {
            return false;
        }

        if path_parts[0] != base_name {
            return false;
        }

        let field_name = path_parts[1];
        !field_name.ends_with("[x]")
    }

    /// Check if the element types contain BackboneElement
    pub(crate) fn is_backbone_element(
        &self,
        element_types: &[crate::fhir_types::ElementType],
    ) -> bool {
        element_types
            .iter()
            .any(|et| et.code.as_deref() == Some("BackboneElement"))
    }

    /// Get the specific nested type for a BackboneElement field
    pub(crate) fn get_nested_type_for_backbone_element(
        &self,
        element: &crate::fhir_types::ElementDefinition,
        is_array: bool,
    ) -> crate::rust_types::RustType {
        let path_parts: Vec<&str> = element.path.split('.').collect();

        if path_parts.len() == 2 {
            let resource_name = path_parts[0];
            let field_name = path_parts[1];

            let field_name_pascal = crate::naming::Naming::to_pascal_case(field_name);
            let nested_type_name = format!("{resource_name}{field_name_pascal}");

            let rust_type = crate::rust_types::RustType::Custom(nested_type_name);

            if is_array {
                crate::rust_types::RustType::Vec(Box::new(rust_type))
            } else {
                rust_type
            }
        } else {
            let rust_type = crate::rust_types::RustType::Custom("BackboneElement".to_string());
            if is_array {
                crate::rust_types::RustType::Vec(Box::new(rust_type))
            } else {
                rust_type
            }
        }
    }

    /// Get the inner type for slice return type
    pub(crate) fn get_inner_type_for_slice(
        &self,
        rust_type: &crate::rust_types::RustType,
    ) -> String {
        match rust_type {
            crate::rust_types::RustType::Vec(inner) => inner.to_string(),
            crate::rust_types::RustType::Option(inner) => {
                if let crate::rust_types::RustType::Vec(vec_inner) = inner.as_ref() {
                    vec_inner.to_string()
                } else {
                    inner.to_string()
                }
            }
            _ => rust_type.to_string(),
        }
    }

    /// Get the type for option return type
    #[allow(dead_code)]
    pub(crate) fn get_type_for_option(&self, rust_type: &crate::rust_types::RustType) -> String {
        match rust_type {
            crate::rust_types::RustType::Option(inner) => inner.to_string(),
            _ => rust_type.to_string(),
        }
    }

    /// Check if a RustType implements the Copy trait
    pub(crate) fn is_copy_type(&self, rust_type: &crate::rust_types::RustType) -> bool {
        match rust_type {
            crate::rust_types::RustType::Boolean
            | crate::rust_types::RustType::Integer
            | crate::rust_types::RustType::Float => true,

            crate::rust_types::RustType::Option(inner) => self.is_copy_type(inner),

            crate::rust_types::RustType::Custom(type_name) => {
                self.is_copy_primitive_type(type_name)
            }

            _ => false,
        }
    }

    /// Check if a custom type name represents a FHIR primitive type that is Copy
    pub(crate) fn is_copy_primitive_type(&self, type_name: &str) -> bool {
        matches!(
            type_name,
            "BooleanType" | "IntegerType" | "UnsignedIntType" | "PositiveIntType" | "DecimalType"
        )
    }

    /// Check if a rust type represents an enum
    #[allow(dead_code)]
    pub(crate) fn is_enum_type(&self, rust_type: &crate::rust_types::RustType) -> bool {
        match rust_type {
            crate::rust_types::RustType::Custom(type_name) => self.is_enum_type_name(type_name),
            _ => false,
        }
    }

    /// Check if a type name represents a FHIR enum type
    #[allow(dead_code)]
    pub(crate) fn is_enum_type_name(&self, type_name: &str) -> bool {
        type_name.ends_with("Status")
            || type_name.ends_with("Kind")
            || type_name.ends_with("Code")
            || type_name.ends_with("Codes")
            || type_name.ends_with("Priority")
            || type_name.ends_with("Intent")
            || matches!(
                type_name,
                "PublicationStatus"
                    | "CapabilityStatementKind"
                    | "CodeSearchSupport"
                    | "FmStatus"
                    | "ReportStatusCodes"
                    | "ReportResultCodes"
                    | "VerificationresultStatus"
                    | "TaskStatus"
                    | "TaskIntent"
                    | "RequestPriority"
                    | "SupplydeliveryStatus"
                    | "SupplyrequestStatus"
            )
    }

    /// Determine the return type for a field accessor method
    #[allow(dead_code)]
    pub(crate) fn determine_method_return_type(
        &self,
        element: &crate::fhir_types::ElementDefinition,
    ) -> String {
        let is_optional = element.min.unwrap_or(0) == 0;

        let is_array = element
            .max
            .as_ref()
            .is_some_and(|max| max == "*" || max.parse::<u32>().unwrap_or(1) > 1);

        let base_type = if let Some(element_types) = &element.element_type {
            if let Some(first_type) = element_types.first() {
                if let Some(code) = &first_type.code {
                    match code.as_str() {
                        "string" | "code" | "id" | "markdown" | "uri" | "url" | "canonical"
                        | "dateTime" | "date" | "time" | "instant" | "base64Binary" | "oid"
                        | "uuid" => "String".to_string(),
                        "boolean" => "bool".to_string(),
                        "integer" | "positiveInt" | "unsignedInt" => "i32".to_string(),
                        "decimal" => "f64".to_string(),
                        "Reference" => "crate::datatypes::reference::Reference".to_string(),
                        "Identifier" => "crate::datatypes::identifier::Identifier".to_string(),
                        "CodeableConcept" => {
                            "crate::datatypes::codeable_concept::CodeableConcept".to_string()
                        }
                        "Coding" => "crate::datatypes::coding::Coding".to_string(),
                        "Address" => "crate::datatypes::address::Address".to_string(),
                        "HumanName" => "crate::datatypes::human_name::HumanName".to_string(),
                        "ContactPoint" => {
                            "crate::datatypes::contact_point::ContactPoint".to_string()
                        }
                        "Attachment" => "crate::datatypes::attachment::Attachment".to_string(),
                        "Annotation" => "crate::datatypes::annotation::Annotation".to_string(),
                        "BackboneElement" => {
                            "crate::datatypes::backbone_element::BackboneElement".to_string()
                        }
                        _ => "String".to_string(),
                    }
                } else {
                    "String".to_string()
                }
            } else {
                "String".to_string()
            }
        } else {
            "String".to_string()
        };

        if is_array {
            format!("Vec<{base_type}>")
        } else if is_optional {
            format!("Option<{base_type}>")
        } else {
            base_type
        }
    }

    /// Generate the method body for a field accessor
    #[allow(dead_code)]
    pub(crate) fn generate_method_body(
        &self,
        field_name: &str,
        element: &crate::fhir_types::ElementDefinition,
    ) -> String {
        let rust_field_name = if field_name == "type" {
            "type_".to_string()
        } else {
            crate::naming::Naming::field_name(field_name)
        };

        let field_access = format!("self.{rust_field_name}");

        let is_optional = element.min.unwrap_or(0) == 0;
        let is_array = element
            .max
            .as_ref()
            .is_some_and(|max| max == "*" || max.parse::<u32>().unwrap_or(1) > 1);

        if is_array {
            format!("{field_access}.clone()")
        } else if let Some(type_def) = element
            .element_type
            .as_ref()
            .and_then(|types| types.first())
        {
            if let Some(code) = &type_def.code {
                match code.as_str() {
                    "string" | "code" | "id" | "markdown" | "uri" | "url" | "canonical"
                    | "dateTime" | "date" | "time" | "instant" | "base64Binary" | "oid"
                    | "uuid" => {
                        if is_optional {
                            format!("{field_access}.as_ref().map(|s| s.to_string())")
                        } else {
                            format!("{field_access}.to_string()")
                        }
                    }
                    "boolean" => {
                        if is_optional {
                            format!("{field_access}.map(|b| b.into())")
                        } else {
                            format!("{field_access}.into()")
                        }
                    }
                    "integer" | "positiveInt" | "unsignedInt" => {
                        if is_optional {
                            format!("{field_access}.map(|i| i.into())")
                        } else {
                            format!("{field_access}.into()")
                        }
                    }
                    "decimal" => {
                        if is_optional {
                            format!("{field_access}.map(|d| d.into())")
                        } else {
                            format!("{field_access}.into()")
                        }
                    }
                    "CodeableConcept" | "Reference" | "Identifier" | "Coding" | "Address"
                    | "HumanName" | "ContactPoint" | "Attachment" | "Annotation"
                    | "BackboneElement" => {
                        format!("{field_access}.clone()")
                    }
                    _ => {
                        if is_optional {
                            format!("{field_access}.as_ref().map(|v| format!(\"{{:?}}\", v))")
                        } else {
                            format!("format!(\"{{:?}}\", {field_access})")
                        }
                    }
                }
            } else {
                format!("{field_access}.clone()")
            }
        } else {
            format!("{field_access}.clone()")
        }
    }

    /// Get the Rust type for a field element, considering ValueSet bindings.
    pub(crate) fn get_field_rust_type(
        &self,
        element: &crate::fhir_types::ElementDefinition,
        field_name: &str,
    ) -> CodegenResult<crate::rust_types::RustType> {
        use crate::rust_types::RustType;

        let Some(element_type) = element.element_type.as_ref().and_then(|t| t.first()) else {
            return Ok(RustType::String);
        };

        let Some(code) = &element_type.code else {
            return Ok(RustType::String);
        };

        if code == "code" {
            if let Some(binding) = &element.binding {
                if binding.strength == "required" {
                    if let Some(value_set_url) = &binding.value_set {
                        if let Some(enum_name) =
                            self.extract_enum_name_from_value_set(value_set_url)
                        {
                            let resource_name = element.path.split('.').next().unwrap_or("");
                            if enum_name != resource_name {
                                return Ok(RustType::Custom(enum_name));
                            }
                        }
                    }
                }
            }
        }

        use crate::generators::TypeUtilities;
        TypeUtilities::map_fhir_type_to_rust(element_type, field_name, &element.path)
    }

    /// Extract enum type name from a ValueSet URL
    pub(crate) fn extract_enum_name_from_value_set(&self, url: &str) -> Option<String> {
        let url_without_version = url.split('|').next().unwrap_or(url);

        let value_set_name = url_without_version.split('/').next_back()?;

        let name = value_set_name
            .split(&['-', '.'][..])
            .filter(|part| !part.is_empty())
            .map(|part| {
                let mut chars = part.chars();
                match chars.next() {
                    None => String::new(),
                    Some(first) => first.to_uppercase().collect::<String>() + chars.as_str(),
                }
            })
            .collect::<String>();

        if name.chars().next().unwrap_or('0').is_ascii_digit() {
            Some(format!("ValueSet{name}"))
        } else {
            Some(name)
        }
    }
}

impl Default for TraitImplGenerator {
    fn default() -> Self {
        Self::new()
    }
}

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

    fn create_test_structure_definition(
        name: &str,
        base_definition: Option<&str>,
    ) -> StructureDefinition {
        StructureDefinition {
            resource_type: "StructureDefinition".to_string(),
            id: name.to_lowercase(),
            url: format!("http://test.com/{name}"),
            version: Some("1.0.0".to_string()),
            name: name.to_string(),
            title: Some(name.to_string()),
            status: "active".to_string(),
            description: None,
            purpose: None,
            kind: "resource".to_string(),
            is_abstract: false,
            base_type: "Resource".to_string(),
            base_definition: base_definition.map(|s| s.to_string()),
            differential: None,
            snapshot: None,
        }
    }

    #[test]
    fn test_resource_type_for_core_resource() {
        let patient = create_test_structure_definition(
            "Patient",
            Some("http://hl7.org/fhir/StructureDefinition/DomainResource"),
        );

        let result = TraitImplGenerator::get_resource_type_for_struct("Patient", &patient);
        assert_eq!(
            result, "Patient",
            "Core resource should return its own name"
        );
    }

    #[test]
    fn test_resource_type_for_group_profile() {
        let group_definition = create_test_structure_definition(
            "GroupDefinition",
            Some("http://hl7.org/fhir/StructureDefinition/Group"),
        );

        let result =
            TraitImplGenerator::get_resource_type_for_struct("GroupDefinition", &group_definition);
        assert_eq!(result, "Group", "Group profile should return 'Group'");
    }

    #[test]
    fn test_resource_type_for_observation_profile() {
        let vital_signs = create_test_structure_definition(
            "VitalSigns",
            Some("http://hl7.org/fhir/StructureDefinition/Observation"),
        );

        let result = TraitImplGenerator::get_resource_type_for_struct("VitalSigns", &vital_signs);
        assert_eq!(
            result, "Observation",
            "Observation profile should return 'Observation'"
        );
    }

    #[test]
    fn test_resource_type_for_profile_on_profile() {
        let bmi = create_test_structure_definition(
            "BMI",
            Some("http://hl7.org/fhir/StructureDefinition/vitalsigns"),
        );

        let result = TraitImplGenerator::get_resource_type_for_struct("BMI", &bmi);
        assert_eq!(
            result, "Observation",
            "BMI profile should resolve to 'Observation' via vitalsigns"
        );
    }

    #[test]
    fn test_resource_type_without_base_definition() {
        let custom_resource = create_test_structure_definition("CustomResource", None);

        let result =
            TraitImplGenerator::get_resource_type_for_struct("CustomResource", &custom_resource);
        assert_eq!(
            result, "CustomResource",
            "Resource without baseDefinition should return struct name"
        );
    }

    #[test]
    fn test_is_core_resource() {
        assert!(TraitImplGenerator::is_core_resource(
            "http://hl7.org/fhir/StructureDefinition/Resource"
        ));
        assert!(TraitImplGenerator::is_core_resource(
            "http://hl7.org/fhir/StructureDefinition/DomainResource"
        ));
        assert!(!TraitImplGenerator::is_core_resource(
            "http://hl7.org/fhir/StructureDefinition/Patient"
        ));
        assert!(!TraitImplGenerator::is_core_resource(
            "http://hl7.org/fhir/StructureDefinition/Group"
        ));
    }

    #[test]
    fn test_extract_base_resource_type() {
        assert_eq!(
            TraitImplGenerator::extract_base_resource_type(
                "http://hl7.org/fhir/StructureDefinition/Group"
            ),
            Some("Group".to_string())
        );
        assert_eq!(
            TraitImplGenerator::extract_base_resource_type(
                "http://hl7.org/fhir/StructureDefinition/Observation"
            ),
            Some("Observation".to_string())
        );
        assert_eq!(
            TraitImplGenerator::extract_base_resource_type(
                "http://hl7.org/fhir/StructureDefinition/vitalsigns"
            ),
            Some("vitalsigns".to_string())
        );
        assert_eq!(
            TraitImplGenerator::extract_base_resource_type("invalid-url"),
            None
        );
    }

    #[test]
    fn test_resolve_to_core_resource_type() {
        assert_eq!(
            TraitImplGenerator::resolve_to_core_resource_type(
                "vitalsigns",
                "http://hl7.org/fhir/StructureDefinition/vitalsigns"
            ),
            "Observation"
        );

        assert_eq!(
            TraitImplGenerator::resolve_to_core_resource_type(
                "Patient",
                "http://hl7.org/fhir/StructureDefinition/Patient"
            ),
            "Patient"
        );
        assert_eq!(
            TraitImplGenerator::resolve_to_core_resource_type(
                "Group",
                "http://hl7.org/fhir/StructureDefinition/Group"
            ),
            "Group"
        );

        assert_eq!(
            TraitImplGenerator::resolve_to_core_resource_type(
                "bmi",
                "http://hl7.org/fhir/StructureDefinition/bmi"
            ),
            "Observation"
        );

        assert_eq!(
            TraitImplGenerator::resolve_to_core_resource_type(
                "UnknownProfile",
                "http://hl7.org/fhir/StructureDefinition/UnknownProfile"
            ),
            "UnknownProfile"
        );
    }

    #[test]
    fn test_empty_trait_implementations_are_filtered() {
        let generator = TraitImplGenerator::new();

        let mut structure_def = create_test_structure_definition("EmptyProfile", None);
        structure_def.differential = Some(StructureDefinitionDifferential { element: vec![] });

        let trait_impls = generator.generate_trait_impls(&structure_def).unwrap();

        assert!(
            !trait_impls.is_empty(),
            "Should have at least Resource trait impl"
        );

        let specific_trait_name = format!(
            "crate::traits::{}::{}Accessors",
            crate::naming::Naming::to_snake_case("EmptyProfile"),
            "EmptyProfile"
        );

        let has_empty_specific_impl = trait_impls
            .iter()
            .any(|impl_| impl_.trait_name == specific_trait_name && impl_.is_empty());

        assert!(
            !has_empty_specific_impl,
            "Should not include empty specific trait implementations"
        );
    }
}