python -m venv venv
source venv/bin/activate
pip install openai docker
export OPENAI_API_KEY="sk-your-key"

# safety.py
import ast
from typing import List

FORBIDDEN_MODULES = {
    "os", "sys", "subprocess", "shutil", "pathlib",
    "socket", "http", "urllib", "requests",
    "ctypes", "multiprocessing", "threading",
}

FORBIDDEN_FUNCTIONS = {
    "exec", "eval", "compile", "__import__",
    "open", "input", "print",  # controlled separately
}

class CodeSafetyChecker:
    def __init__(self, allowed_modules: set = None):
        self.allowed_modules = allowed_modules or {"math", "json", "re", "datetime", "collections"}

    def check(self, code: str) -> tuple[bool, List[str]]:
        issues = []
        try:
            tree = ast.parse(code)
        except SyntaxError as e:
            return False, [f"Syntax error: {str(e)}"]

        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    if alias.name.split(".")[0] in FORBIDDEN_MODULES:
                        issues.append(f"Forbidden import: {alias.name}")

            elif isinstance(node, ast.ImportFrom):
                if node.module and node.module.split(".")[0] in FORBIDDEN_MODULES:
                    issues.append(f"Forbidden import from: {node.module}")

            elif isinstance(node, ast.Call):
                if isinstance(node.func, ast.Name):
                    if node.func.id in FORBIDDEN_FUNCTIONS:
                        issues.append(f"Forbidden function call: {node.func.id}")

            elif isinstance(node, ast.Attribute):
                if isinstance(node.value, ast.Name):
                    if node.value.id in FORBIDDEN_MODULES:
                        issues.append(f"Forbidden attribute access: {node.value.id}.{node.attr}")

        return len(issues) == 0, issues

    def sanitize_code(self, code: str) -> str:
        tree = ast.parse(code)
        sanitized = []
        for node in ast.iter_child_nodes(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
                sanitized.append(ast.get_source_segment(code, node) or "")
        return "\n\n".join(sanitized)

# sandbox/local_sandbox.py
import subprocess
import tempfile
import os
from typing import Dict

class LocalSandbox:
    def __init__(
        self,
        timeout: int = 30,
        memory_limit_mb: int = 256,
        max_output_bytes: int = 1024 * 1024,
    ):
        self.timeout = timeout
        self.memory_limit_mb = memory_limit_mb
        self.max_output_bytes = max_output_bytes

    def execute(self, code: str) -> Dict:
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".py", delete=False, dir="/tmp"
        ) as f:
            f.write(code)
            temp_path = f.name

        try:
            result = subprocess.run(
                ["python", temp_path],
                capture_output=True,
                text=True,
                timeout=self.timeout,
                env={
                    **os.environ,
                    "PYTHONDONTWRITEBYTECODE": "1",
                },
                cwd=tempfile.gettempdir(),
            )
            return {
                "success": result.returncode == 0,
                "stdout": result.stdout[:self.max_output_bytes],
                "stderr": result.stderr[:self.max_output_bytes],
                "return_code": result.returncode,
                "execution_time_ms": 0,
            }
        except subprocess.TimeoutExpired:
            return {
                "success": False,
                "stdout": "",
                "stderr": f"Execution timed out after {self.timeout}s",
                "return_code": -1,
                "execution_time_ms": self.timeout * 1000,
            }
        finally:
            os.unlink(temp_path)

# sandbox/docker_sandbox.py
import docker
import tempfile
import os
from typing import Dict

class DockerSandbox:
    def __init__(
        self,
        image: str = "python:3.11-slim",
        timeout: int = 30,
        memory_limit: str = "256m",
    ):
        self.client = docker.from_env()
        self.image = image
        self.timeout = timeout
        self.memory_limit = memory_limit

    def execute(self, code: str) -> Dict:
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".py", delete=False
        ) as f:
            f.write(code)
            temp_path = f.name

        try:
            container = self.client.containers.run(
                self.image,
                command=f"python /code/script.py",
                volumes={
                    os.path.dirname(temp_path): {
                        "bind": "/code",
                        "mode": "ro",
                    }
                },
                mem_limit=self.memory_limit,
                network_disabled=True,
                remove=True,
                detach=True,
            )
            result = container.wait(timeout=self.timeout)
            logs = container.logs().decode("utf-8", errors="ignore")
            return {
                "success": result.get("StatusCode", 1) == 0,
                "stdout": logs,
                "stderr": "" if result.get("StatusCode", 1) == 0 else logs,
                "return_code": result.get("StatusCode", 1),
                "execution_time_ms": 0,
            }
        except Exception as e:
            return {
                "success": False,
                "stdout": "",
                "stderr": f"Docker error: {str(e)}",
                "return_code": -1,
                "execution_time_ms": 0,
            }
        finally:
            os.unlink(temp_path)

