"""
One pyrepl session: a Python process holding one namespace.

Spawned by pyrepld, never by hand. Talks to the daemon over two inherited
pipes (fd numbers in argv), one JSON object per line. stdin is /dev/null so
input() gets EOF instead of reading the protocol; fd 1/2 go to a pipe the
daemon reads, which catches output from subprocesses and C code. Python-level
print() goes over the control pipe instead, so it is ordered with "done".
A forked child (multiprocessing, os.fork) writes to fd 1/2 instead: its
messages could interleave with ours and corrupt the control stream.

`worker.py --stdio` is the remote and root form: the protocol is on
stdin/stdout and fd 1/2 are moved to stderr, which ssh carries back
separately. Anything printed before "ready" (rc files, banners) is skipped by
the daemon. On another host it doesn't run from disk: boot.py arrives over
ssh with this file, tk.py and vendor/ as source and runs it from memory, so
PYREPL_FILES is set and nothing here is read from HERE.

Code runs in the main thread so that SIGINT from the daemon lands as a
KeyboardInterrupt in the user's code -- that is what makes interrupt work
even inside time.sleep(). A reader thread takes every control line: work
(exec/put/get/call) is queued for the main thread, "interrupt" raises SIGINT
here (the only way to reach a remote worker), and replies to hub requests
are handed to whoever is waiting for them.

put/get/call move objects between sessions as cloudpickle bytes (vendor/),
so functions and closures travel by value. Both ends need the same Python
minor version for code objects; plain data loads anywhere. swapout pickles
the whole namespace the same way, when pyrepld swaps an idle session to disk.

Must stay runnable on python3.6: a session can be started with any
interpreter the caller names (a venv, /usr/bin/python3).
"""
import ast
import base64
import codecs
import importlib.abc
import importlib.util
import io
import itertools
import json
import linecache
import os
import pickle
import queue
import signal
import sys
import threading
import time
import traceback
import types

HERE = os.path.dirname(os.path.abspath(__file__))
FILES = globals().get("PYREPL_FILES")     # set by boot.py on a remote host
if sys.argv[1:] == ["--stdio"]:
    CTL_IN, CTL_OUT = os.dup(0), os.dup(1)
    _null = os.open(os.devnull, os.O_RDONLY)
    os.dup2(_null, 0)
    os.close(_null)
    os.dup2(2, 1)
else:
    CTL_IN, CTL_OUT = int(sys.argv[1]), int(sys.argv[2])
os.set_inheritable(CTL_IN, False)
os.set_inheritable(CTL_OUT, False)
_ctl_in = os.fdopen(CTL_IN, "r", encoding="utf-8")


def _json_escape(err):
    # A lone surrogate (a non-UTF-8 filename, say) in output or in an error
    # message must not make the protocol raise. They only occur inside JSON
    # strings, where an escaped backslash makes them arrive as visible "\udce9".
    bad = err.object[err.start:err.end]
    return "".join("\\\\u%04x" % ord(c) for c in bad), err.end


codecs.register_error("pyrepl-json", _json_escape)
_ctl_out = os.fdopen(CTL_OUT, "w", encoding="utf-8", errors="pyrepl-json")
_ctl_lock = threading.Lock()
_tls = threading.local()
_PID = os.getpid()
_SIGINT = {signal.SIGINT}

REPR_CAP = 10000
SHOW_CAP = 30000         # results that define __pyrepl_show__ (tk.py's) print as text
OUT_CHUNK = 1 << 20      # one print() of 500 MB must not become one 500 MB line


def send(**msg):
    line = json.dumps(msg, ensure_ascii=False) + "\n"
    pending = getattr(_tls, "pending", None)
    if pending is not None:
        # Re-entered on this thread: a __del__ run by the GC while we held the
        # lock printed something. Queue it for the outer call.
        pending.append(line)
        return
    _tls.pending = pending = [line]
    try:
        with _ctl_lock:
            # A KeyboardInterrupt in the middle of a write would leave half a
            # line on the pipe and wedge the daemon's reader, so SIGINT waits.
            old = signal.pthread_sigmask(signal.SIG_BLOCK, _SIGINT)
            try:
                while pending:
                    _ctl_out.write(pending.pop(0))
                    _ctl_out.flush()
            finally:
                signal.pthread_sigmask(signal.SIG_SETMASK, old)
    finally:
        _tls.pending = None       # nothing allocates between the last check and here


