from env.models import Observation, Action, Reward
from env.tasks import TASKS
from env.grader import grade_detect, grade_risk, grade_fix

class CodeGuardianEnv:
    def __init__(self):
        self.tasks = TASKS
        self.current_task_index = 0
        self.current_step = 0  # 0 = detect, 1 = risk, 2 = fix
        self.partial_score = 0
        self.done = False

    def reset(self):
        self.current_task_index = 0
        self.current_step = 0
        self.done = False
        self.partial_score = 0
        return self._get_obs()

    def _get_obs(self):
        task = self.tasks[self.current_task_index]
        return Observation(
            code=task["code"],
            intent=task["intent"],
            constraints=task["constraints"],
            task_id=task["task_id"],
            difficulty=task["difficulty"],
            step_index=self.current_step
        )

    def step(self, action: Action):
        """Execute one step in the environment with structured reward tracking"""
        # Early guard to prevent IndexError when environment is done
        if self.done:
            reward_obj = Reward(
                value=0.0,
                feedback="Environment is complete. Reset to continue.",
                partial_scores={},
                breakdown={}
            )
            return None, reward_obj, self.done
        
        task = self.tasks[self.current_task_index]
        weights = [0.3, 0.2, 0.5]  # detect, risk, fix weights
        step_types = ['detect', 'risk', 'fix']
        
        # Check if action order is valid
        if action.action_type != step_types[self.current_step]:
            score = 0.0
            feedback = f"Invalid action order: expected {step_types[self.current_step]}, got {action.action_type}"
            partial_scores = {}
        else:
            # Use deterministic graders
            original_code = task["code"]
            
            if action.action_type == "detect":
                score = grade_detect(action.detected_issues or [], task["expected_issues"])
            elif action.action_type == "risk":
                score = grade_risk(action.risk_level or "", task["expected_risk"])
            elif action.action_type == "fix":
                score = grade_fix(action.suggested_fix or "", task["canonical_fix"], original_code)
            else:
                score = 0.0
            
            # Weighted reward
            feedback = str(score)
            partial_scores = {action.action_type: score}
        
        # Calculate weighted reward and update partial score
        reward_value = score * weights[self.current_step] if partial_scores else -0.3
        self.partial_score += reward_value
        
        # Create structured reward with breakdown
        reward_obj = Reward(
            value=reward_value,
            feedback=feedback,
            partial_scores=partial_scores,
            breakdown=partial_scores.copy()  # Track detailed breakdown
        )
        
        # Advance step_index % 3, task if done
        if partial_scores:  # Only advance if action was valid
            self.current_step = (self.current_step + 1) % 3
            if self.current_step == 0:  # Completed full cycle
                self.current_task_index += 1
                if self.current_task_index >= len(self.tasks):
                    self.done = True
        
        next_obs = self._get_obs() if not self.done else None
        return next_obs, reward_obj, self.done

    def state(self):
        return {
            "current_task": self.current_task_index,
            "current_step": self.current_step,
            "partial_score": self.partial_score,
            "done": self.done
        }