qihuaiflow/backend/python/app/api/admin/code_exec_admin.py

"""
Code Execution Admin API
Manages code execution sandbox containers, task submission, and result retrieval.

Endpoints:
  GET  /stats        - Execution statistics
  GET  /runtimes     - Available runtime info from sandbox
  POST /execute      - Submit code for execution
  GET  /tasks        - List execution tasks (with filters)
  GET  /tasks/{id}   - Get specific task result
  DELETE /tasks/{id} - Cancel/delete task
"""
import logging
import uuid
import json
import time
from datetime import datetime, timezone, timedelta
from typing import Optional, List

import httpx
from fastapi import APIRouter, Depends, HTTPException, Query
from pydantic import BaseModel, Field
from sqlalchemy.orm import Session
from sqlalchemy import text

from ...database import get_db, SessionLocal
from ...auth import get_current_admin

logger = logging.getLogger(__name__)

router = APIRouter()

CST = timezone(timedelta(hours=8))

# Code execution sandbox configuration
# The executor container is on the same Docker network, accessible by container name
EXECUTOR_URL = "http://qihuaiflow-executor:8000"
EXECUTOR_TIMEOUT = 120  # max execution timeout

# Supported languages
SUPPORTED_LANGUAGES = ["python", "javascript"]


def _execute_raw(db: Session, sql, params=None):
    """Execute raw SQL and return results."""
    result = db.execute(text(sql), params or {})
    if sql.strip().upper().startswith("SELECT") or sql.strip().upper().startswith("WITH"):
        rows = result.fetchall()
        if rows:
            return [dict(r._mapping) for r in rows]
        return []
    db.commit()
    return result


def _format_cst_time(dt):
    """Format datetime to CST string."""
    if dt is None:
        return None
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=CST)
    return dt.strftime("%Y-%m-%d %H:%M:%S")


def _now_cst():
    """Get current time as CST-aware datetime."""
    return datetime.now(CST)


# ===== Pydantic Models =====

class ExecuteRequest(BaseModel):
    language: str = "python"
    code: str
    stdin: str = ""
    timeout: int = Field(default=30, ge=1, le=120)


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


class TaskFilter(BaseModel):
    language: Optional[str] = None
    status: Optional[str] = None
    page: int = 1
    page_size: int = 20


# ===== API Endpoints =====

@router.get("/stats")
async def get_execution_stats(db: Session = Depends(get_db), _=Depends(get_current_admin)):
    """Get code execution statistics."""
    try:
        total = _execute_raw(db, "SELECT COUNT(*) as cnt FROM code_exec_tasks")[0]["cnt"]
        today = _execute_raw(db,
            "SELECT COUNT(*) as cnt FROM code_exec_tasks WHERE DATE(created_at) = CURRENT_DATE")[0]["cnt"]
        success_count = _execute_raw(db,
            "SELECT COUNT(*) as cnt FROM code_exec_tasks WHERE exit_code = 0 AND error IS NULL")[0]["cnt"]
        failed_count = _execute_raw(db,
            "SELECT COUNT(*) as cnt FROM code_exec_tasks WHERE exit_code != 0 OR error IS NOT NULL")[0]["cnt"]
        timeout_count = _execute_raw(db,
            "SELECT COUNT(*) as cnt FROM code_exec_tasks WHERE error LIKE '%timeout%'")[0]["cnt"]
        by_language = _execute_raw(db,
            "SELECT language, COUNT(*) as cnt FROM code_exec_tasks GROUP BY language")

        avg_time = _execute_raw(db,
            "SELECT AVG(execution_time_ms) as avg_ms FROM code_exec_tasks WHERE execution_time_ms > 0")[0]
        total_avg = round(avg_time.get("avg_ms", 0) or 0, 2)

        return {
            "code": 0,
            "data": {
                "total": total,
                "today": today,
                "success": success_count,
                "failed": failed_count,
                "timeout": timeout_count,
                "avg_execution_time_ms": total_avg,
                "by_language": {r["language"]: r["cnt"] for r in by_language} if by_language else {},
            }
        }
    except Exception as e:
        # Table might not exist yet
        return {
            "code": 0,
            "data": {
                "total": 0, "today": 0, "success": 0, "failed": 0, "timeout": 0,
                "avg_execution_time_ms": 0, "by_language": {}
            }
        }


