Entropyk/crates/entropyk/src/builder.rs

use std::collections::HashMap;
use thiserror::Error;

use crate::ThermoError;

/// Error type for system builder operations.
#[derive(Error, Debug, Clone)]
pub enum SystemBuilderError {
    /// A component with the given name already exists in the builder.
    #[error("Component '{0}' already exists")]
    ComponentExists(String),

    /// The specified component name was not found in the builder.
    #[error("Component '{0}' not found")]
    ComponentNotFound(String),

    /// Failed to create an edge between two components.
    #[error("Failed to create edge from '{from}' to '{to}': {reason}")]
    EdgeFailed {
        /// Name of the source component.
        from: String,
        /// Name of the target component.
        to: String,
        /// Reason for the failure.
        reason: String,
    },

    /// The system must be finalized before this operation.
    #[error("System must be finalized before solving")]
    NotFinalized,

    /// Cannot build a system with no components.
    #[error("Cannot build an empty system")]
    EmptySystem,
}

/// A builder for creating thermodynamic systems with a fluent API.
///
/// The `SystemBuilder` provides an ergonomic way to construct thermodynamic
/// systems by adding components and edges with human-readable names.
///
/// # Example
///
/// ```
/// use entropyk::SystemBuilder;
///
/// let builder = SystemBuilder::new();
/// assert_eq!(builder.component_count(), 0);
/// ```
///
/// For real components, see the crate-level documentation.
pub struct SystemBuilder {
    system: entropyk_solver::System,
    component_names: HashMap<String, petgraph::graph::NodeIndex>,
    fluid_name: Option<String>,
}

impl SystemBuilder {
    /// Creates a new empty system builder.
    pub fn new() -> Self {
        Self {
            system: entropyk_solver::System::new(),
            component_names: HashMap::new(),
            fluid_name: None,
        }
    }

    /// Sets the default fluid for the system.
    ///
    /// This stores the fluid name for reference. The actual fluid assignment
    /// to components is handled at the component/port level.
    ///
    /// # Arguments
    ///
    /// * `fluid` - The fluid name (e.g., "R134a", "R410A", "CO2")
    #[inline]
    pub fn with_fluid(mut self, fluid: impl Into<String>) -> Self {
        self.fluid_name = Some(fluid.into());
        self
    }

    /// Adds a named component to the system.
    ///
    /// The name is used for later reference when creating edges.
    /// Returns an error if a component with the same name already exists.
    ///
    /// # Arguments
    ///
    /// * `name` - A unique identifier for this component
    /// * `component` - The component to add
    #[inline]
    pub fn component(
        mut self,
        name: &str,
        component: Box<dyn entropyk_components::Component>,
    ) -> Result<Self, SystemBuilderError> {
        if self.component_names.contains_key(name) {
            return Err(SystemBuilderError::ComponentExists(name.to_string()));
        }

        let idx = self.system.add_component(component);
        self.component_names.insert(name.to_string(), idx);

        Ok(self)
    }

    /// Creates an edge between two named components.
    ///
    /// The edge represents a fluid connection from the source component's
    /// outlet to the target component's inlet.
    ///
    /// # Arguments
    ///
    /// * `from` - Name of the source component
    /// * `to` - Name of the target component
    ///
    /// # Errors
    ///
    /// Returns an error if either component name is not found.
    #[inline]
    pub fn edge(mut self, from: &str, to: &str) -> Result<Self, SystemBuilderError> {
        let from_idx = self
            .component_names
            .get(from)
            .ok_or_else(|| SystemBuilderError::ComponentNotFound(from.to_string()))?;

        let to_idx = self
            .component_names
            .get(to)
            .ok_or_else(|| SystemBuilderError::ComponentNotFound(to.to_string()))?;

        self.system
            .add_edge(*from_idx, *to_idx)
            .map_err(|e| SystemBuilderError::EdgeFailed {
                from: from.to_string(),
                to: to.to_string(),
                reason: e.to_string(),
            })?;

        Ok(self)
    }

    /// Gets the underlying system without finalizing.
    ///
    /// This is useful when you need to perform additional operations
    /// on the system before finalizing.
    pub fn into_inner(self) -> entropyk_solver::System {
        self.system
    }

