yui_code_runner.py

"""
yui_code_runner.py — YUI CODE EXECUTION ENGINE v1.0
=====================================================
Hỗ trợ chạy code qua nhiều backend theo thứ tự ưu tiên:

  1. PISTON API   — Không cần key, 5 req/s, ~30+ ngôn ngữ (ưu tiên nhất)
  2. JUDGE0 API   — Cần key (RapidAPI free: 50/ngày), chính xác hơn
  3. LOCAL SANDBOX — Python/JS/Bash local (fallback cuối)

Rate limiting:
  - Guest/friend: 3 lần/phút, 20 lần/giờ
  - Admin (Bạch): không giới hạn

Yui tự dùng khi detect code trong reply cần verify,
hoặc user gọi /chaycode, hoặc AI tự test code mình vừa viết.
"""

import os, time, random, threading, asyncio
from typing import Optional

# ── RATE LIMITER ──────────────────────────────────────────────────
_rate_lock = threading.Lock()
_rate_data: dict = {}  # {user_id: {"minute": [ts...], "hour": [ts...]}}

RATE_LIMIT_MINUTE = 3   # requests / phút cho guest
RATE_LIMIT_HOUR   = 20  # requests / giờ cho guest

def _check_rate(user_id: str, is_admin: bool = False) -> tuple[bool, str]:
    """
    Kiểm tra rate limit.
    Returns: (allowed: bool, reason: str)
    """
    if is_admin:
        return True, ""
    now = time.time()
    with _rate_lock:
        d = _rate_data.setdefault(user_id, {"minute": [], "hour": []})
        # Dọn timestamps cũ
        d["minute"] = [t for t in d["minute"] if now - t < 60]
        d["hour"]   = [t for t in d["hour"]   if now - t < 3600]
        if len(d["minute"]) >= RATE_LIMIT_MINUTE:
            wait = int(60 - (now - d["minute"][0])) + 1
            return False, f"⏳ Quá {RATE_LIMIT_MINUTE} lần/phút — chờ {wait}s nhé!"
        if len(d["hour"]) >= RATE_LIMIT_HOUR:
            wait = int(3600 - (now - d["hour"][0])) // 60 + 1
            return False, f"⏳ Quá {RATE_LIMIT_HOUR} lần/giờ — chờ {wait} phút nhé!"
        d["minute"].append(now)
        d["hour"].append(now)
    return True, ""


# ── LANGUAGE MAPPING ──────────────────────────────────────────────
# Piston language slugs + versions
PISTON_LANGS = {
    "python":     ("python",     "3.10.0"),
    "py":         ("python",     "3.10.0"),
    "python3":    ("python",     "3.10.0"),
    "javascript": ("javascript", "18.15.0"),
    "js":         ("javascript", "18.15.0"),
    "node":       ("javascript", "18.15.0"),
    "typescript": ("typescript", "5.0.3"),
    "ts":         ("typescript", "5.0.3"),
    "java":       ("java",       "15.0.2"),
    "c":          ("c",          "10.2.0"),
    "cpp":        ("c++",        "10.2.0"),
    "c++":        ("c++",        "10.2.0"),
    "rust":       ("rust",       "1.50.0"),
    "go":         ("go",         "1.16.2"),
    "ruby":       ("ruby",       "3.0.1"),
    "php":        ("php",        "8.2.3"),
    "bash":       ("bash",       "5.2.0"),
    "sh":         ("bash",       "5.2.0"),
    "lua":        ("lua",        "5.4.4"),
    "kotlin":     ("kotlin",     "1.8.20"),
    "swift":      ("swift",      "5.8.0"),
    "r":          ("r",          "4.1.1"),
    "csharp":     ("csharp",     "6.12.0"),
    "cs":         ("csharp",     "6.12.0"),
    "perl":       ("perl",       "5.36.0"),
}

