env/codedebugger_env.py

import time
import re
from env.executor import CodeExecutor
from agents.reward import calculate_reward
from agents.critic import run_all_anti_hack_checks
from fastapi import FastAPI
import threading

class CodeDebuggerEnv:
    def __init__(self, max_iterations=5):
        self.executor = CodeExecutor(timeout_seconds=10)
        self.max_iterations = max_iterations
        self.problem = None
        self.iteration = 0
        self.history = []
        self.prev_tests_passed = 0
    
    def reset(self, problem: dict) -> dict:
        """Reset for new problem. Returns initial observation."""
        self.problem = problem
        self.iteration = 0
        self.history = []
        self.prev_tests_passed = 0
        return {
            "problem_id": problem["id"],
            "difficulty": problem["difficulty"],
            "description": problem["description"],
            "buggy_code": problem["buggy_code"],
            "error_type": problem.get("error_type", ""),
            "test_cases": problem["test_cases"],
            "iteration": 0
        }
    
    def step(self, action: str) -> tuple:
        """
        action = submitted fixed code string
        
        Process:
        1. Increment iteration
        2. Safety check on submitted code
        3. If unsafe: skip tests, return penalty reward
        4. Run all tests via executor
        5. Run anti-hack checks
        6. Calculate reward
        7. Store in history
        8. Return (observation, reward_dict, done, info)
        """
        self.iteration += 1
        
        # 2. Safety check
        safety = self.executor.check_code_safety(action)
        
        func_name = extract_function_name(self.problem["buggy_code"])
        
        if not safety["safe"]:
            # 3. If unsafe: skip tests
            test_results = {
                "tests_passed": 0,
                "tests_total": len(self.problem["test_cases"]),
                "pass_rate": 0.0,
                "results": [],
                "execution_time_total": 0.0
            }
            execution_time = 0.0
            anti_hack = {"all_passed": True, "total_penalty": 0.0, "checks": {}}
        else:
            # 4. Run all tests via executor
            start_t = time.time()
            test_results = self.executor.run_all_tests(
                code=action,
                function_name=func_name,
                test_cases=self.problem["test_cases"]
            )
            execution_time = time.time() - start_t
            
            # 5. Run anti-hack checks
            anti_hack = run_all_anti_hack_checks(
                code=action,
                original_code=self.problem["buggy_code"],
                test_cases=self.problem["test_cases"],
                test_results=test_results,
                execution_time=test_results.get("execution_time_total", execution_time)
            )
            
        # 6. Calculate reward
        reward_dict = calculate_reward(
            code=action,
            test_results=test_results,
            original_buggy_code=self.problem["buggy_code"],
            iteration=self.iteration,
            execution_time=test_results.get("execution_time_total", execution_time),
            safety_check=safety,
            previous_tests_passed=self.prev_tests_passed,
            error_type=self.problem.get("error_type", "")
        )
        
        # incorporate critic penalties into reward
        critic_penalty = anti_hack.get("total_penalty", 0.0)
        reward_dict["anti_hack_penalty"] += critic_penalty
        reward_dict["total"] += critic_penalty
        reward_dict["total"] = max(0.0, reward_dict["total"]) # clamp to >= 0
        
        # 7. Store in history
        self.history.append({
            "iteration": self.iteration,
            "code": action,
            "reward": reward_dict,
            "test_results": test_results
        })
        
        self.prev_tests_passed = test_results.get("tests_passed", 0)
        
        # done = True if solved or out of iterations
        done = False
        if test_results.get("pass_rate", 0.0) == 1.0:
            done = True
        elif self.iteration >= self.max_iterations:
            done = True
            
        # 8. Return
        observation = {
            "problem_id": self.problem["id"],
            "iteration": self.iteration,
            "code": action,
            "test_results": test_results,
            "safety": safety,
            "anti_hack": anti_hack
        }
        
        info = {
            "solved": test_results.get("pass_rate", 0.0) == 1.0,
            "safety_violations": safety.get("violations", []),
            "anti_hack_failed": not anti_hack.get("all_passed", True)
        }
        
        return observation, reward_dict, done, info

    def render(self) -> str:
        """
        Returns formatted display string showing:
        - Problem title [difficulty]
        - Iteration N/max
        - Test results table
        - Reward breakdown
        - Issues list
        - Best reward so far
        """
        if not self.problem:
            return "Environment not initialized."
            
        title = self.problem.get("title", "Unknown")
        diff = self.problem.get("difficulty", "unknown").upper()
        
        lines = [
            f"=== {title} [{diff}] ===",
            f"Iteration: {self.iteration} / {self.max_iterations}"
        ]
        
        if not self.history:
            lines.append("No actions taken yet.")
            return "\n".join(lines)
            
        last = self.history[-1]
        tr = last["test_results"]
        rw = last["reward"]
        
        passed = tr.get("tests_passed", 0)
        total = tr.get("tests_total", 0)
        rate = tr.get("pass_rate", 0.0)
        lines.append(f"\nTest Results: {passed}/{total} passed ({rate*100:.1f}%)")
        
        lines.append("\nReward Breakdown:")
        for k, v in rw.items():
            if k != "total":
                lines.append(f"  {k}: {v:.2f}")
        lines.append(f"  --> TOTAL: {rw['total']:.2f}")
        
        issues = []
        if not last.get("safety", {}).get("safe", True):
            issues.append(f"Unsafe code detected: {last['safety'].get('violations')}")
        if not last.get("anti_hack", {}).get("all_passed", True):
            for name, chk in last["anti_hack"].get("checks", {}).items():
                if not chk.get("passed", True):
                    issues.append(f"Anti-hack failure ({name}): {chk.get('message')}")
                    
        if issues:
            lines.append("\nIssues:")
            for issue in issues:
                lines.append(f"  - {issue}")
                
        best_reward = max(h["reward"]["total"] for h in self.history)
        lines.append(f"\nBest reward so far: {best_reward:.2f}")
        
        return "\n".join(lines)

def extract_function_name(code: str) -> str:
    """Use regex to find 'def function_name(' in code"""
    match = re.search(r"def\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*\(", code)
    if match:
        return match.group(1)
    return "unknown"

# FastAPI app
app = FastAPI(title="CodeDebugger RL Environment")

env_instance = CodeDebuggerEnv()

@app.post("/reset")
def api_reset(problem: dict):
    return env_instance.reset(problem)

@app.post("/step") 
def api_step(payload: dict):
    action = payload.get("code", "")
    obs, reward, done, info = env_instance.step(action)
    return {"observation": obs, "reward": reward, 
            "done": done, "info": info}

@app.get("/render")
def api_render():
    return {"display": env_instance.render()}

@app.get("/health")
def health():
    return {"status": "ok", "version": "1.0.0"}