"""
Code Execution Sandbox Server
Runs inside a Docker container, accepts code execution requests.
Supports: Python, JavaScript
"""
import asyncio
import subprocess
import tempfile
import os
import time
import resource
import json
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI(title="Code Execution Sandbox", version="1.0.0")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

# Use /tmp for temp files (writable by executor user)
# /workspace is a tmpfs mount owned by root
WORKSPACE = Path("/tmp")
WORKSPACE.mkdir(exist_ok=True)

# Max code length: 50KB
MAX_CODE_LENGTH = 50 * 1024
# Default timeout: 30s, max: 120s
DEFAULT_TIMEOUT = 30
MAX_TIMEOUT = 120


class ExecuteRequest(BaseModel):
    language: str = Field(..., description="Programming language: python, javascript")
    code: str = Field(..., description="Source code to execute")
    stdin: str = Field(default="", description="Standard input for the program")
    timeout: int = Field(default=DEFAULT_TIMEOUT, ge=1, le=MAX_TIMEOUT, description="Execution timeout in seconds")


class ExecuteResponse(BaseModel):
    success: bool
    stdout: str = ""
    stderr: str = ""
    exit_code: int = 0
    execution_time_ms: float = 0
    error: Optional[str] = None
    truncated: bool = False


@app.get("/health")
async def health():
    return {"status": "healthy", "languages": ["python", "javascript"]}


@app.get("/info")
async def info():
    """Return sandbox runtime info."""
    py_ver = subprocess.run(["python3", "--version"], capture_output=True, text=True).stdout.strip()
    try:
        node_ver = subprocess.run(["node", "--version"], capture_output=True, text=True).stdout.strip()
    except Exception:
        node_ver = "not installed"
    return {
        "python": py_ver,
        "nodejs": node_ver,
        "workspace": str(WORKSPACE),
        "max_code_length": MAX_CODE_LENGTH,
        "max_timeout": MAX_TIMEOUT,
    }


@app.post("/execute", response_model=ExecuteResponse)
async def execute_code(req: ExecuteRequest):
    """Execute code in a sandboxed subprocess."""
    language = req.language.lower().strip()
    code = req.code
    stdin_input = req.stdin
    timeout = min(req.timeout, MAX_TIMEOUT)

    # Validate language
    if language not in ("python", "py", "python3", "javascript", "js", "node"):
        raise HTTPException(status_code=400, detail=f"Unsupported language: {language}. Supported: python, javascript")

    # Normalize language
    if language in ("py", "python3"):
        language = "python"
    if language in ("js", "node"):
        language = "javascript"

    # Validate code length
    if len(code.encode("utf-8")) > MAX_CODE_LENGTH:
        raise HTTPException(status_code=400, detail=f"Code too long. Max {MAX_CODE_LENGTH // 1024}KB")

    start_time = time.time()

    try:
        if language == "python":
            result = await _run_python(code, stdin_input, timeout)
        elif language == "javascript":
            result = await _run_javascript(code, stdin_input, timeout)
        else:
            raise HTTPException(status_code=400, detail=f"Unsupported language: {language}")
    except asyncio.TimeoutError:
        execution_time = (time.time() - start_time) * 1000
        return ExecuteResponse(
            success=False,
            stdout="",
            stderr="",
            exit_code=-1,
            execution_time_ms=execution_time,
            error=f"Execution timed out after {timeout}s",
        )

    execution_time_ms = (time.time() - start_time) * 1000
    result["execution_time_ms"] = round(execution_time_ms, 2)
    return result


async def _run_python(code: str, stdin_input: str, timeout: int) -> dict:
    """Execute Python code."""
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".py", dir=WORKSPACE, delete=False, encoding="utf-8"
    ) as f:
        f.write(code)
        script_path = f.name

    try:
        # Use resource limits
        def set_limits():
            # Memory limit: 256MB
            resource.setrlimit(resource.RLIMIT_AS, (256 * 1024 * 1024, 256 * 1024 * 1024))
            # CPU time limit: timeout + 5s grace
            resource.setrlimit(resource.RLIMIT_CPU, (timeout + 5, timeout + 5))

        proc = await asyncio.create_subprocess_exec(
            "python3",
            "-u",  # unbuffered
            script_path,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            preexec_fn=set_limits,
        )

        try:
            stdout, stderr = await asyncio.wait_for(
                proc.communicate(input=stdin_input.encode("utf-8") if stdin_input else None),
                timeout=timeout,
            )
        except asyncio.TimeoutError:
            proc.kill()
            await proc.wait()
            raise

        stdout_str = stdout.decode("utf-8", errors="replace")
        stderr_str = stderr.decode("utf-8", errors="replace")

        # Truncate output to 500KB
        truncated = False
        if len(stdout_str.encode("utf-8")) > 500 * 1024:
            stdout_str = stdout_str[:500 * 1024]
            truncated = True
        if len(stderr_str.encode("utf-8")) > 500 * 1024:
            stderr_str = stderr_str[:500 * 1024]
            truncated = True

        return {
            "success": proc.returncode == 0,
            "stdout": stdout_str,
            "stderr": stderr_str,
            "exit_code": proc.returncode or 0,
            "error": None,
            "truncated": truncated,
        }

    finally:
        # Cleanup temp file
        try:
            os.unlink(script_path)
        except OSError:
            pass


async def _run_javascript(code: str, stdin_input: str, timeout: int) -> dict:
    """Execute JavaScript code via Node.js."""
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".js", dir=WORKSPACE, delete=False, encoding="utf-8"
    ) as f:
        f.write(code)
        script_path = f.name

    try:
        def set_limits():
            # Only limit CPU time for JS — V8 needs large virtual address space
            # for CodeRange; rely on Docker --memory for actual RSS limit.
            resource.setrlimit(resource.RLIMIT_CPU, (timeout + 5, timeout + 5))

        # Use Node.js to run the script with tight heap limits
        env = os.environ.copy()
        node_flags = "--max-old-space-size=32 --max-semi-space-size=8"
        env["NODE_OPTIONS"] = node_flags

        proc = await asyncio.create_subprocess_exec(
            "node",
            "--max-old-space-size=32",
            "--max-semi-space-size=8",
            script_path,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=env,
            preexec_fn=set_limits,
        )

        try:
            stdout, stderr = await asyncio.wait_for(
                proc.communicate(input=stdin_input.encode("utf-8") if stdin_input else None),
                timeout=timeout,
            )
        except asyncio.TimeoutError:
            proc.kill()
            await proc.wait()
            raise

        stdout_str = stdout.decode("utf-8", errors="replace")
        stderr_str = stderr.decode("utf-8", errors="replace")

        truncated = False
        if len(stdout_str.encode("utf-8")) > 500 * 1024:
            stdout_str = stdout_str[:500 * 1024]
            truncated = True
        if len(stderr_str.encode("utf-8")) > 500 * 1024:
            stderr_str = stderr_str[:500 * 1024]
            truncated = True

        return {
            "success": proc.returncode == 0,
            "stdout": stdout_str,
            "stderr": stderr_str,
            "exit_code": proc.returncode or 0,
            "error": None,
            "truncated": truncated,
        }

    finally:
        try:
            os.unlink(script_path)
        except OSError:
            pass


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)