# Judge0 language IDs (RapidAPI)
JUDGE0_LANG_IDS = {
    "python":     71,   # Python 3.8.1
    "py":         71,
    "python3":    71,
    "javascript": 63,   # JavaScript (Node.js 12.14.0)
    "js":         63,
    "typescript": 74,   # TypeScript 3.7.4
    "ts":         74,
    "java":       62,   # Java (OpenJDK 13.0.1)
    "c":          50,   # C (GCC 9.2.0)
    "cpp":        54,   # C++ (GCC 9.2.0)
    "c++":        54,
    "rust":       73,   # Rust 1.40.0
    "go":         60,   # Go 1.13.5
    "ruby":       72,   # Ruby 2.7.0
    "php":        68,   # PHP 7.4.1
    "bash":       46,   # Bash 5.0.0
    "csharp":     51,   # C# Mono 6.6.0.161
    "cs":         51,
    "r":          80,   # R 4.0.0
    "kotlin":     78,   # Kotlin 1.3.70
    "swift":      83,   # Swift 5.2.3
}


# ── SECURITY FILTER ───────────────────────────────────────────────
_BLOCKED_PATTERNS = [
    "rm -rf", "del /f", "format c:", "shutdown", "reboot",
    "import os; os.system", "__import__('os')",
    "subprocess.call", "os.popen", "sys.exit",
    "socket.connect", "urllib.request.urlopen",
    "open('/etc/passwd')", "open('/etc/shadow')",
    ":(){:|:&};:", "fork bomb",
    "nc -e", "netcat", "mkfifo",
]

def _security_check(code: str) -> tuple[bool, str]:
    """Kiểm tra code có chứa pattern nguy hiểm không."""
    code_lower = code.lower()
    for pattern in _BLOCKED_PATTERNS:
        if pattern.lower() in code_lower:
            return False, f"🚫 Code chứa lệnh nguy hiểm (`{pattern}`) — bị chặn!"
    # Giới hạn độ dài
    if len(code) > 5000:
        return False, "🚫 Code quá dài (tối đa 5000 ký tự)"
    return True, ""


# ── PISTON API ────────────────────────────────────────────────────
PISTON_URL = "https://emkc.org/api/v2/piston/execute"
_piston_last_req = [0.0]
_piston_lock = threading.Lock()

async def _run_piston(code: str, lang: str, stdin: str = "") -> dict:
    """
    Gọi Piston API — không cần key, 5 req/s rate limit.
    Returns: {"ok": bool, "output": str, "error": str, "time": float}
    """
    lang_info = PISTON_LANGS.get(lang.lower())
    if not lang_info:
        return {"ok": False, "output": "", "error": f"Piston không hỗ trợ ngôn ngữ: {lang}"}

    piston_lang, piston_ver = lang_info

    # Rate limit 5 req/s — dùng sleep nhẹ nếu cần
    with _piston_lock:
        now = time.time()
        gap = now - _piston_last_req[0]
        if gap < 0.22:   # ~4.5 req/s để an toàn
            await asyncio.sleep(0.22 - gap)
        _piston_last_req[0] = time.time()

    payload = {
        "language": piston_lang,
        "version":  piston_ver,
        "files":    [{"name": "main", "content": code}],
        "stdin":    stdin,
        "args":     [],
        "compile_timeout": 10000,
        "run_timeout":     10000,
        "compile_memory_limit": -1,
        "run_memory_limit":     -1,
    }
    try:
        import httpx
        t0 = time.time()
        async with httpx.AsyncClient(timeout=20) as client:
            resp = await client.post(PISTON_URL, json=payload)
        elapsed = round(time.time() - t0, 2)

        if resp.status_code == 429:
            return {"ok": False, "output": "", "error": "Piston rate limit (5/s) — thử lại sau 1s", "time": elapsed}
        if resp.status_code != 200:
            return {"ok": False, "output": "", "error": f"Piston HTTP {resp.status_code}", "time": elapsed}

        data = resp.json()
        run  = data.get("run", {})
        output = (run.get("stdout") or "").strip()
        stderr = (run.get("stderr") or "").strip()
        code_ret = run.get("code", 0)

        # Compile stage (C, Java, etc.)
        compile_info = data.get("compile", {})
        if compile_info:
            cerr = (compile_info.get("stderr") or "").strip()
            if cerr and code_ret != 0:
                return {"ok": False, "output": output, "error": f"Compile error:\n{cerr[:800]}", "time": elapsed}

        if code_ret != 0 and not output:
            return {"ok": False, "output": "", "error": stderr[:800] or f"Exit code {code_ret}", "time": elapsed}

        return {
            "ok": True,
            "output": output[:1500],
            "error": stderr[:200] if stderr else "",
            "time": elapsed,
            "lang": f"{piston_lang} {piston_ver}",
        }
    except Exception as e:
        return {"ok": False, "output": "", "error": f"Piston exception: {str(e)[:100]}"}