# generator.py
from openai import OpenAI

CODE_SYSTEM = """You are an expert Python programmer. Generate clean, 
working Python code based on user requirements.

Rules:
1. Write complete, runnable code
2. Include necessary imports
3. Add a main() function or script entry point
4. Handle errors gracefully
5. Use type hints where helpful
6. Print results to stdout for verification"""

class CodeGenerator:
    def __init__(self, model: str = "gpt-4-turbo-preview"):
        self.client = OpenAI()
        self.model = model

    def generate(
        self,
        requirement: str,
        context: str = "",
        previous_code: str = "",
        error: str = "",
    ) -> str:
        prompt = f"Requirement: {requirement}"
        if context:
            prompt += f"\n\nContext: {context}"
        if previous_code:
            prompt += f"\n\nPrevious code that failed:\n```python\n{previous_code}\n```\n\nError:\n{error}\n\nFix the code:"
        else:
            prompt += "\n\nGenerate the code:"

        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": CODE_SYSTEM},
                {"role": "user", "content": prompt},
            ],
            temperature=0.0,
        )
        return self._extract_code(response.choices[0].message.content)

    def _extract_code(self, content: str) -> str:
        if "```python" in content:
            parts = content.split("```python")
            if len(parts) > 1:
                code = parts[1].split("```")[0]
                return code.strip()
        return content.strip()

# debugger.py
from openai import OpenAI
from typing import Tuple

DEBUG_SYSTEM = """You are a debugging expert. Analyze the error and provide 
a fix. Return ONLY the corrected Python code, no explanation."""

class CodeDebugger:
    def __init__(self, model: str = "gpt-4-turbo-preview"):
        self.client = OpenAI()
        self.model = model

    def debug(self, code: str, error: str, stdout: str = "") -> str:
        prompt = f"""Code:
```python
{code}

# Test generation
class TestGenerator:
    def __init__(self, model: str = "gpt-4-turbo-preview"):
        from openai import OpenAI
        self.client = OpenAI()
        self.model = model

    def generate_tests(self, code: str, requirement: str) -> str:
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "Generate pytest tests for the given code."},
                {"role": "user", "content": f"Code:\n{code}\n\nRequirement: {requirement}\n\nGenerate tests:"},
            ],
            temperature=0.0,
        )
        return response.choices[0].message.content

import pytest
from code_agent import CodeAgent

@pytest.fixture
def agent():
    return CodeAgent(max_retries=3)

def test_simple_code(agent):
    result = agent.run("Write a function that calculates fibonacci numbers")
    assert result["success"]
    assert "fibonacci" in result["code"].lower() or "def " in result["code"]

def test_safety_check():
    from safety import CodeSafetyChecker
    checker = CodeSafetyChecker()
    safe, issues = checker.check("import math\nprint(math.sqrt(4))")
    assert safe
    safe, issues = checker.check("import os\nos.system('rm -rf /')")
    assert not safe

def test_debugging():
    from debugger import CodeDebugger
    debugger = CodeDebugger()
    fixed = debugger.debug(
        "def add(a, b):\n    return a + b",
        "TypeError: unsupported operand",
    )
    assert "def " in fixed

# main.py
from code_agent import CodeAgent

def main():
    agent = CodeAgent(model="gpt-4-turbo-preview", use_docker=True)
    print("Code Agent Ready. Type 'quit' to exit.\n")

    while True:
        req = input("Describe what you want to build: ").strip()
        if req.lower() == "quit":
            break
        result = agent.run(req)
        print(f"\n{'SUCCESS' if result['success'] else 'FAILED'}")
        print(f"Attempts: {result['attempts']}")
        print(f"\nCode:\n{result['code']}")
        if result["success"]:
            print(f"\nOutput:\n{result['output']}")
        else:
            print(f"\nError: {result['error']}")

if __name__ == "__main__":
    main()