@router.get("/runtimes")
async def get_runtimes(_=Depends(get_current_admin)):
    """Get sandbox runtime info."""
    try:
        async with httpx.AsyncClient(timeout=5) as client:
            resp = await client.get(f"{EXECUTOR_URL}/info")
            if resp.status_code == 200:
                return {"code": 0, "data": resp.json()}
    except Exception as e:
        pass

    return {
        "code": 0,
        "data": {
            "python": "Python 3.11 (sandbox)",
            "nodejs": "Node.js (sandbox)",
            "workspace": "/workspace",
            "max_code_length": 51200,
            "max_timeout": 120,
        }
    }


@router.post("/execute", response_model=ExecuteResponse)
async def execute_code(req: ExecuteRequest, db: Session = Depends(get_db), _=Depends(get_current_admin)):
    """Submit code for execution."""
    if req.language not in SUPPORTED_LANGUAGES:
        raise HTTPException(status_code=400, detail=f"Unsupported language: {req.language}")

    task_id = str(uuid.uuid4())[:8]
    now = _now_cst()

    # Create task record
    _execute_raw(db, """
        INSERT INTO code_exec_tasks (id, language, code_preview, status, timeout_sec, created_at)
        VALUES (:id, :lang, :preview, 'running', :timeout, :now)
    """, {
        "id": task_id,
        "lang": req.language,
        "preview": req.code[:200],
        "timeout": req.timeout,
        "now": now,
    })

    # Execute via sandbox
    result = {
        "success": False,
        "stdout": "",
        "stderr": "",
        "exit_code": -1,
        "execution_time_ms": 0,
        "error": None,
    }

    try:
        async with httpx.AsyncClient(timeout=req.timeout + 10) as client:
            resp = await client.post(
                f"{EXECUTOR_URL}/execute",
                json={
                    "language": req.language,
                    "code": req.code,
                    "stdin": req.stdin,
                    "timeout": req.timeout,
                },
            )
            if resp.status_code == 200:
                data = resp.json()
                result = {
                    "success": data.get("success", False),
                    "stdout": data.get("stdout", ""),
                    "stderr": data.get("stderr", ""),
                    "exit_code": data.get("exit_code", -1),
                    "execution_time_ms": data.get("execution_time_ms", 0),
                    "error": data.get("error"),
                    "truncated": data.get("truncated", False),
                }
            else:
                result["error"] = f"Sandbox returned {resp.status_code}: {resp.text[:200]}"
    except httpx.TimeoutException:
        result["error"] = f"Execution timed out after {req.timeout}s"
    except httpx.ConnectError:
        result["error"] = "Execution sandbox is not available"
    except Exception as e:
        result["error"] = str(e)[:500]

    # Determine status
    if result["error"]:
        status = "failed"
    elif result["exit_code"] != 0:
        status = "error"
    else:
        status = "completed"

    # Update task record
    now2 = _now_cst()
    _execute_raw(db, """
        UPDATE code_exec_tasks SET
            status = :status, exit_code = :exit_code, execution_time_ms = :time_ms,
            stdout = :stdout, stderr = :stderr, error_msg = :error,
            truncated = :truncated, completed_at = :now
        WHERE id = :id
    """, {
        "status": status,
        "exit_code": result["exit_code"],
        "time_ms": result["execution_time_ms"],
        "stdout": result["stdout"][:50000] if result.get("stdout") else "",
        "stderr": result["stderr"][:50000] if result.get("stderr") else "",
        "error": result["error"][:2000] if result.get("error") else None,
        "truncated": 1 if result.get("truncated") else 0,
        "now": now2,
        "id": task_id,
    })

    return ExecuteResponse(
        task_id=task_id,
        success=result["success"],
        stdout=result.get("stdout", ""),
        stderr=result.get("stderr", ""),
        exit_code=result["exit_code"],
        execution_time_ms=result["execution_time_ms"],
        error=result.get("error"),
        truncated=result.get("truncated", False),
        created_at=_format_cst_time(now),
    )


