import re

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

content = content.replace("use std::cell::Cell;", "use std::cell::{Cell, RefCell};")
content = content.replace("cache: Cell<MovingBoundaryCache>,", "cache: RefCell<MovingBoundaryCache>,")
content = content.replace("cache: Cell::new(MovingBoundaryCache::default()),", "cache: RefCell::new(MovingBoundaryCache::default()),")

# Patch compute_residuals
old_compute_residuals = """    fn compute_residuals(
        &self,
        state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        // For a moving boundary HX, we need to:
        // 1. Identify zones based on current inlet/outlet enthalpies
        // 2. Calculate UA for each zone
        // 3. Update nominal UA in the inner model
        // 4. Compute residuals using the standard model (e.g. EpsNtu)

        // HACK: For now, we use placeholder enthalpies to test the identification logic.
        // Proper port extraction will be added in Story 4.1.
        let h_in = 400_000.0;
        let h_out = 200_000.0;
        let p = 500_000.0;
        let m_refrig = 0.1; // Placeholder mass flow
        let t_sec_in = 300.0;
        let t_sec_out = 320.0;

        let mut cache = self.cache.take();
        let use_cache = cache.is_valid_for(p, m_refrig);

        let _discretization = if use_cache {
            cache.discretization.clone()
        } else {
            let (disc, h_sat_l, h_sat_v) = self.identify_zones(h_in, h_out, p, t_sec_in, t_sec_out)?;
            cache.valid = true;
            cache.p_ref = p;
            cache.m_ref = m_refrig;
            cache.h_sat_l = h_sat_l;
            cache.h_sat_v = h_sat_v;
            cache.discretization = disc.clone();
            disc
        };
        
        self.cache.set(cache);

        // Update total UA in the inner model (EpsNtuModel)
        // Note: HeatExchanger/Model are often immutable, but calibration indices can be used.
        // For now, we use Cell or similar if we need to store internal state, 
        // but typically the Model handles the UA.
        // self.inner.model.set_ua(discretization.total_ua); 
        // Wait, EpsNtuModel's UA is fixed. We might need a custom model or use ua_scale.
        
        self.inner.compute_residuals(state, residuals)
    }"""

new_compute_residuals = """    fn compute_residuals(
        &self,
        state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        let (p, m_refrig, t_sec_in, t_sec_out) = if let (Some(hot), Some(cold)) = (self.inner.hot_conditions(), self.inner.cold_conditions()) {
            (hot.pressure_pa(), hot.mass_flow_kg_s(), cold.temperature_k(), cold.temperature_k() + 5.0)
        } else {
            (500_000.0, 0.1, 300.0, 320.0)
        };

        // Extract enthalpies exactly as HeatExchanger does:
        let enthalpies = self.port_enthalpies(state)?;
        let h_in = enthalpies[0].to_joules_per_kg();
        let h_out = enthalpies[1].to_joules_per_kg();

        let mut cache = self.cache.borrow_mut();
        let use_cache = cache.is_valid_for(p, m_refrig);

        if !use_cache {
            let (disc, h_sat_l, h_sat_v) = self.identify_zones(h_in, h_out, p, t_sec_in, t_sec_out)?;
            cache.valid = true;
            cache.p_ref = p;
            cache.m_ref = m_refrig;
            cache.h_sat_l = h_sat_l;
            cache.h_sat_v = h_sat_v;
            cache.discretization = disc;
        }

        let total_ua = cache.discretization.total_ua;
        let base_ua = self.inner.ua_nominal();
        let custom_ua_scale = if base_ua > 0.0 { total_ua / base_ua } else { 1.0 };

        self.inner.compute_residuals_with_ua_scale(state, residuals, custom_ua_scale)
    }"""

content = content.replace(old_compute_residuals, new_compute_residuals)

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