import re

with open("crates/components/src/heat_exchanger/exchanger.rs", "r") as f:
    content = f.read()

# Add getters for hot/cold conditions
content = content.replace("pub fn hot_fluid_id(&self) -> Option<&FluidsFluidId> {",
"""pub fn hot_conditions(&self) -> Option<&HxSideConditions> {
        self.hot_conditions.as_ref()
    }

    pub fn cold_conditions(&self) -> Option<&HxSideConditions> {
        self.cold_conditions.as_ref()
    }

    pub fn hot_fluid_id(&self) -> Option<&FluidsFluidId> {""")


# Add compute_residuals_with_ua_scale
# Find compute_residuals inside impl Component
# We'll just add a method to `impl<Model: HeatTransferModel> HeatExchanger<Model>`
method_to_add = """
    pub fn compute_residuals_with_ua_scale(
        &self,
        _state: &StateSlice,
        residuals: &mut ResidualVector,
        custom_ua_scale: f64,
    ) -> Result<(), ComponentError> {
        self.do_compute_residuals(_state, residuals, Some(custom_ua_scale))
    }
"""
content = content.replace("impl<Model: HeatTransferModel + 'static> Component for HeatExchanger<Model> {",
method_to_add + "\nimpl<Model: HeatTransferModel + 'static> Component for HeatExchanger<Model> {")

# Then modify compute_residuals to call do_compute_residuals
content = content.replace("fn compute_residuals(", "fn do_compute_residuals(\n        &self,\n        _state: &StateSlice,\n        residuals: &mut ResidualVector,\n        custom_ua_scale: Option<f64>,\n    ) -> Result<(), ComponentError> {\n        if residuals.len() < self.n_equations() {")

# Wait, replace the top of compute_residuals:
old_start = """    fn compute_residuals(
        &self,
        _state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {"""

new_start = """    fn compute_residuals(
        &self,
        _state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        self.do_compute_residuals(_state, residuals, None)
    }

    fn do_compute_residuals(
        &self,
        _state: &StateSlice,
        residuals: &mut ResidualVector,
        custom_ua_scale: Option<f64>,
    ) -> Result<(), ComponentError> {"""

content = content.replace(old_start, new_start)

# inside do_compute_residuals, replace `let dynamic_f_ua = self.calib_indices.f_ua.map(|idx| _state[idx]);`
old_f_ua = "let dynamic_f_ua = self.calib_indices.f_ua.map(|idx| _state[idx]);"
new_f_ua = "let dynamic_f_ua = custom_ua_scale.or_else(|| self.calib_indices.f_ua.map(|idx| _state[idx]));"
content = content.replace(old_f_ua, new_f_ua)

with open("crates/components/src/heat_exchanger/exchanger.rs", "w") as f:
    f.write(content)