@router.get("/tasks")
async def list_tasks(
    language: Optional[str] = Query(None),
    status: Optional[str] = Query(None),
    page: int = Query(1, ge=1),
    page_size: int = Query(20, ge=1, le=100),
    db: Session = Depends(get_db),
    _=Depends(get_current_admin),
):
    """List execution tasks with filters."""
    where_clauses = []
    params = {}

    if language:
        where_clauses.append("language = :lang")
        params["lang"] = language
    if status:
        where_clauses.append("status = :status")
        params["status"] = status

    where_sql = ""
    if where_clauses:
        where_sql = "WHERE " + " AND ".join(where_clauses)

    # Count total
    count_result = _execute_raw(db, f"SELECT COUNT(*) as cnt FROM code_exec_tasks {where_sql}", params)
    total = count_result[0]["cnt"] if count_result else 0

    # Fetch page
    offset = (page - 1) * page_size
    tasks = _execute_raw(db, f"""
        SELECT id, language, code_preview, status, exit_code, execution_time_ms,
               error_msg, truncated, created_at, completed_at
        FROM code_exec_tasks {where_sql}
        ORDER BY created_at DESC
        LIMIT :limit OFFSET :offset
    """, {**params, "limit": page_size, "offset": offset})

    items = []
    for t in (tasks or []):
        items.append({
            "id": t["id"],
            "language": t["language"],
            "code_preview": t.get("code_preview", ""),
            "status": t["status"],
            "exit_code": t.get("exit_code"),
            "execution_time_ms": t.get("execution_time_ms"),
            "error_msg": t.get("error_msg"),
            "truncated": bool(t.get("truncated", 0)),
            "created_at": _format_cst_time(t.get("created_at")),
            "completed_at": _format_cst_time(t.get("completed_at")),
        })

    return {
        "code": 0,
        "data": {
            "total": total,
            "page": page,
            "page_size": page_size,
            "items": items,
        }
    }


@router.get("/tasks/{task_id}")
async def get_task(task_id: str, db: Session = Depends(get_db), _=Depends(get_current_admin)):
    """Get a specific task with full output."""
    results = _execute_raw(db, """
        SELECT id, language, code_preview, status, exit_code, execution_time_ms,
               stdout, stderr, error_msg, truncated, created_at, completed_at
        FROM code_exec_tasks WHERE id = :id
    """, {"id": task_id})

    if not results:
        raise HTTPException(status_code=404, detail="Task not found")

    t = results[0]
    return {
        "code": 0,
        "data": {
            "id": t["id"],
            "language": t["language"],
            "code_preview": t.get("code_preview", ""),
            "status": t["status"],
            "exit_code": t.get("exit_code"),
            "execution_time_ms": t.get("execution_time_ms"),
            "stdout": t.get("stdout", ""),
            "stderr": t.get("stderr", ""),
            "error_msg": t.get("error_msg"),
            "truncated": bool(t.get("truncated", 0)),
            "created_at": _format_cst_time(t.get("created_at")),
            "completed_at": _format_cst_time(t.get("completed_at")),
        }
    }


@router.delete("/tasks/{task_id}")
async def delete_task(task_id: str, db: Session = Depends(get_db), _=Depends(get_current_admin)):
    """Delete a task record."""
    result = _execute_raw(db, "DELETE FROM code_exec_tasks WHERE id = :id", {"id": task_id})
    return {"code": 0, "msg": "deleted"}