# -- requests to the daemon (the hub) from code running in this session ------
_replies = {}             # id -> [Event, reply]
_ids = itertools.count(1)


def hub(req, timeout=None):
    """Send a request to pyrepld and wait for its answer.

    Waits in short steps so a KeyboardInterrupt can land; the caller decides
    what an interrupted request means (tk forwards it to the other session).
    """
    rid = next(_ids)
    slot = _replies[rid] = [threading.Event(), None]
    send(t="hub", id=rid, req=req)
    deadline = None if timeout is None else time.time() + timeout
    try:
        while not slot[0].wait(0.2):
            if deadline is not None and time.time() > deadline:
                raise TimeoutError("pyrepld did not answer %r within %ss (a pyrepld started before "
                                   "2026-10-20 has no hub: `pyrepl stop`, then use it again)"
                                   % (req.get("op"), timeout))
    finally:
        _replies.pop(rid, None)
    return slot[1]


class Stream(object):
    """sys.stdout / sys.stderr replacement: every write becomes a message."""

    def __init__(self, name):
        self.name = name
        self.encoding = "utf-8"
        self.errors = "replace"
        self.buffer = BinaryStream(self)

    def write(self, s):
        if not isinstance(s, str):
            s = str(s)
        if not s:
            return 0
        if os.getpid() != _PID:
            os.write(self.fileno(), s.encode("utf-8", "replace"))
        else:
            for i in range(0, len(s), OUT_CHUNK):
                send(t="out", s=s[i:i + OUT_CHUNK])
        return len(s)

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

    def flush(self):
        pass

    def isatty(self):
        return False

    def fileno(self):
        # Real fd, so subprocess.run(stdout=sys.stdout) and faulthandler work;
        # anything written there reaches the daemon through the raw pipe.
        return 1 if self.name == "stdout" else 2

    def readable(self):
        return False

    def writable(self):
        return True


class BinaryStream(object):
    """sys.stdout.buffer: bytes in, text out."""

    def __init__(self, text):
        self.text = text
        self.decoder = codecs.getincrementaldecoder("utf-8")("replace")

    def write(self, b):
        self.text.write(self.decoder.decode(bytes(b)))
        return len(b)

    def flush(self):
        pass

    def fileno(self):
        return self.text.fileno()

    def isatty(self):
        return False

    def writable(self):
        return True


# A real __main__ module, so functions and classes defined in the session
# pickle (multiprocessing, concurrent.futures) and look like REPL objects.
_main = types.ModuleType("__main__")
_main.__builtins__ = __builtins__
sys.modules["__main__"] = _main
NS = _main.__dict__

# Imports resolve from the session's cwd, like the interactive interpreter,
# not from the directory this file lives in. vendor/ (cloudpickle) goes last
# so a session's own packages win.
sys.path[0] = ""
if not FILES:                 # boot.py serves vendor/ from memory, also last
    sys.path.append(os.path.join(HERE, "vendor"))
sys.stdout = Stream("stdout")
sys.stderr = Stream("stderr")
sys.stdin = open(os.devnull)


def load_toolkit():
    """Preload tk.py from next to this file; PYREPL_NO_TK=1 skips it.

    Loaded by path rather than imported, so a tk.py in the session's cwd
    can't stand in for it. A broken toolkit costs the toolkit, not the session.
    """
    if os.environ.get("PYREPL_NO_TK"):
        return {}
    path = os.path.join(os.path.dirname(__file__), "tk.py") if FILES else os.path.join(HERE, "tk.py")
    try:
        mod = types.ModuleType("tk")
        mod.__file__ = path
        if FILES:
            src = FILES["tk.py"]
            linecache.cache[path] = (len(src), None, src.splitlines(True), path)
        else:
            with open(path, encoding="utf-8") as fh:
                src = fh.read()
        exec(compile(src, path, "exec"), mod.__dict__)
    except Exception:
        sys.stderr.write("pyrepl: tk.py failed to load; this session runs without it\n"
                         + traceback.format_exc())
        return {}
    mod._hub, mod._dumps, mod._loads = hub, _dumps, _loads
    sys.modules["tk"] = mod
    NS["tk"] = mod
    for k in mod.__all__:
        NS[k] = getattr(mod, k)
    return dict((k, NS[k]) for k in mod.__all__ + ["tk"])


def user_names():
    out = []
    for k, v in list(NS.items()):   # a user thread may be adding globals
        if k.startswith("__") or k == "_" or TOOLKIT.get(k, NS) is v:
            continue
        out.append([k, type(v).__name__])
    out.sort()
    return out


