env/environment.py

from typing import Dict, Any
from env.models import Observation, Action, EpisodeState
from env.tasks import TASKS
from env.grader import grade_identify, grade_fix


def _clamp(reward: float) -> float:
    """Clamp reward strictly to [0.0, 1.0] per OpenEnv spec."""
    return round(max(0.0, min(reward, 1.0)), 4)


class CodeReviewEnvironment:
    """
    OpenEnv-compliant code review environment.
    Agent reads buggy code, identifies issues, and suggests fixes.
    3 tasks: easy (syntax) → medium (logic) → hard (performance)

    All rewards are clamped to [0.0, 1.0] before being returned.
    """

    def __init__(self):
        self._tasks = TASKS
        self._current_task_index = 0
        self._current_step = 0
        self._history = []
        self._phase = "identify"   # "identify" → "fix"
        self._identify_reward = 0.0
        self._total_reward = 0.0
        self._done = False

    # ------------------------------------------------------------------
    # OpenEnv required: reset()
    # ------------------------------------------------------------------
    def reset(self) -> Observation:
        self._current_task_index = 0
        self._current_step = 0
        self._history = []
        self._phase = "identify"
        self._identify_reward = 0.0
        self._total_reward = 0.0
        self._done = False
        return self._make_observation()

    # ------------------------------------------------------------------
    # OpenEnv required: step(action)
    # ------------------------------------------------------------------
    def step(self, action: Action) -> Dict[str, Any]:
        if self._done:
            return {
                "observation": self._make_observation(),
                "reward": 0.0,
                "done": True,
                "info": {"error": "Episode already done. Call reset()."},
            }

        task = self._current_task()
        reward = 0.0
        info = {}

        self._current_step += 1

        # ── IDENTIFY phase ──────────────────────────────────────────────
        if action.action_type == "identify":
            if self._phase == "fix":
                # Repeated identify after fix phase — penalise (clamped to 0.0)
                reward = 0.0
                info["warning"] = "Already in fix phase. Skipping repeated identify."
            else:
                reward = grade_identify(task["identify_keywords"], action.content)
                self._identify_reward = reward
                self._phase = "fix"
                info["phase_transition"] = "identify → fix"
                info["identify_score"] = reward

        # ── FIX phase ───────────────────────────────────────────────────
        elif action.action_type == "fix":
            if self._phase == "identify":
                # Jumped straight to fix without identifying — partial credit only
                fix_score = grade_fix(task["fixed_code"], action.content)
                reward = fix_score * 0.5   # halved because no identify step
                info["warning"] = "Skipped identify phase. Partial fix credit."
            else:
                fix_score = grade_fix(task["fixed_code"], action.content)
                # Identify bonus for continuous signal — final reward clamped to 1.0
                bonus = 0.1 if self._identify_reward >= 0.4 else 0.0
                reward = fix_score + bonus
                info["fix_score"] = fix_score
                info["identify_bonus"] = bonus

            # Clamp before recording and returning
            reward = _clamp(reward)
            self._total_reward += reward

            # Move to next task or end episode
            done = self._advance_task()
            self._phase = "identify"
            self._identify_reward = 0.0

            obs = self._make_observation()
            self._history.append(
                f"[task={task['id']}] action={action.action_type} reward={reward:.2f}"
            )
            info["task_completed"] = task["id"]
            info["task_difficulty"] = task["difficulty"]

            return {
                "observation": obs,
                "reward": reward,
                "done": done,
                "info": info,
            }

        else:
            # Unknown action type — penalise, clamped to 0.0
            reward = 0.0
            info["error"] = f"Unknown action_type '{action.action_type}'. Use 'identify' or 'fix'."

        reward = _clamp(reward)
        self._total_reward += reward
        self._history.append(
            f"[task={task['id']}] action={action.action_type} reward={reward:.2f}"
        )

        return {
            "observation": self._make_observation(),
            "reward": reward,
            "done": self._done,
            "info": info,
        }

    # ------------------------------------------------------------------
    # OpenEnv required: state (property)
    # ------------------------------------------------------------------
    @property
    def state(self) -> EpisodeState:
        task = self._current_task()
        return EpisodeState(
            current_task_id=task["id"],
            current_task_difficulty=task["difficulty"],
            current_task_category=task["category"],
            phase=self._phase,
            step=self._current_step,
            total_reward=round(self._total_reward, 4),
            done=self._done,
            history=list(self._history),
        )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------
    def _current_task(self) -> dict:
        # Clamp to last task when episode is done to avoid index error
        idx = min(self._current_task_index, len(self._tasks) - 1)
        return self._tasks[idx]

    def _make_observation(self) -> Observation:
        task = self._current_task()
        task_prompt = (
            f"[{task['difficulty'].upper()} | {task['category']}] {task['title']}\n"
            f"{task['description']}\n\n"
            f"Phase: {self._phase.upper()}\n"
            f"{'Identify the bug.' if self._phase == 'identify' else 'Fix the code.'}"
        )
        return Observation(
            code=task["code"],
            task=task_prompt,
            history=list(self._history),
            task_id=task["id"],
            language=task["language"],
            difficulty=task["difficulty"],
            category=task["category"],
        )

    def _advance_task(self) -> bool:
        """Move to next task. Returns True if episode is done."""
        self._current_task_index += 1
        if self._current_task_index >= len(self._tasks):
            self._done = True
            return True
        return False