# ── JUDGE0 API ────────────────────────────────────────────────────
JUDGE0_URL      = "https://judge0-ce.p.rapidapi.com"
JUDGE0_ALT_URL  = "https://judge0-extra-ce.p.rapidapi.com"  # Extra languages
JUDGE0_HOST     = "judge0-ce.p.rapidapi.com"
_judge0_daily   = {}   # {key_prefix: {"count": int, "reset_ts": float}}

def _get_judge0_key() -> Optional[str]:
    """Lấy Judge0 key còn hạn ngạch từ env."""
    keys = []
    for k, v in os.environ.items():
        if k.startswith("JUDGE0_KEY") or k.startswith("RAPIDAPI_KEY"):
            if v.strip():
                keys.append(v.strip())
    if not keys:
        return None
    now = time.time()
    for key in keys:
        prefix = key[:8]
        d = _judge0_daily.get(prefix, {"count": 0, "reset_ts": now + 86400})
        if now > d["reset_ts"]:
            _judge0_daily[prefix] = {"count": 0, "reset_ts": now + 86400}
            return key
        if d["count"] < 45:   # Buffer 5 để tránh hit giới hạn đột ngột
            return key
    return None  # Tất cả key hết quota hôm nay

def _mark_judge0_used(key: str):
    prefix = key[:8]
    d = _judge0_daily.get(prefix, {"count": 0, "reset_ts": time.time() + 86400})
    d["count"] += 1
    _judge0_daily[prefix] = d

async def _run_judge0(code: str, lang: str, stdin: str = "") -> dict:
    """
    Gọi Judge0 API — cần JUDGE0_KEY (RapidAPI free: 50/ngày).
    Chính xác hơn Piston, hỗ trợ nhiều ngôn ngữ hơn.
    """
    key = _get_judge0_key()
    if not key:
        return {"ok": False, "output": "", "error": "Judge0: Hết quota hôm nay (50/ngày) — reset 00:00 UTC"}

    lang_id = JUDGE0_LANG_IDS.get(lang.lower())
    if not lang_id:
        return {"ok": False, "output": "", "error": f"Judge0 không hỗ trợ: {lang}"}

    import base64, httpx
    payload = {
        "source_code":   base64.b64encode(code.encode()).decode(),
        "language_id":   lang_id,
        "stdin":         base64.b64encode(stdin.encode()).decode() if stdin else "",
        "cpu_time_limit": 5,
        "memory_limit":   128000,
    }
    headers = {
        "X-RapidAPI-Key":  key,
        "X-RapidAPI-Host": JUDGE0_HOST,
        "Content-Type":    "application/json",
    }
    try:
        t0 = time.time()
        async with httpx.AsyncClient(timeout=30) as client:
            # Submit
            sub = await client.post(
                f"{JUDGE0_URL}/submissions?base64_encoded=true&wait=true",
                json=payload, headers=headers
            )
        elapsed = round(time.time() - t0, 2)

        if sub.status_code == 429:
            return {"ok": False, "output": "", "error": "Judge0: 429 rate limit — thử lại sau"}
        if sub.status_code not in (200, 201):
            return {"ok": False, "output": "", "error": f"Judge0 HTTP {sub.status_code}"}

        _mark_judge0_used(key)
        result = sub.json()
        stdout  = base64.b64decode(result.get("stdout")  or "").decode(errors="replace").strip()
        stderr  = base64.b64decode(result.get("stderr")  or "").decode(errors="replace").strip()
        compile_out = base64.b64decode(result.get("compile_output") or "").decode(errors="replace").strip()
        status  = result.get("status", {}).get("description", "")
        time_ms = result.get("time", "?")

        if status == "Accepted":
            return {"ok": True, "output": stdout[:1500], "error": "", "time": elapsed, "exec_time": time_ms}
        elif compile_out:
            return {"ok": False, "output": "", "error": f"Compile error:\n{compile_out[:800]}", "time": elapsed}
        else:
            err = stderr or f"Status: {status}"
            return {"ok": False, "output": stdout[:500], "error": err[:500], "time": elapsed}
    except Exception as e:
        return {"ok": False, "output": "", "error": f"Judge0 exception: {str(e)[:100]}"}