def run(cell, code):
    """Exec statements; if the last one is an expression, return its repr."""
    fname = "<cell %d>" % cell
    linecache.cache[fname] = (len(code), None, code.splitlines(True), fname)
    tree = ast.parse(code, fname, "exec")
    last = None
    if tree.body and isinstance(tree.body[-1], ast.Expr):
        last = ast.Expression(tree.body.pop().value)
    exec(compile(tree, fname, "exec"), NS)
    if last is None:
        return None
    value = eval(compile(last, fname, "eval"), NS)
    if value is None:
        return None
    NS["_"] = value
    show = getattr(type(value), "__pyrepl_show__", None)
    if show is not None:
        r = show(value)
        if len(r) > SHOW_CAP:
            head, tail = SHOW_CAP // 3, SHOW_CAP * 2 // 3
            r = r[:head] + "\n[... %d chars omitted ...]\n" % (len(r) - head - tail) + r[-tail:]
        return r
    r = repr(value)
    if len(r) > REPR_CAP:
        r = r[:REPR_CAP] + "... [repr truncated, %d chars total]" % len(r)
    return r


def clean_tb(exc_type, exc, tb):
    """Traceback without the frames that belong to this file."""
    frames = traceback.extract_tb(tb)
    me = os.path.abspath(__file__)
    keep = [f for f in frames if os.path.abspath(f.filename) != me]
    lines = ["Traceback (most recent call last):\n"] if keep else []
    lines += traceback.format_list(keep)
    lines += traceback.format_exception_only(exc_type, exc)
    return "".join(lines)


# -- moving objects between sessions ----------------------------------------
def _pickler():
    try:
        import cloudpickle
        return cloudpickle
    except ImportError:
        return pickle     # data still moves; functions defined in a session do not


MINOR = "%d.%d" % sys.version_info[:2]


class _DataOnly(pickle.Unpickler):
    """For a blob from another minor version: plain data yes, code no."""

    def find_class(self, module, name):
        if module.split(".")[0] == "cloudpickle":
            raise pickle.UnpicklingError(
                "it holds code (a function or class defined in a session), and code objects "
                "don't carry across minor versions")
        return pickle.Unpickler.find_class(self, module, name)


def _dumps(obj):
    """Blob = "<sender's major.minor>:" + base64 pickle."""
    return MINOR + ":" + base64.b64encode(_pickler().dumps(obj, protocol=4)).decode("ascii")


def _loads(blob):
    minor, _, data = blob.partition(":")
    raw = base64.b64decode(data)
    if minor == MINOR:
        _pickler()        # cloudpickle's reconstructors must be importable first
        return pickle.loads(raw)
    try:
        return _DataOnly(io.BytesIO(raw)).load()
    except pickle.UnpicklingError as exc:
        raise TypeError("can't load this here: it came from Python %s and this session runs %s; %s"
                        % (minor, MINOR, exc))


def _adopt(f, seen=None):
    """Make a function that arrived by value live in this namespace.

    cloudpickle gives it a private snapshot of the sender's globals; instead
    it gets NS, so names defined here resolve here. Names it uses that this
    session lacks come along from the snapshot (helpers travel with it);
    names both sides have stay ours. Only top-level functions are adopted:
    one tucked inside a container keeps its snapshot.
    """
    if (not isinstance(f, types.FunctionType) or f.__globals__ is NS
            or f.__module__ != "__main__"):
        return f
    seen = seen if seen is not None else {}
    if id(f) in seen:
        return seen[id(f)]
    g = types.FunctionType(f.__code__, NS, f.__name__, f.__defaults__, f.__closure__)
    g.__kwdefaults__, g.__qualname__, g.__doc__ = f.__kwdefaults__, f.__qualname__, f.__doc__
    g.__dict__.update(f.__dict__)
    seen[id(f)] = g
    for k, v in list(f.__globals__.items()):
        if not k.startswith("__") and k not in NS:
            NS[k] = _adopt(v, seen)
    return g


def _short(v):
    r = repr(v)
    return r if len(r) <= 200 else r[:200] + "..."


def do_put(req):
    objs = _loads(req["blob"])
    seen = {}
    for k, v in objs.items():
        NS[k] = _adopt(v, seen)
    return "put " + ", ".join("%s: %s" % (k, type(v).__name__) for k, v in objs.items()), None