    /// Gets a reference to the component name to index mapping.
    pub fn component_names(&self) -> &HashMap<String, petgraph::graph::NodeIndex> {
        &self.component_names
    }

    /// Returns the number of components added so far.
    pub fn component_count(&self) -> usize {
        self.component_names.len()
    }

    /// Returns the number of edges created so far.
    pub fn edge_count(&self) -> usize {
        self.system.edge_count()
    }

    /// Builds and finalizes the system.
    ///
    /// This method consumes the builder and returns a finalized [`entropyk_solver::System`]
    /// ready for solving.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The system is empty (no components)
    /// - Finalization fails (e.g., invalid topology)
    pub fn build(self) -> Result<entropyk_solver::System, ThermoError> {
        if self.component_names.is_empty() {
            return Err(ThermoError::Builder(SystemBuilderError::EmptySystem));
        }

        let mut system = self.system;
        system.finalize()?;

        Ok(system)
    }

    /// Builds the system without finalizing.
    ///
    /// Use this when you need to perform additional operations
    /// that require an unfinalized system.
    pub fn build_unfinalized(self) -> Result<entropyk_solver::System, SystemBuilderError> {
        if self.component_names.is_empty() {
            return Err(SystemBuilderError::EmptySystem);
        }

        Ok(self.system)
    }
}

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

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

    struct MockComponent {
        n_eqs: usize,
    }

    impl entropyk_components::Component for MockComponent {
        fn compute_residuals(
            &self,
            _state: &entropyk_components::SystemState,
            _residuals: &mut entropyk_components::ResidualVector,
        ) -> Result<(), ComponentError> {
            Ok(())
        }

        fn jacobian_entries(
            &self,
            _state: &entropyk_components::SystemState,
            _jacobian: &mut entropyk_components::JacobianBuilder,
        ) -> Result<(), ComponentError> {
            Ok(())
        }

        fn n_equations(&self) -> usize {
            self.n_eqs
        }

        fn get_ports(&self) -> &[entropyk_components::ConnectedPort] {
            &[]
        }
    }

    #[test]
    fn test_builder_creates_system() {
        let builder = SystemBuilder::new();
        assert_eq!(builder.component_count(), 0);
        assert_eq!(builder.edge_count(), 0);
    }

    #[test]
    fn test_add_component() {
        let builder = SystemBuilder::new()
            .component("comp1", Box::new(MockComponent { n_eqs: 2 }))
            .unwrap();

        assert_eq!(builder.component_count(), 1);
    }

    #[test]
    fn test_duplicate_component_error() {
        let result = SystemBuilder::new()
            .component("comp", Box::new(MockComponent { n_eqs: 1 }))
            .unwrap()
            .component("comp", Box::new(MockComponent { n_eqs: 1 }));

        assert!(result.is_err());
        if let Err(SystemBuilderError::ComponentExists(name)) = result {
            assert_eq!(name, "comp");
        } else {
            panic!("Expected ComponentExists error");
        }
    }

    #[test]
    fn test_add_edge() {
        let builder = SystemBuilder::new()
            .component("a", Box::new(MockComponent { n_eqs: 1 }))
            .unwrap()
            .component("b", Box::new(MockComponent { n_eqs: 1 }))
            .unwrap()
            .edge("a", "b")
            .unwrap();

        assert_eq!(builder.edge_count(), 1);
    }

    #[test]
    fn test_edge_missing_component() {
        let result = SystemBuilder::new()
            .component("a", Box::new(MockComponent { n_eqs: 1 }))
            .unwrap()
            .edge("a", "nonexistent");

        assert!(result.is_err());
        if let Err(SystemBuilderError::ComponentNotFound(name)) = result {
            assert_eq!(name, "nonexistent");
        } else {
            panic!("Expected ComponentNotFound error");
        }
    }

    #[test]
    fn test_build_empty_system() {
        let result = SystemBuilder::new().build();
        assert!(result.is_err());
    }

    #[test]
    fn test_default() {
        let builder = SystemBuilder::default();
        assert_eq!(builder.component_count(), 0);
    }
}