# ── LOCAL SANDBOX ─────────────────────────────────────────────────
async def _run_local(code: str, lang: str) -> dict:
    """Fallback: chạy local (chỉ Python, JS, Bash trên HF Space)."""
    import subprocess, tempfile, os as _os
    lang_map = {
        "python": ("python3", ".py"), "py": ("python3", ".py"),
        "python3": ("python3", ".py"),
        "javascript": ("node", ".js"), "js": ("node", ".js"),
        "bash": ("bash", ".sh"), "sh": ("bash", ".sh"),
    }
    if lang.lower() not in lang_map:
        return {"ok": False, "output": "", "error": f"Local sandbox chỉ hỗ trợ Python/JS/Bash"}
    runner, ext = lang_map[lang.lower()]
    t0 = time.time()
    with tempfile.NamedTemporaryFile(mode="w", suffix=ext, delete=False, encoding="utf-8") as f:
        f.write(code); tmp = f.name
    try:
        res = subprocess.run(
            [runner, tmp], capture_output=True, text=True, timeout=10,
            env={"PATH": "/usr/bin:/bin:/usr/local/bin", "HOME": "/tmp"}
        )
        elapsed = round(time.time() - t0, 2)
        out = res.stdout.strip()[:1500]
        err = res.stderr.strip()[:400]
        if res.returncode == 0:
            return {"ok": True, "output": out, "error": err, "time": elapsed, "source": "local"}
        return {"ok": False, "output": out, "error": err or f"Exit {res.returncode}", "time": elapsed, "source": "local"}
    except subprocess.TimeoutExpired:
        return {"ok": False, "output": "", "error": "Timeout (>10s)", "time": 10.0, "source": "local"}
    except Exception as e:
        return {"ok": False, "output": "", "error": str(e)[:80], "source": "local"}
    finally:
        try: _os.unlink(tmp)
        except: pass


# ── MAIN ENTRY POINT ──────────────────────────────────────────────
async def run_code(
    code:     str,
    lang:     str     = "python",
    stdin:    str     = "",
    user_id:  str     = "unknown",
    is_admin: bool    = False,
    source:   str     = "manual",   # "manual"|"auto" (Yui tự gọi)
) -> dict:
    """
    Chạy code với fallback chain: Piston → Judge0 → Local
    
    Returns:
      {
        "ok": bool,
        "output": str,
        "error": str,
        "time": float,
        "backend": str,   # "piston"|"judge0"|"local"
        "rate_blocked": bool,
        "rate_msg": str,
      }
    """
    # Rate limit (auto-calls từ Yui không bị limit khi is_admin)
    if source == "manual":
        allowed, rate_msg = _check_rate(user_id, is_admin)
        if not allowed:
            return {"ok": False, "output": "", "error": rate_msg,
                    "rate_blocked": True, "rate_msg": rate_msg, "backend": "none"}

    # Security check
    safe, reason = _security_check(code)
    if not safe:
        return {"ok": False, "output": "", "error": reason, "rate_blocked": False, "backend": "none"}

    lang = lang.lower().strip()

    # 1. Thử Piston (không cần key)
    result = await _run_piston(code, lang, stdin)
    if result["ok"]:
        result["backend"] = "piston"
        return result
    piston_err = result.get("error", "")
    print(f"[CODE-RUN] Piston fail ({piston_err[:50]}) → Judge0")

    # 2. Thử Judge0 (cần key, 50/ngày)
    if _get_judge0_key():
        result = await _run_judge0(code, lang, stdin)
        if result["ok"]:
            result["backend"] = "judge0"
            return result
        judge0_err = result.get("error", "")
        print(f"[CODE-RUN] Judge0 fail ({judge0_err[:50]}) → local")
    else:
        judge0_err = "Không có Judge0 key"

    # 3. Local sandbox (fallback cuối)
    result = await _run_local(code, lang)
    result["backend"] = result.get("source", "local")
    return result