def do_get(req):
    missing = [n for n in req["names"] if n not in NS]
    if missing:
        raise NameError("not defined in this session: %s" % ", ".join(missing))
    vals = dict((n, NS[n]) for n in req["names"])
    try:
        blob = _dumps(vals)
    except Exception as exc:
        raise TypeError("can't copy %s out of this session (%s). Keep it here and call "
                        "functions on it instead." % (", ".join(vals), exc))
    return "get " + ", ".join(vals), blob


def do_call(req):
    fn, args, kw, into = _loads(req["blob"])
    if isinstance(fn, str):
        # A name here, or module.name for a loaded module (tk._fread): that
        # one goes through sys.modules, so `import tkinter as tk` can't hide it.
        mod, _, attr = fn.rpartition(".")
        if mod:
            if mod not in sys.modules:
                raise NameError("no module %r loaded in this session to call %s from" % (mod, attr))
            fn = getattr(sys.modules[mod], attr)
        elif fn in NS:
            fn = NS[fn]
        else:
            raise NameError("no %r in this session to call" % fn)
    else:
        fn = _adopt(fn)
    value = fn(*args, **kw)
    label = getattr(fn, "__name__", "call")
    if into:
        NS[into] = value
        return "%s -> %s = %s" % (label, into, _short(value)), _dumps(None)
    try:
        blob = _dumps(value)
    except Exception as exc:
        NS["_"] = value
        raise TypeError("%s returned a %s, which can't be copied back (%s). It is kept "
                        "here as `_`; pass into='name' to keep it under a name and call "
                        "functions on it." % (label, type(value).__name__, exc))
    return "%s -> %s" % (label, _short(value)), blob


def do_swapout(req):
    """The namespace as one blob, for pyrepld to keep on disk while the
    session is swapped out (it comes back through do_put). Names that won't
    pickle are left behind and listed, so the restore can say what was lost."""
    tk = sys.modules.get("tk")
    running = [j for j in getattr(tk, "_JOBS", []) if not j.done]
    if running:
        raise RuntimeError("bg job(s) still running: %s" % ", ".join(repr(j) for j in running))
    vals = dict((k, NS[k]) for k, _ in user_names())
    lost = []
    try:
        blob = _dumps(vals)
    except Exception:
        for k in sorted(vals):
            try:
                _dumps(vals[k])
            except Exception as exc:
                lost.append([k, type(vals[k]).__name__, str(exc)[:200]])
                del vals[k]
        blob = _dumps(vals)
    threads = [t.name for t in threading.enumerate()
               if t is not threading.main_thread() and t is not _READER and t.is_alive()]
    return json.dumps({"cwd": os.getcwd(), "lost": lost, "threads": threads}), blob


def do_swapin(req):
    """A swapout blob back into the namespace, in the cwd it had if that still exists."""
    try:
        os.chdir(req["cwd"])
    except (OSError, KeyError, TypeError):
        pass
    return do_put(req)[0], json.dumps({"cwd": os.getcwd()})


OPS = {"put": do_put, "get": do_get, "call": do_call, "swapout": do_swapout, "swapin": do_swapin}


# -- modules: the `pyrepl` package ------------------------------------------
_serving = False          # hub replies need the reader thread, which main() starts
# Same places as tk.module_path (pyrepld lists and serves them from tk).
MODULES = os.path.join(os.environ.get("XDG_CONFIG_HOME") or os.path.expanduser("~/.config"),
                       "pyrepl", "modules")
SHIPPED = os.path.join(HERE, "modules")


