# /// script
# requires-python = ">=3.11"
# dependencies = ["exedev>=0.4,<0.5", "openai>=2,<3", "fastapi>=0.115,<1", "uvicorn>=0.30,<1"]
# ///
"""A chat that answers with code: the model writes Python, a sandbox VM runs it.

Each conversation gets its own exe.dev VM with a Jupyter kernel (sandbox.py),
created at its first cell, paused when idle, and deleted after a few idle
days or with the conversation. The model is reached through exe.dev's LLM
gateway, with no API key. Conversations are JSON files in ./data.

https://exe.dev/deploy/code-interpreter/app.py
"""

import json
import os
import re
import threading
import time
import uuid
from pathlib import Path

import exedev
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import FileResponse, Response, StreamingResponse
from starlette.concurrency import run_in_threadpool
from openai import OpenAI

from sandbox import Sandbox

HERE = Path(__file__).parent
DATA = HERE / "data"
MODEL = os.environ.get("MODEL", "gpt-5.5")
EGRESS = os.environ.get("SANDBOX_EGRESS") or None  # "integrations" cuts sandboxes off the internet
PAUSE_AFTER = float(os.environ.get("PAUSE_AFTER_MINUTES", "30")) * 60
DELETE_AFTER = float(os.environ.get("DELETE_AFTER_DAYS", "3")) * 86400
CELL_TIMEOUT = float(os.environ.get("CELL_TIMEOUT", "300"))
MAX_STEPS = 20  # cells per reply
TAG = "code-interpreter"  # sandboxes' comments start with this
BASE = "code-interpreter-base"  # the comment of the VM sandboxes are copied from

exe = exedev.Client(os.environ.get("EXE_TOKEN"), base_url=os.environ.get("EXE_API_URL", "https://exe.dev"))
llm = OpenAI(base_url=os.environ.get("LLM_URL", "https://llm.int.exe.xyz/v1"), api_key=os.environ.get("LLM_API_KEY", "unused"))

SYSTEM = f"""You are a data analyst with a Python 3.12 Jupyter kernel, on a Linux VM of
its own. Use the run_python tool to compute, check, and plot rather than guess.
- Variables, imports, and files persist between calls, as in a notebook.
- The working directory is ~/work. Files the user uploads are there, and files
  you save there are offered to the user for download.
- numpy, pandas, matplotlib, scipy, sympy, scikit-learn, openpyxl, and requests are
  installed. Install others with %pip install <package>.
- A chart shown with plt.show() is displayed to the user; you only see a note
  that it was shown.
- A cell is interrupted after {CELL_TIMEOUT:g} seconds.
After running code, answer in plain prose; the user sees your code and its output."""

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "run_python",
            "description": "Run a cell in the Jupyter kernel. Returns its output, the value of its last line, and any error.",
            "parameters": {
                "type": "object",
                "properties": {"code": {"type": "string", "description": "Python code, or IPython magics such as %pip"}},
                "required": ["code"],
            },
        },
    }
]

app = FastAPI()
locks: dict[str, threading.Lock] = {}
base_lock = threading.Lock()


# Conversations: {id, title, updated, sandbox, messages (for the model), events (for the page)}

def path(cid: str) -> Path:
    if not re.fullmatch(r"[0-9a-f]{32}", cid):
        raise HTTPException(404)
    return DATA / f"{cid}.json"


def load(cid: str) -> dict:
    try:
        return json.loads(path(cid).read_text())
    except FileNotFoundError:
        raise HTTPException(404) from None


def save(conv: dict) -> None:
    tmp = path(conv["id"]).with_suffix(".tmp")
    tmp.write_text(json.dumps(conv))
    tmp.replace(path(conv["id"]))


def lock(cid: str) -> threading.Lock:
    return locks.setdefault(cid, threading.Lock())


def base_vm() -> exedev.VM:
    """The VM with the packages installed that sandboxes are copied from, made once and kept paused."""
    with base_lock:
        bases = sorted((vm for vm in exe.vms() if vm.comment == BASE), key=lambda vm: vm.created_at)
        if bases:
            return bases[0]
        print("setting up the base sandbox", flush=True)
        vm = Sandbox.make_base(exe, comment=BASE)
        try:
            vm.pause()
        except exedev.PermissionDeniedError:
            pass  # the token can't pause: it stays running
        return vm