def format_run_result(result: dict, lang: str = "", code_preview: bool = True) -> str:
    """Format kết quả chạy code thành string đẹp cho Discord."""
    if result.get("rate_blocked"):
        return result["rate_msg"]

    backend  = result.get("backend", "?")
    elapsed  = result.get("time", "?")
    exec_t   = result.get("exec_time", "")
    ok       = result.get("ok", False)
    output   = result.get("output", "").strip()
    error    = result.get("error", "").strip()

    icon = "✅" if ok else "❌"
    time_str = f"{elapsed}s" if elapsed != "?" else ""
    if exec_t:
        time_str += f" (exec: {exec_t}s)"

    lines = [f"{icon} **[{backend.upper()}]** {lang} • {time_str}"]

    if ok and output:
        if len(output) > 1000:
            output = output[:1000] + "\n... (truncated)"
        lines.append(f"```\n{output}\n```")
    elif ok and not output:
        lines.append("*(Không có output)*")
    else:
        lines.append(f"```\n{error or 'Lỗi không xác định'}\n```")
        if output:
            lines.append(f"stdout: `{output[:200]}`")

    return "\n".join(lines)


# ── YUI AUTO CODE TEST ────────────────────────────────────────────
import re as _re_coderun

_CODE_BLOCK_RE = _re_coderun.compile(
    r"```(?P<lang>\w+)?\n(?P<code>.*?)```",
    _re_coderun.DOTALL
)
_AUTO_TEST_LANGS = {"python", "py", "javascript", "js", "typescript", "ts"}
_AUTO_TEST_TRIGGERS = [
    "def ", "function ", "const ", "let ", "var ",
    "class ", "import ", "from ", "print(", "console.log(",
    "return ", "async def ", "async function",
]

def should_auto_test(reply: str, task_type: str) -> Optional[tuple[str, str]]:
    """
    Kiểm tra xem reply của Yui có code block đáng test không.
    Returns: (code, lang) hoặc None
    """
    if task_type not in ("code", "debug", "math"):
        return None
    match = _CODE_BLOCK_RE.search(reply)
    if not match:
        return None
    lang = (match.group("lang") or "python").lower().strip()
    code = match.group("code").strip()
    if lang not in _AUTO_TEST_LANGS:
        return None
    if len(code) < 10:
        return None
    if not any(t in code for t in _AUTO_TEST_TRIGGERS):
        return None
    return code, lang


async def auto_test_reply(reply: str, task_type: str, timeout: int = 8) -> Optional[str]:
    """
    Yui tự test code trong reply của mình.
    Nếu có lỗi → báo ngắn gọn để Yui tự sửa.
    Returns: None (OK) hoặc error string
    """
    info = should_auto_test(reply, task_type)
    if not info:
        return None
    code, lang = info
    try:
        result = await asyncio.wait_for(
            run_code(code, lang, user_id="yui_auto", is_admin=True, source="auto"),
            timeout=timeout
        )
        if not result["ok"]:
            err = result.get("error", "unknown error")
            return f"[AUTO-TEST] Code lỗi: {err[:200]}"
        return None   # OK
    except asyncio.TimeoutError:
        return "[AUTO-TEST] Timeout khi test code"
    except Exception as e:
        return f"[AUTO-TEST] Exception: {str(e)[:80]}"