class _Modules(importlib.abc.MetaPathFinder, importlib.abc.Loader):
    """First on sys.meta_path, and only for the `pyrepl` package: pyrepl.NAME
    is NAME.py from the user's folder (~/.config/pyrepl/modules), else from
    the modules/ shipped with pyrepl. Owning the name means nothing on
    sys.path can stand in for it (a pyrepl/ directory in the cwd, say), and
    a module can't collide with a real package. Nothing is preloaded, and
    importlib.reload() fetches the file again.

    A local session as the user reads the files itself: no round trip, never
    stale. A remote or root one asks pyrepld, which reads them on its own
    host and sends the source (an older pyrepld without the hub: this host's
    folders are read directly)."""

    def __init__(self):
        self.src = {}
        self.hub_ok = True

    def _from_hub(self, name):
        """(source, path), False for "not there", None for "ask the disk"."""
        if not (self.hub_ok and _serving and os.getpid() == _PID):
            return None
        try:
            r = hub({"op": "module", "name": name}, 10)
        except TimeoutError:
            r = {"error": "unknown hub op"}
        if "unknown hub op" in (r.get("error") or ""):
            self.hub_ok = False
            return None
        return (r["source"], r["path"]) if "source" in r else False

    def _from_disk(self, name):
        for folder in (MODULES, SHIPPED):
            path = os.path.join(folder, name + ".py")
            try:
                with open(path, encoding="utf-8") as fh:
                    return fh.read(), path
            except (OSError, UnicodeDecodeError):
                pass
        return None

    def find_spec(self, name, path=None, target=None):
        if name == "pyrepl":
            return importlib.util.spec_from_loader(name, self, is_package=True)
        if not name.startswith("pyrepl.") or not name[7:].isidentifier():
            return None
        sub = name[7:]
        found = self._from_hub(sub) if (FILES or os.geteuid() == 0) else None
        if found is None:
            found = self._from_disk(sub)
        if not found:
            return None
        self.src[name] = found[0]
        spec = importlib.util.spec_from_loader(name, self, origin=found[1])
        spec.has_location = True
        return spec

    def create_module(self, spec):
        return None

    def exec_module(self, module):
        name = module.__spec__.name
        if name == "pyrepl":
            module.__doc__ = "pyrepl's modules: `from pyrepl import NAME`. tk.modules() lists them."
            return
        path = module.__spec__.origin
        src = self.src.pop(name)
        linecache.cache[path] = (len(src), None, src.splitlines(True), path)
        exec(compile(src, path, "exec"), module.__dict__)


sys.meta_path.insert(0, _Modules())
TOOLKIT = load_toolkit()      # last: tk is handed hub and the codec above
WORK = queue.Queue()
_READER = None


def reader():
    """Every control line lands here first; see the module docstring."""
    for line in _ctl_in:
        try:
            req = json.loads(line)
        except ValueError:      # not ours (a stray password line?): never echo it
            sys.stderr.write("pyrepl: ignored a control line that was not JSON\n")
            continue
        if not isinstance(req, dict):
            continue
        op = req.get("op")
        if op == "interrupt":
            os.kill(os.getpid(), signal.SIGINT)
        elif op == "reply":
            slot = _replies.get(req.get("id"))
            if slot:                        # nobody waiting: an interrupted request
                slot[1] = req.get("reply")
                slot[0].set()
        else:
            WORK.put(req)
    os._exit(0)           # the daemon (or the ssh link) is gone; so is this session


_armed = False


def on_sigint(signum, frame):
    # Inside a run SIGINT is a KeyboardInterrupt in the user's code. Outside
    # one it is a no-op: an interrupt that loses the race with the end of a
    # cell must not land in this loop, or in the middle of reading a request.
    if _armed:
        raise KeyboardInterrupt


def main():
    global _armed, _serving, _READER
    send(t="ready", pid=os.getpid(), python=sys.version.split()[0],
         executable=sys.executable, cwd=os.getcwd(), host=os.uname()[1].split(".")[0])
    _READER = threading.Thread(target=reader, daemon=True)
    _READER.start()
    _serving = True
    while True:
        try:
            req = WORK.get()
        except KeyboardInterrupt:
            continue
        op = req.get("op")
        t0 = time.time()
        result, error, blob = None, None, None
        _armed = True
        try:
            try:
                if op == "exec":
                    result = run(req["cell"], req["code"])
                elif op in OPS:
                    result, blob = OPS[op](req)
                else:
                    error = "unknown op %r\n" % op
            finally:
                _armed = False   # a SIGINT before this line is caught just below
        except SystemExit as exc:
            error = "SystemExit(%r) -- ignored, the session is still alive\n" % (exc.code,)
        except BaseException:  # KeyboardInterrupt included: report, keep going
            error = clean_tb(*sys.exc_info())
        if os.getpid() != _PID:
            # A child forked by the cell fell out of it: it must not take
            # part in the protocol, so it ends here.
            if error:
                sys.stderr.write(error)
            os._exit(1 if error else 0)
        try:
            names = user_names()[:500]
        except Exception:  # noqa: BLE001 -- the names are a courtesy, the result is not
            names = []
        msg = dict(t="done", result=result, error=error, seconds=round(time.time() - t0, 3),
                   names=names)
        if blob is not None:
            msg["blob"] = blob
        send(**msg)


if __name__ == "__main__":
    signal.signal(signal.SIGINT, on_sigint)
    sys.exit(main())