"""
Code Executor Pro — ultimate mashup server
FastAPI, session persistence, JSON output, multi-language, file workspace
"""
from fastapi import FastAPI, HTTPException, UploadFile, File, Form
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
import uvicorn
import sys
import io
import traceback
import subprocess
import os
import json
import tempfile
import time
import uuid
from pathlib import Path

app = FastAPI(title="Code Executor Pro", version="1.0.0")

WORKSPACE_ROOT = Path("/workspace")
WORKSPACE_ROOT.mkdir(exist_ok=True)

# ── Session Manager ──────────────────────────────────────────────────────────
class Session:
    def __init__(self):
        self.vars = {}
        self.created_at = time.time()
        self.last_access = time.time()

sessions: dict[str, Session] = {}
SESSION_TIMEOUT = 3600  # 1 hour

def get_or_create_session(session_id: str | None) -> tuple[str, Session]:
    if not session_id:
        session_id = uuid.uuid4().hex[:12]
    if session_id not in sessions:
        sessions[session_id] = Session()
    sess = sessions[session_id]
    sess.last_access = time.time()
    # Clean stale sessions
    now = time.time()
    stale = [k for k, v in sessions.items() if now - v.last_access > SESSION_TIMEOUT]
    for k in stale:
        del sessions[k]
    return session_id, sess

# ── Request Models ───────────────────────────────────────────────────────────
class RunRequest(BaseModel):
    code: str = Field(..., description="Code to execute")
    language: str = Field("python", description="python | javascript | bash")
    session_id: str | None = Field(None, description="Session ID for state persistence")
    timeout: int = Field(30, description="Max execution time in seconds")

class RunResponse(BaseModel):
    success: bool
    stdout: str
    stderr: str
    error: str | None = None
    execution_time_ms: int
    session_id: str
    files_created: list[str] = []

# ── Executors ────────────────────────────────────────────────────────────────
def execute_python(code: str, session_vars: dict, workspace: Path, timeout: int):
    stdout_capture = io.StringIO()
    stderr_capture = io.StringIO()
    old_stdout, old_stderr = sys.stdout, sys.stderr

    try:
        sys.stdout = stdout_capture
        sys.stderr = stderr_capture

        exec_globals = {
            '__builtins__': __builtins__,
            'WORKSPACE': str(workspace),
        }
        exec_globals.update(session_vars)

        exec(code, exec_globals)

        session_vars.clear()
        session_vars.update({
            k: v for k, v in exec_globals.items()
            if not k.startswith('__') and k not in ['WORKSPACE']
        })

        return stdout_capture.getvalue(), stderr_capture.getvalue(), None
    except Exception as e:
        return stdout_capture.getvalue(), stderr_capture.getvalue(), f"{type(e).__name__}: {str(e)}\n{traceback.format_exc()}"
    finally:
        sys.stdout, sys.stderr = old_stdout, old_stderr

def execute_javascript(code: str, timeout: int):
    try:
        result = subprocess.run(
            ["node", "-e", code],
            capture_output=True, text=True, timeout=timeout
        )
        return result.stdout, result.stderr, None if result.returncode == 0 else f"exit code {result.returncode}"
    except subprocess.TimeoutExpired:
        return "", "", f"Timeout after {timeout}s"
    except FileNotFoundError:
        return "", "", "Node.js not installed"
    except Exception as e:
        return "", "", str(e)

def execute_bash(code: str, timeout: int):
    try:
        result = subprocess.run(
            ["bash", "-c", code],
            capture_output=True, text=True, timeout=timeout
        )
        return result.stdout, result.stderr, None if result.returncode == 0 else f"exit code {result.returncode}"
    except subprocess.TimeoutExpired:
        return "", "", f"Timeout after {timeout}s"
    except Exception as e:
        return "", "", str(e)

# ── API Routes ───────────────────────────────────────────────────────────────
@app.get("/health")
def health():
    return {"status": "ok", "sessions_active": len(sessions)}

@app.post("/run", response_model=RunResponse)
def run_code(req: RunRequest):
    session_id, session = get_or_create_session(req.session_id)
    workspace = WORKSPACE_ROOT / session_id
    workspace.mkdir(exist_ok=True)

    start = time.time()
    stdout, stderr, error = "", "", None

    try:
        if req.language == "python":
            stdout, stderr, error = execute_python(req.code, session.vars, workspace, req.timeout)
        elif req.language == "javascript":
            stdout, stderr, error = execute_javascript(req.code, req.timeout)
        elif req.language == "bash":
            stdout, stderr, error = execute_bash(req.code, req.timeout)
        else:
            raise HTTPException(400, f"Unsupported language: {req.language}")
    except Exception as e:
        error = f"{type(e).__name__}: {str(e)}"
        stdout = ""
        stderr = ""

    elapsed_ms = int((time.time() - start) * 1000)

    # List files created in workspace
    files = [str(f.name) for f in workspace.iterdir() if f.is_file()]

    return RunResponse(
        success=error is None,
        stdout=stdout,
        stderr=stderr,
        error=error,
        execution_time_ms=elapsed_ms,
        session_id=session_id,
        files_created=files,
    )

@app.post("/session/reset")
def reset_session(session_id: str = Form(...)):
    if session_id in sessions:
        del sessions[session_id]
    return {"status": "reset", "session_id": session_id}

@app.get("/session/{session_id}/vars")
def list_session_vars(session_id: str):
    if session_id not in sessions:
        raise HTTPException(404, "Session not found")
    return {"session_id": session_id, "vars": list(sessions[session_id].vars.keys())}

@app.post("/files/upload/{session_id}")
async def upload_file(session_id: str, file: UploadFile = File(...)):
    workspace = WORKSPACE_ROOT / session_id
    workspace.mkdir(exist_ok=True)
    dest = workspace / file.filename
    content = await file.read()
    dest.write_bytes(content)
    return {"status": "ok", "file": file.filename, "size": len(content)}

@app.get("/files/{session_id}")
def list_files(session_id: str):
    workspace = WORKSPACE_ROOT / session_id
    if not workspace.exists():
        return {"session_id": session_id, "files": []}
    files = []
    for f in workspace.iterdir():
        if f.is_file():
            files.append({"name": f.name, "size": f.stat().st_size})
    return {"session_id": session_id, "files": files}

@app.get("/")
def root():
    return {
        "service": "Code Executor Pro",
        "version": "1.0.0",
        "languages": ["python", "javascript", "bash"],
        "endpoints": {
            "POST /run": "Execute code",
            "POST /session/reset": "Reset session",
            "GET /session/{id}/vars": "List session variables",
            "POST /files/upload/{id}": "Upload file",
            "GET /files/{id}": "List workspace files",
            "GET /health": "Health check",
        }
    }

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)