#!/usr/bin/env python3
"""
pi-rlm kernel shim.

A persistent Python 3 REPL that talks to the pi-rlm extension over a unix
socket using newline-delimited JSON. Standard library only — no ipykernel,
no jupyter_client, no pip install required.

The kernel owns a single global namespace that survives across executions.
The namespace contains:

  rlm(prompt, context=None) -> str
      Call a recursive sub-LLM through the host extension and return its
      final answer as a string. Blocks until the host responds.

  user_prompt : str
      The current user request, injected by the host before each turn.

Protocol (kernel -> host), one JSON object per line:
  {"type": "ready", "python": str, "pid": int, "depth": int}
  {"type": "result", "id": int, "ok": bool, "stdout": str, "stderr": str,
   "result": str | None}
  {"type": "stream", "id": int, "stream": "stdout" | "stderr", "data": str}
  {"type": "rlm_request", "id": int, "prompt": str, "depth": int}
  {"type": "set_done", "id": int}
  {"type": "vars", "id": int, "vars": [{"name", "type", "repr"}]}

Protocol (host -> kernel):
  {"type": "exec", "id": int, "code": str}
  {"type": "set", "id": int, "name": str, "value": any}
  {"type": "vars_request", "id": int}
  {"type": "snapshot_request", "id": int}   → {"type": "snapshot_done", ...}
  {"type": "rlm_response", "id": int, "ok": bool,
   "result"?: str, "error"?: str}
  {"type": "refine_response", "id": int, "scheduled": bool, "reason"?: str}

Snapshots: if PI_RLM_SNAPSHOT points at a path, the kernel restores
picklable namespace values from it at boot (advertised in "ready") and
writes it on snapshot_request. Only data is pickled — cell-defined
functions/classes, modules and builtins are skipped.

Concurrency note: the main thread is the only socket reader. rlm() calls
(including from user-spawned threads) are serialized with a lock; while an
rlm() call is blocked waiting, the lock holder is the only reader, and the
host only sends rlm_response messages during that window.
"""

import ast
import io
import json
import os
import socket
import sys
import threading
import traceback

SOCKET_PATH = os.environ.get("PI_RLM_SOCKET", "")
RLM_DEPTH = int(os.environ.get("PI_RLM_DEPTH", "0") or "0")
SNAPSHOT_PATH = os.environ.get("PI_RLM_SNAPSHOT", "")

# Cap repr() of trailing expressions. stdout/stderr are capped host-side.
MAX_REPR_CHARS = 100_000
VAR_SNAPSHOT_LIMIT = 200
# Never pickle a single value larger than this.
SNAPSHOT_VALUE_CAP = 64 * 1024 * 1024
# Names injected by the host; never snapshotted.
_HOST_NAMES = frozenset(("rlm", "refine", "user_prompt"))


def _die(msg: str, code: int = 2) -> "None":
    sys.stderr.write(f"pi-rlm kernel: {msg}\n")
    sys.exit(code)


if not SOCKET_PATH:
    _die("PI_RLM_SOCKET is not set (this process is started by the pi-rlm extension)")

try:
    _sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    _sock.connect(SOCKET_PATH)
except OSError as exc:
    _die(f"cannot connect to {SOCKET_PATH}: {exc}")

_rfile = _sock.makefile("r", encoding="utf-8", errors="replace", newline="\n")
_wlock = threading.Lock()
_rlm_lock = threading.Lock()
_rlm_counter = 0


def _send(msg: dict) -> None:
    data = (json.dumps(msg, ensure_ascii=False) + "\n").encode("utf-8")
    with _wlock:
        _sock.sendall(data)


def _readline() -> dict:
    line = _rfile.readline()
    if not line:
        raise RuntimeError("pi-rlm host closed the connection")
    return json.loads(line)