def open_sandbox(conv: dict):
    """The conversation's sandbox, resumed or made. A generator: it yields status
    events for the page and returns the sandbox and a note for the model."""
    if conv["sandbox"]:
        try:
            vm = exe.vm(conv["sandbox"])
            if vm.status == "paused":
                yield {"type": "status", "text": "Resuming the sandbox\u2026"}
                vm.resume()
            sb = Sandbox(vm)
            if sb.kernel_running():
                return sb, ""
            yield {"type": "status", "text": "Restarting the kernel\u2026"}
            sb.start_kernel()
            return sb, "[The kernel was restarted: variables and imports are gone, files remain.]\n"
        except exedev.NotFoundError:
            note = "[The sandbox was deleted after being idle: variables and files from before are gone.]\n"
    else:
        note = ""
    yield {"type": "status", "text": "Starting a sandbox VM\u2026"}
    sb = Sandbox.create(exe, base=base_vm(), comment=f"{TAG} {conv['id']}", egress=EGRESS)
    conv["sandbox"] = sb.vm.name
    save(conv)
    return sb, note


def reply(conv: dict, text: str):
    """Run the tool-calling loop for one user message, yielding events for the page."""

    def event(e: dict) -> dict:
        conv["events"].append(e)
        return e

    yield event({"type": "user", "text": text})
    conv["messages"].append({"role": "user", "content": text})
    if conv["title"] == "New conversation":
        conv["title"] = text[:60]
    sb = None
    for _ in range(MAX_STEPS):
        turn = len(conv["messages"])
        try:
            msg = llm.chat.completions.create(
                model=MODEL, messages=[{"role": "system", "content": SYSTEM}, *conv["messages"]], tools=TOOLS
            ).choices[0].message
            conv["messages"].append(msg.model_dump(exclude_none=True, include={"role", "content", "tool_calls"}))
            if msg.content:
                yield event({"type": "assistant", "text": msg.content})
            if not msg.tool_calls:
                break
            for call in msg.tool_calls:
                try:
                    code = json.loads(call.function.arguments)["code"]
                except (ValueError, KeyError, TypeError):
                    result = "Error: call run_python with a JSON object with a code string."
                else:
                    yield event({"type": "code", "code": code})
                    note = ""
                    if sb is None:
                        sb, note = yield from open_sandbox(conv)
                    ex = sb.run_code(code, timeout=CELL_TIMEOUT)
                    plain = [r["text/plain"] for r in ex.results if "text/plain" in r and "image/png" not in r]
                    yield event({"type": "output", "stdout": ex.stdout, "stderr": ex.stderr, "results": plain,
                                 "images": ex.pngs, "error": ex.error})
                    result = note + ex.for_model()
                conv["messages"].append({"role": "tool", "tool_call_id": call.id, "content": result})
        except Exception as err:
            # Forget the unfinished turn, which the model API would refuse next time, and show the error.
            del conv["messages"][turn:]
            print(f"{conv['id']}: {err!r}", flush=True)
            yield event({"type": "error", "text": f"{type(err).__name__}: {err}"})
            break
        save(conv)
    else:
        yield event({"type": "assistant", "text": f"(Stopped after {MAX_STEPS} rounds of code. Say \u201cgo on\u201d to continue.)"})
    if sb is not None:
        try:
            yield {"type": "files", "files": sb.files()}
        except exedev.ExeError:
            pass


@app.get("/")
def index():
    return FileResponse(HERE / "index.html")


@app.get("/api/conversations")
def conversations():
    convs = [json.loads(p.read_text()) for p in DATA.glob("*.json")]
    convs.sort(key=lambda c: c["updated"], reverse=True)
    return [{"id": c["id"], "title": c["title"]} for c in convs]


@app.post("/api/conversations")
def new_conversation():
    conv = {"id": uuid.uuid4().hex, "title": "New conversation", "updated": time.time(), "sandbox": None,
            "messages": [], "events": []}
    save(conv)
    return {"id": conv["id"]}


@app.get("/api/conversations/{cid}")
def get_conversation(cid: str):
    conv = load(cid)
    return {"id": cid, "title": conv["title"], "sandbox": conv["sandbox"], "events": conv["events"]}


@app.delete("/api/conversations/{cid}")
def delete_conversation(cid: str):
    with lock(cid):
        conv = load(cid)
        if conv["sandbox"]:
            try:
                exe.vm(conv["sandbox"]).delete()
            except exedev.NotFoundError:
                pass
        path(cid).unlink()
    return {"ok": True}


