import re

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

# 1. Change the signature of `calculate_step_reward`
old_sig = """    def calculate_step_reward(
        self,
        state: dict[str, Any] | None,
        action: str,
        result: dict[str, Any] | None,
        original_config: str | None = None,
        fixed_config: str | None = None,
        error_message: str | None = None,
        expected_config: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> float:"""

new_sig = """    def calculate_step_reward(
        self,
        state: dict[str, Any] | None,
        action: str,
        result: dict[str, Any] | None,
        original_config: str | None = None,
        fixed_config: str | None = None,
        error_message: str | None = None,
        expected_config: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> tuple[float, dict[str, float]]:"""

code = code.replace(old_sig, new_sig)

# 2. Change the return and body of `calculate_step_reward`
old_body = """        reward = 0.0
        reward += self._progress_reward(action, result)
        reward += self._execution_reward(result)
        reward += self._quality_reward(
            action=action,
            current_config=current_config,
            expected_config=expected_config,
            original_config=original_config,
            error_message=error_message,
            result=result,
            metadata=metadata,
        )
        reward += self._penalty_reward(state=state, result=result, current_config=current_config)

        return round(self._clamp_01(reward), 4)"""

new_body = """        prog = self._progress_reward(action, result)
        exec = self._execution_reward(result)
        det_score, hide_score, llm_score, qual = self._quality_reward(
            action=action,
            current_config=current_config,
            expected_config=expected_config,
            original_config=original_config,
            error_message=error_message,
            result=result,
            metadata=metadata,
        )
        pen = self._penalty_reward(state=state, result=result, current_config=current_config)
        
        reward = prog + exec + qual + pen
        reward_clamped = round(self._clamp_01(reward), 4)
        
        components = {
            "progress": prog,
            "execution": exec,
            "deterministic": float(det_score or 0.0),
            "hidden": float(hide_score or 0.0),
            "llm_judge": float(llm_score or 0.0),
            "penalty": pen,
            "total": reward_clamped
        }
        return reward_clamped, components"""

code = code.replace(old_body, new_body)

# 3. Change `_quality_reward` return signature and body
old_qual_sig = """    def _quality_reward(
        self,
        action: str,
        current_config: str,
        expected_config: str,
        original_config: str,
        error_message: str,
        result: dict[str, Any],
        metadata: dict[str, Any],
    ) -> float:"""

new_qual_sig = """    def _quality_reward(
        self,
        action: str,
        current_config: str,
        expected_config: str,
        original_config: str,
        error_message: str,
        result: dict[str, Any],
        metadata: dict[str, Any],
    ) -> tuple[float, float, float, float]:"""

code = code.replace(old_qual_sig, new_qual_sig)

old_qual_early_return = """        if not current_config or not expected_config:
            return 0.0010101"""
new_qual_early_return = """        if not current_config or not expected_config:
            return 0.0, 0.0, 0.0, 0.0010101"""
code = code.replace(old_qual_early_return, new_qual_early_return)

old_qual_return = """        quality_reward = 0.0
        quality_reward += self.QUALITY_WEIGHTS["deterministic"] * self._clamp_01(deterministic_score)
        quality_reward += self.QUALITY_WEIGHTS["hidden"] * self._clamp_01(hidden_pass_rate or 0.0)
        quality_reward += self.QUALITY_WEIGHTS["llm"] * self._clamp_01(llm_average)

        return quality_reward"""

new_qual_return = """        quality_reward = 0.0
        quality_reward += self.QUALITY_WEIGHTS["deterministic"] * self._clamp_01(deterministic_score)
        quality_reward += self.QUALITY_WEIGHTS["hidden"] * self._clamp_01(hidden_pass_rate or 0.0)
        quality_reward += self.QUALITY_WEIGHTS["llm"] * self._clamp_01(llm_average)

        return self._clamp_01(deterministic_score), self._clamp_01(hidden_pass_rate or 0.0), self._clamp_01(llm_average), quality_reward"""

code = code.replace(old_qual_return, new_qual_return)

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