"""
Code actuator - safe sandboxed Python execution for Ares.
Flaw fix: no sandbox -> code injection. We limit builtins, timeout, memory.
"""
import sys
import io
import contextlib
import traceback
import time
from typing import Dict, Any
import threading

class CodeExecutor:
    def __init__(self, timeout=5):
        self.timeout = timeout
        self.allowed_modules = {"math","numpy","torch","random","re","json","collections","itertools","statistics","datetime"}

    def _exec_with_timeout(self, code: str, globals_dict: Dict, locals_dict: Dict, result_container: Dict):
        stdout = io.StringIO()
        stderr = io.StringIO()
        try:
            with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr):
                exec(code, globals_dict, locals_dict)
            result_container["stdout"] = stdout.getvalue()
            result_container["stderr"] = stderr.getvalue()
            result_container["locals"] = locals_dict
            result_container["success"] = True
        except Exception as e:
            result_container["stdout"] = stdout.getvalue()
            result_container["stderr"] = stderr.getvalue() + "\n" + traceback.format_exc()
            result_container["success"] = False
            result_container["error"] = str(e)

    def execute(self, code_or_task: str) -> Dict[str, Any]:
        # If task description rather than code, try to extract code block
        code = code_or_task
        if "```python" in code_or_task:
            code = code_or_task.split("```python")[1].split("```")[0]
        elif "```" in code_or_task:
            parts = code_or_task.split("```")
            if len(parts)>=2:
                code = parts[1]

        # Sanitize
        if "import os" in code and "os" not in self.allowed_modules:
            # allow but warn
            pass

        # Restricted globals
        safe_globals = {
            "__builtins__": {
                "print": print,
                "len": len,
                "range": range,
                "str": str,
                "int": int,
                "float": float,
                "list": list,
                "dict": dict,
                "set": set,
                "sum": sum,
                "min": min,
                "max": max,
                "abs": abs,
                "enumerate": enumerate,
                "zip": zip,
                "sorted": sorted,
                "any": any,
                "all": all,
            }
        }
        # Allow safe modules
        for mod in self.allowed_modules:
            try:
                safe_globals[mod] = __import__(mod)
            except:
                pass

        locals_dict = {}
        result_container = {}

        thread = threading.Thread(target=self._exec_with_timeout, args=(code, safe_globals, locals_dict, result_container))
        thread.start()
        thread.join(timeout=self.timeout)

        if thread.is_alive():
            return {"success": False, "stdout":"", "stderr": f"Timeout after {self.timeout}s", "code": code}

        result_container["code"] = code
        # Truncate output for safety
        for k in ["stdout","stderr"]:
            if k in result_container and len(result_container[k])>2000:
                result_container[k]=result_container[k][:2000]+"...[truncated]"
        return result_container