import re

with open("env/environment.py", "r") as f:
    code = f.read()

# 1. Look for the reward calculation call in `env/environment.py`
old_call = """        reward = self.reward_calculator.calculate_step_reward(
            state=self._state.model_dump(),
            action=tool,
            result=result,
            original_config=self._state.original_config,
            fixed_config=self._state.current_config,
            error_message=self._state.last_error,
            expected_config=self._state.task.expected_config,
            metadata=self._state.task.metadata,
        )

        reward_model = Reward(value=float(reward), components={"total": float(reward)})"""

new_call = """        reward, comps = self.reward_calculator.calculate_step_reward(
            state=self._state.model_dump(),
            action=tool,
            result=result,
            original_config=self._state.original_config,
            fixed_config=self._state.current_config,
            error_message=self._state.last_error,
            expected_config=self._state.task.expected_config,
            metadata=self._state.task.metadata,
        )

        reward_model = Reward(value=float(reward), components=comps)"""

code = code.replace(old_call, new_call)

with open("env/environment.py", "w") as f:
    f.write(code)