"""A code interpreter on an exe.dev VM: a Jupyter kernel that keeps its state.

    import exedev
    from sandbox import Sandbox

    exe = exedev.Client()
    with Sandbox.create(exe) as sb:          # a new VM, deleted at the end
        sb.run_code("import pandas as pd; df = pd.DataFrame({'x': [1, 2, 3]})")
        print(sb.run_code("df.x.sum()").text)   # np.int64(6): variables carry over

Each Sandbox is a whole VM with an IPython kernel in ~/work, started by
exe.dev's Python SDK (https://pypi.org/project/exedev/). run_code() runs a
cell in it and returns what Jupyter would show: printed text, the value of
the last line, charts as PNG, and the traceback if it failed.

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

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field

import exedev

# Installed in each sandbox's ~/.venv. Add what your users' code needs.
PACKAGES = ["ipykernel", "numpy", "pandas", "matplotlib", "scipy", "sympy", "scikit-learn", "openpyxl", "requests"]

ANSI = re.compile(r"\x1b\[[0-9;]*m")  # colors in tracebacks and pip's output

PYTHON = ".venv/bin/python"  # relative to the home directory, where commands start

# The kernel side: runs on the sandbox as ~/.kernel.py. "ready" waits for a
# new kernel; "run" sends it one cell and prints what came back, as JSON.
KERNEL_PY = r'''
import json, os, signal, sys, time
from queue import Empty
from jupyter_client import BlockingKernelClient

kc = BlockingKernelClient(connection_file=os.path.expanduser("~/.kernel.json"))
kc.load_connection_file()
kc.start_channels()

if sys.argv[1] == "ready":
    kc.wait_for_ready(timeout=60)
    kc.execute_interactive("import os; open(os.path.expanduser('~/.kernel.pid'), 'w').write(str(os.getpid()))")
    sys.exit()

code, timeout = sys.argv[2], float(sys.argv[3])
msg_id = kc.execute(code)
deadline, interrupted, out = time.monotonic() + timeout, False, []
while True:
    try:
        msg = kc.get_iopub_msg(timeout=max(deadline - time.monotonic(), 0.1))
    except Empty:
        if interrupted:
            out.append({"type": "error", "name": "KernelDied", "value": "the kernel stopped answering", "traceback": []})
            break
        os.kill(int(open(os.path.expanduser("~/.kernel.pid")).read()), signal.SIGINT)
        interrupted, deadline = True, time.monotonic() + 10
        out.append({"type": "stderr", "text": f"Interrupted: the cell ran for {timeout:g} seconds.\n"})
        continue
    if msg["parent_header"].get("msg_id") != msg_id:
        continue
    kind, c = msg["msg_type"], msg["content"]
    if kind == "stream":
        out.append({"type": c["name"], "text": c["text"]})
    elif kind in ("execute_result", "display_data"):
        out.append({"type": "result", "data": c["data"]})
    elif kind == "error":
        out.append({"type": "error", "name": c["ename"], "value": c["evalue"], "traceback": c["traceback"]})
    elif kind == "status" and c["execution_state"] == "idle":
        break
print(json.dumps(out))
'''

SETUP = f"""#!/bin/bash
set -e
uv venv -q --seed ~/.venv
uv pip install -q --python ~/.venv {' '.join(PACKAGES)}
mkdir -p ~/work
"""


@dataclass
class Execution:
    """What one cell produced, in order, as E2B's run_code() returns it."""

    stdout: str = ""
    stderr: str = ""
    results: list[dict[str, str]] = field(default_factory=list)  # MIME type -> data, as Jupyter sends it
    error: dict | None = None  # name, value, traceback

    @property
    def text(self) -> str | None:
        """The text of the last result, such as a cell's final value."""
        texts = [r["text/plain"] for r in self.results if "text/plain" in r]
        return texts[-1] if texts else None

    @property
    def pngs(self) -> list[str]:
        """Charts and other images, base64-encoded PNG."""
        return [r["image/png"] for r in self.results if "image/png" in r]

    def for_model(self, limit: int = 10_000) -> str:
        """A plain-text summary for an LLM: output, results, and the error."""
        parts = [self.stdout, self.stderr]
        for r in self.results:
            if "image/png" in r:
                parts.append("[a chart was shown to the user]\n")
            elif "text/plain" in r:
                parts.append(r["text/plain"] + "\n")
        if self.error:
            parts.append(f"{self.error['name']}: {self.error['value']}\n")
        text = ANSI.sub("", "".join(parts)) or "(no output)"
        return text if len(text) <= limit else "[...]\n" + text[-limit:]


