#!/usr/bin/env python3
import os, sys, subprocess, tempfile, json, traceback, time, resource

class ExecutionSandbox:
    def __init__(self, engine=None, timeout=5, max_memory_mb=256):
        self.engine = engine
        self.timeout = timeout
        self.max_memory = max_memory_mb * 1024 * 1024
        self.execution_history = []

    def _run_in_subprocess(self, code, test_input=None):
        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
            f.write(code)
            tmp = f.name
        try:
            script = f"""
import sys, json, traceback, builtins
import resource
resource.setrlimit(resource.RLIMIT_CPU, ({self.timeout}, {self.timeout}))
resource.setrlimit(resource.RLIMIT_AS, ({self.max_memory}, {self.max_memory}))
_original = builtins.__import__
def _restricted(name, *args, **kwargs):
    blocked = ['os.system', 'subprocess', 'socket', 'urllib', 'http.client', 'ftplib']
    if any(b in name for b in blocked):
        raise ImportError(f"Blocked: {{name}}")
    return _original(name, *args, **kwargs)
builtins.__import__ = _restricted
result = {{"stdout": "", "stderr": "", "output": None, "error": None, "success": False}}
old_out, old_err = sys.stdout, sys.stderr
try:
    from io import StringIO
    sys.stdout, sys.stderr = StringIO(), StringIO()
    ns = {{"__builtins__": builtins}}
    with open("{tmp}") as f:
        exec(f.read(), ns)
    for name, obj in ns.items():
        if callable(obj) and name in ["main", "run", "solve", "test"]:
            try:
                result["output"] = obj("{test_input or ''}") if "{test_input or ''}" else obj()
            except Exception as e:
                result["error"] = str(e)
    result["stdout"] = sys.stdout.getvalue()
    result["stderr"] = sys.stderr.getvalue()
    result["success"] = result["error"] is None
except Exception as e:
    result["error"] = traceback.format_exc()
    result["success"] = False
finally:
    sys.stdout, sys.stderr = old_out, old_err
print(json.dumps(result))
"""
            proc = subprocess.run([sys.executable, '-c', script], capture_output=True, text=True, timeout=self.timeout+2, cwd='/tmp')
            out = proc.stdout.strip()
            if out:
                try:
                    return json.loads(out.split('\n')[-1])
                except:
                    return {"stdout": proc.stdout, "stderr": proc.stderr, "error": proc.stderr or "Failed", "success": False}
            return {"stdout": "", "stderr": proc.stderr, "error": proc.stderr or "No output", "success": False}
        except subprocess.TimeoutExpired:
            return {"error": f"Timeout after {self.timeout}s", "success": False}
        except Exception as e:
            return {"error": str(e), "success": False}
        finally:
            try: os.unlink(tmp)
            except: pass

    def execute_and_verify(self, code, expected_output=None, test_input=None):
        result = self._run_in_subprocess(code, test_input)
        if expected_output is not None:
            actual = result.get("output")
            result["verified"] = True
            result["match"] = actual == expected_output
            result["expected"] = expected_output
            result["actual"] = actual
        else:
            result["verified"] = False
        if self.engine and hasattr(self.engine, '_update_pheromone'):
            success = result.get("success") and result.get("match", True)
            for idx in range(1000):
                self.engine._update_pheromone(idx, success)
        self.execution_history.append({"success": result.get("success"), "match": result.get("match")})
        return result

    def auto_fix(self, code, error, max_attempts=3):
        if not self.engine:
            return code, {"fixed": False}
        current = code
        for attempt in range(max_attempts):
            prompt = f"Fix this Python code. Error: {error}\n\n```python\n{current}\n```\n\nProvide ONLY the corrected code."
            try:
                fixed = self.engine.generate(prompt, max_tokens=800)
                if "```python" in fixed:
                    fixed = fixed.split("```python")[1].split("```")[0]
                elif "```" in fixed:
                    fixed = fixed.split("```")[1].split("```")[0]
                result = self._run_in_subprocess(fixed)
                if result.get("success"):
                    return fixed, {"fixed": True, "attempts": attempt + 1}
                current = fixed
                error = result.get("error", "Unknown")
            except Exception as e:
                return current, {"fixed": False, "reason": str(e)}
        return current, {"fixed": False, "reason": f"Failed after {max_attempts}"}

    def get_stats(self):
        total = len(self.execution_history)
        if total == 0:
            return {"total": 0, "success_rate": 0}
        successes = sum(1 for h in self.execution_history if h.get("success"))
        matches = sum(1 for h in self.execution_history if h.get("match"))
        return {"total": total, "successes": successes, "success_rate": successes/total*100, "matches": matches}

if __name__ == "__main__":
    s = ExecutionSandbox()
    r1 = s.execute_and_verify("def add(a,b): return a+b\nprint(add(2,3))")
    print(f"Test 1: success={r1['success']}, stdout={r1.get('stdout','')[:50]}")
    r2 = s.execute_and_verify("def div(a,b): return a/b\ndiv(10,0)")
    print(f"Test 2: success={r2['success']}, error={r2.get('error','')[:50]}")
    print(f"Stats: {s.get_stats()}")