import os
import re
import shutil
import subprocess
import tempfile
import uuid
import resource
import time

MAX_CODE_SIZE = 50_000       # 50 KB max user code
MAX_EXECUTION_TIME = 5        # 5 seconds max execution
JAVA_HEAP_LIMIT = "1024m"     # 1 GB max heap for Java user code

def extract_class_name(code: str) -> str:
    match = re.search(r'public\s+class\s+(\w+)', code)
    return match.group(1) if match else "Main"

def set_limits():
    # Limit only CPU time; no artificial memory limit
    import resource
    resource.setrlimit(resource.RLIMIT_CPU, (MAX_EXECUTION_TIME, MAX_EXECUTION_TIME))

def run_java(code: str):
    import os, tempfile, subprocess, uuid, shutil, time, re

    if len(code.encode("utf-8")) > MAX_CODE_SIZE:
        return {"status": "error", "message": "Code too large"}

    run_id = str(uuid.uuid4())
    workdir = os.path.join(tempfile.gettempdir(), f"runner_{run_id}")
    os.makedirs(workdir, exist_ok=True)

    try:
        class_name = re.search(r'public\s+class\s+(\w+)', code)
        class_name = class_name.group(1) if class_name else "Main"

        java_file = os.path.join(workdir, f"{class_name}.java")
        with open(java_file, "w") as f:
            f.write(code)

        # Compile
        start_compile = time.time()
        compile_proc = subprocess.run(
            ["javac", java_file],
            capture_output=True, text=True, cwd=workdir
        )
        compile_time = (time.time() - start_compile) * 1000
        if compile_proc.returncode != 0:
            return {
                "status": "error",
                "stage": "compile",
                "compile_time_ms": round(compile_time, 2),
                "exit_code": compile_proc.returncode,
                "stdout": compile_proc.stdout,
                "stderr": compile_proc.stderr
            }

        # Run
        start_run = time.time()
        run_proc = subprocess.run(
            ["java", f"-Xmx{JAVA_HEAP_LIMIT}", "-cp", workdir, class_name],
            capture_output=True, text=True,
            cwd=workdir,
            timeout=MAX_EXECUTION_TIME,
            preexec_fn=set_limits
        )
        run_time = (time.time() - start_run) * 1000

        return {
            "status": "success",
            "compile_time_ms": round(compile_time, 2),
            "execution_time_ms": round(run_time, 2),
            "exit_code": run_proc.returncode,
            "stdout": run_proc.stdout,
            "stderr": run_proc.stderr
        }

    except subprocess.TimeoutExpired:
        return {"status": "error", "stage": "run", "message": "Execution timed out"}

    finally:
        shutil.rmtree(workdir, ignore_errors=True)