class Sandbox:
    """An exe.dev VM running an IPython kernel in ~/work."""

    def __init__(self, vm: exedev.VM) -> None:
        self.vm = vm

    @classmethod
    def create(
        cls, exe: exedev.Client, *, base: exedev.VM | None = None, comment: str = "code sandbox", **new_args
    ) -> Sandbox:
        """A new sandbox: a copy of base (see make_base), or a new VM set up from scratch.

        new_args go to exe.new() or base.copy(), such as egress="integrations".
        """
        if base is not None:
            vm = base.copy(**new_args)
        else:
            vm = exe.new(comment=comment, setup_script=SETUP, wait_until="setup", timeout=600, **new_args)
        try:
            vm.set_comment(comment)
            sandbox = cls(vm)
            sandbox.start_kernel()
            return sandbox
        except BaseException:
            vm.delete()
            raise

    @staticmethod
    def make_base(exe: exedev.Client, comment: str = "code sandbox base") -> exedev.VM:
        """A VM with the packages installed, to copy() sandboxes from in a second or two."""
        return exe.new(comment=comment, setup_script=SETUP, wait_until="setup", timeout=600)

    def start_kernel(self) -> None:
        """Start a new kernel, stopping the old one: variables and imports are lost, files stay."""
        kernel = self._kernel()
        if kernel is not None:
            kernel.kill()
            kernel.wait(timeout=10)
        self.vm.write_file(".kernel.py", KERNEL_PY)
        self.vm.run("rm -f .kernel.json .kernel.pid", check=True)
        self.vm.spawn("cd ~/work && exec ~/.venv/bin/python -m ipykernel_launcher -f ~/.kernel.json", name="kernel")
        self.vm.run(
            f"for i in $(seq 100); do [ -s .kernel.json ] && break; sleep 0.1; done; {PYTHON} .kernel.py ready",
            check=True,
            timeout=90,
        )

    def kernel_running(self) -> bool:
        return self._kernel() is not None

    def _kernel(self) -> exedev.Process | None:
        for p in self.vm.processes():
            try:
                if p.name == "kernel" and p.poll() is None:
                    return p
            except exedev.ProcessLostError:  # killed, or the VM restarted
                pass
        return None

    def run_code(self, code: str, timeout: float = 120) -> Execution:
        """Run a cell. Past timeout seconds it's interrupted, like Ctrl-C, and the kernel lives on."""
        r = self.vm.run([PYTHON, ".kernel.py", "run", code, str(timeout)], timeout=timeout + 60, check=True)
        ex = Execution()
        for item in json.loads(r.stdout):
            if item["type"] == "stdout":
                ex.stdout += item["text"]
            elif item["type"] == "stderr":
                ex.stderr += item["text"]
            elif item["type"] == "result":
                ex.results.append(item["data"])
            else:
                ex.error = {k: item[k] for k in ("name", "value", "traceback")}
        return ex

    def files(self) -> list[str]:
        """Files under ~/work, relative to it, skipping hidden ones."""
        r = self.vm.run("cd work && find . -type f -not -path '*/.*' -printf '%P\\n' | sort | head -200", check=True)
        return r.stdout.splitlines()

    def write_file(self, path: str, data: str | bytes) -> None:
        self.vm.write_file("work/" + path.lstrip("/"), data)

    def read_file(self, path: str) -> bytes:
        return self.vm.read_file("work/" + path.lstrip("/"))

    def delete(self) -> None:
        self.vm.delete()

    def __enter__(self) -> Sandbox:
        return self

    def __exit__(self, *exc) -> None:
        try:
            self.vm.delete()
        except exedev.NotFoundError:
            pass


__all__ = ["Sandbox", "Execution", "PACKAGES"]