def rlm(prompt, context=None):
    """Call a recursive sub-LLM. Returns its final answer as a string.

    The sub-agent sees only what you pass it — build prompts with f-strings
    or pass data via ``context``. Raises RuntimeError if the call fails or
    the maximum recursion depth is reached.
    """
    global _rlm_counter
    text = str(prompt)
    if context is not None:
        if isinstance(context, str):
            payload = context
        else:
            try:
                payload = json.dumps(context, ensure_ascii=False, default=str)
            except Exception:
                payload = str(context)
        text = f"{text}\n\n{payload}"
    with _rlm_lock:
        _rlm_counter += 1
        rid = _rlm_counter
        _send({"type": "rlm_request", "id": rid, "prompt": text, "depth": RLM_DEPTH + 1})
        while True:
            msg = _readline()
            if msg.get("type") == "rlm_response" and msg.get("id") == rid:
                if msg.get("ok"):
                    return msg.get("result", "")
                raise RuntimeError(f"rlm() failed: {msg.get('error', 'unknown error')}")
            # Anything else while waiting is unexpected; ignore it.


def refine(instructions=None, global_=False):
    """Schedule a harness refinement for when the current turn ends.

    Returns {"scheduled": True} immediately, or {"scheduled": False,
    "reason": ...}. Use after observing a repeated failure, a reusable
    tactic, or a behavior policy worth persisting. global_=True targets
    the cross-session harness store.
    """
    global _rlm_counter
    with _rlm_lock:
        _rlm_counter += 1
        rid = _rlm_counter
        _send({
            "type": "refine_request",
            "id": rid,
            "instructions": None if instructions is None else str(instructions),
            "global": bool(global_),
        })
        while True:
            msg = _readline()
            if msg.get("type") == "refine_response" and msg.get("id") == rid:
                out = {"scheduled": bool(msg.get("scheduled"))}
                if msg.get("reason"):
                    out["reason"] = msg["reason"]
                return out


class _StreamProxy:
    """file-like object that buffers output and mirrors it to the host."""

    def __init__(self, name: str, exec_id: int, buffer: io.StringIO):
        self._name = name
        self._exec_id = exec_id
        self._buffer = buffer

    def write(self, s):
        if not isinstance(s, str):
            s = str(s)
        self._buffer.write(s)
        if s:
            try:
                _send({"type": "stream", "id": self._exec_id, "stream": self._name, "data": s})
            except Exception:
                pass  # never let streaming break user code
        return len(s)

    def flush(self):
        pass

    def isatty(self):
        return False

    def writelines(self, lines):
        for line in lines:
            self.write(line)


_NAMESPACE: dict = {"__name__": "__rlm__", "rlm": rlm, "refine": refine}


def _execute(exec_id: int, code: str) -> dict:
    out_buf, err_buf = io.StringIO(), io.StringIO()
    result_repr = None
    ok = True
    old_stdout, old_stderr = sys.stdout, sys.stderr
    sys.stdout = _StreamProxy("stdout", exec_id, out_buf)
    sys.stderr = _StreamProxy("stderr", exec_id, err_buf)
    try:
        tree = ast.parse(code, "<cell>", "exec")
        if tree.body and isinstance(tree.body[-1], ast.Expr):
            # Interactive behavior: eval a trailing expression and report its repr.
            last = ast.Expression(tree.body.pop().value)
            ast.fix_missing_locations(last)
            if tree.body:
                exec(compile(tree, "<cell>", "exec"), _NAMESPACE)
            value = eval(compile(last, "<cell>", "eval"), _NAMESPACE)
            if value is not None:
                result_repr = repr(value)
                if len(result_repr) > MAX_REPR_CHARS:
                    result_repr = (
                        result_repr[:MAX_REPR_CHARS]
                        + f"... [repr truncated at {MAX_REPR_CHARS} chars]"
                    )
        else:
            exec(compile(tree, "<cell>", "exec"), _NAMESPACE)
    except KeyboardInterrupt:
        ok = False
        err_buf.write("KeyboardInterrupt: cell interrupted\n")
    except SystemExit as exc:
        ok = False
        err_buf.write(f"SystemExit: {exc}\n")
    except BaseException:
        ok = False
        err_buf.write(traceback.format_exc())
    finally:
        sys.stdout = old_stdout
        sys.stderr = old_stderr
    return {
        "type": "result",
        "id": exec_id,
        "ok": ok,
        "stdout": out_buf.getvalue(),
        "stderr": err_buf.getvalue(),
        "result": result_repr,
    }