@app.post("/api/conversations/{cid}/messages")
async def send(cid: str, request: Request):
    text = (await request.json())["text"].strip()
    load(cid)

    def stream():
        with lock(cid):
            conv = load(cid)
            try:
                for e in reply(conv, text):
                    yield json.dumps(e) + "\n"
            finally:
                conv["updated"] = time.time()
                save(conv)

    return StreamingResponse(stream(), media_type="application/x-ndjson")


@app.put("/api/conversations/{cid}/files/{name}")
async def upload(cid: str, name: str, request: Request):
    if not re.fullmatch(r"[\w][\w .-]*", name):
        raise HTTPException(400, "use a plain file name")
    data = await request.body()

    def put():
        with lock(cid):
            conv = load(cid)
            sb, _ = drain(open_sandbox(conv))
            sb.write_file(name, data)
            conv["events"].append({"type": "upload", "name": name, "size": len(data)})
            conv["messages"].append({"role": "user", "content": f"[I uploaded {name} ({len(data)} bytes) to ~/work.]"})
            conv["updated"] = time.time()
            save(conv)

    await run_in_threadpool(put)
    return {"ok": True}


@app.get("/api/conversations/{cid}/files/{name:path}")
def download(cid: str, name: str):
    conv = load(cid)
    if not conv["sandbox"] or ".." in name.split("/"):
        raise HTTPException(404)
    try:
        vm = exe.vm(conv["sandbox"])
        if vm.status == "paused":
            vm.resume()
        data = Sandbox(vm).read_file(name)
    except (exedev.NotFoundError, FileNotFoundError, IsADirectoryError):
        raise HTTPException(404) from None
    filename = name.rsplit("/", 1)[-1].replace('"', "")
    return Response(data, media_type="application/octet-stream",
                    headers={"Content-Disposition": f'attachment; filename="{filename}"'})


@app.get("/api/health")
def health():
    status = {}
    try:
        status["sandboxes"] = len(exe.vms())
        status["exe"] = "ok"
    except exedev.ExeError as err:
        status["exe"] = str(err)
    try:
        status["llm"] = "ok" if MODEL in {m.id for m in llm.models.list()} else f"no model {MODEL}"
    except Exception as err:
        status["llm"] = str(err)
    return status


def drain(gen):
    """Run a generator to its end, ignoring what it yields, and return its value."""
    while True:
        try:
            next(gen)
        except StopIteration as stop:
            return stop.value


def reap() -> None:
    """Pause idle sandboxes, and delete long-idle and orphaned ones."""
    global can_pause
    vms = {vm.name: vm for vm in exe.vms() if (vm.comment or "").startswith(TAG + " ")}
    owned = set()
    for p in DATA.glob("*.json"):
        conv = json.loads(p.read_text())
        vm = vms.get(conv["sandbox"] or "")
        if vm is None:
            continue
        owned.add(vm.name)
        idle = time.time() - conv["updated"]
        if not lock(conv["id"]).acquire(blocking=False):
            continue  # in use
        try:
            if idle > DELETE_AFTER:
                vm.delete()
                print(f"{conv['id']}: deleted idle sandbox {vm.name}", flush=True)
            elif idle > PAUSE_AFTER and vm.status == "running" and can_pause:
                vm.pause()
                print(f"{conv['id']}: paused idle sandbox {vm.name}", flush=True)
        except exedev.PermissionDeniedError as err:
            print(f"not pausing idle sandboxes: {err}", flush=True)
            can_pause = False
        finally:
            lock(conv["id"]).release()
    # Sandboxes whose conversation is gone, made more than an hour ago.
    for vm in vms.values():
        if vm.name not in owned and vm.created_at and time.time() - vm.created_at.timestamp() > 3600:
            vm.delete()
            print(f"deleted orphaned sandbox {vm.name}", flush=True)


def reaper() -> None:
    while True:
        time.sleep(60)
        try:
            reap()
        except Exception as err:
            print(f"reaper: {err!r}", flush=True)


can_pause = True

if __name__ == "__main__":
    DATA.mkdir(exist_ok=True)
    threading.Thread(target=reaper, daemon=True).start()
    uvicorn.run(app, host=os.environ.get("HOST", "0.0.0.0"), port=int(os.environ.get("PORT", "8000")))