def _vars_snapshot() -> list:
    items = []
    for name, value in sorted(_NAMESPACE.items()):
        if name.startswith("__"):
            continue
        type_name = type(value).__name__
        try:
            r = repr(value)
        except Exception:
            r = f"<{type_name}>"
        if len(r) > 120:
            r = r[:120] + "..."
        items.append({"name": name, "type": type_name, "repr": r})
        if len(items) >= VAR_SNAPSHOT_LIMIT:
            break
    return items


def _snapshot() -> dict:
    """Pickle the data values of the namespace to SNAPSHOT_PATH (atomic)."""
    import pickle

    if not SNAPSHOT_PATH:
        return {"ok": False, "saved": [], "skipped": [], "error": "no snapshot path configured"}
    payload = {}
    saved, skipped = [], []
    for name, value in list(_NAMESPACE.items()):
        if name.startswith("__") or name in _HOST_NAMES:
            continue
        try:
            blob = pickle.dumps(value, protocol=pickle.HIGHEST_PROTOCOL)
        except Exception:
            skipped.append(name)  # modules, cell-defined functions/classes, …
            continue
        if len(blob) > SNAPSHOT_VALUE_CAP:
            skipped.append(name)
            continue
        payload[name] = value
        saved.append(name)
    try:
        tmp = SNAPSHOT_PATH + ".tmp"
        with open(tmp, "wb") as fh:
            pickle.dump(payload, fh, protocol=pickle.HIGHEST_PROTOCOL)
        os.replace(tmp, SNAPSHOT_PATH)
    except OSError as exc:
        return {"ok": False, "saved": [], "skipped": skipped, "error": str(exc)}
    return {"ok": True, "saved": saved, "skipped": skipped}


def _restore() -> list:
    """Load a previous snapshot into the namespace. Returns restored names."""
    import pickle

    if not SNAPSHOT_PATH or not os.path.exists(SNAPSHOT_PATH):
        return []
    try:
        with open(SNAPSHOT_PATH, "rb") as fh:
            payload = pickle.load(fh)
    except Exception:
        return []
    if not isinstance(payload, dict):
        return []
    restored = []
    for name, value in payload.items():
        if isinstance(name, str) and name.isidentifier() and name not in _HOST_NAMES:
            _NAMESPACE[name] = value
            restored.append(name)
    return restored


def main() -> None:
    restored = _restore()
    _send({
        "type": "ready",
        "python": sys.version.split()[0],
        "pid": os.getpid(),
        "depth": RLM_DEPTH,
        "restored": restored,
    })
    while True:
        try:
            line = _rfile.readline()
        except (OSError, ValueError):
            break
        if not line:
            break  # host gone
        try:
            msg = json.loads(line)
        except json.JSONDecodeError:
            continue
        mtype = msg.get("type")
        mid = msg.get("id")
        if mtype == "exec":
            try:
                _send(_execute(mid, msg.get("code", "")))
            except Exception:
                try:
                    _send({
                        "type": "result",
                        "id": mid,
                        "ok": False,
                        "stdout": "",
                        "stderr": "kernel internal error:\n" + traceback.format_exc(),
                        "result": None,
                    })
                except Exception:
                    break
        elif mtype == "set":
            name = msg.get("name")
            if isinstance(name, str) and name.isidentifier():
                _NAMESPACE[name] = msg.get("value")
            try:
                _send({"type": "set_done", "id": mid})
            except Exception:
                break
        elif mtype == "vars_request":
            try:
                _send({"type": "vars", "id": mid, "vars": _vars_snapshot()})
            except Exception:
                break
        elif mtype == "snapshot_request":
            try:
                _send({"type": "snapshot_done", "id": mid, **_snapshot()})
            except Exception:
                break
        # rlm_response / refine_response with no waiter are ignored.


if __name__ == "__main__":
    main()