"""Remote kernel management and execution for Prime Agent."""

from __future__ import annotations

import ast
import asyncio
import base64
import json
import os
import time
import uuid
from pathlib import PurePosixPath
from typing import Any, AsyncIterator, Awaitable, Callable, Generator, Generic, Sequence, TypeVar

from rlm import host_request
from rlm.bash import BashResult

T = TypeVar("T")

KNOWN_PRIME_SKILLS: frozenset[str] = frozenset({
    "agent_message",
    "agent_observe",
    "atproto",
    "attach_image",
    "browser",
    "compact",
    "dependency_recon",
    "edit",
    "git",
    "goal",
    "ideonomy_plain",
    "ideonomy_rich",
    "jj",
    "linear",
    "lucid_code",
    "lucid_review",
    "nix",
    "notion",
    "prime_agent_update",
    "prime_intellect",
    "refine",
    "remote",
    "repo",
    "review_fanout",
    "rlm_heartbeat",
    "search",
    "skill_creator",
    "ssh_control",
    "tangled",
    "tangled_lexicon_codegen",
    "tangled_linux_validation",
    "tangled_spindle_ops",
    "threads",
    "websearch",
    "yt",
})


def _decode_response_text(raw: str) -> str:
    if raw.startswith(("'", '"')):
        try:
            val = ast.literal_eval(raw)
            if isinstance(val, str):
                return val
        except (ValueError, SyntaxError):
            pass
    return raw


def _coerce_path(path: Any, kernel: RemoteKernelHandle) -> str:
    if isinstance(path, RemotePath):
        if path._kernel is not kernel:
            raise ValueError("RemotePath belongs to a different remote connection")
        return str(path._path)
    if isinstance(path, os.PathLike):
        fspath = os.fspath(path)
        if isinstance(fspath, bytes):
            raise TypeError("Bytes paths are not supported")
        return str(fspath)
    if isinstance(path, bytes):
        raise TypeError("Bytes paths are not supported")
    if isinstance(path, str):
        return path
    raise TypeError(f"Expected str, os.PathLike, or RemotePath, got {type(path).__name__}")


class RemoteExecutionError(RuntimeError):
    """Raised when remote code execution fails."""

    def __init__(self, message: str, ename: str = "", evalue: str = "", traceback: list[str] | None = None):
        super().__init__(message)
        self.ename = ename
        self.evalue = evalue
        self.traceback = traceback or []


class RemoteHandle(Generic[T]):
    """Live handle to remote execution; behaves like BashHandle."""

    def __init__(
        self,
        req_id: str,
        task: asyncio.Task[T],
        interrupt_fn: Callable[[str], asyncio.Future[None] | None],
        stdout_chunks: list[str],
    ):
        self._req_id = req_id
        self._task = task
        self._interrupt_fn = interrupt_fn
        self._stdout_chunks = stdout_chunks
        self._started = time.monotonic()

    @property
    def running(self) -> bool:
        return not self._task.done()

    def output(self) -> str:
        return "".join(self._stdout_chunks)

    def tail(self, n: int = 50) -> str:
        lines = self.output().splitlines(keepends=True)
        return "".join(lines[-n:]) if lines else ""

    def poll(self) -> T | None:
        if self._task.done():
            return self._task.result()
        return None

    def kill(self) -> None:
        if not self._task.done():
            self._interrupt_fn(self._req_id)
            self._task.cancel()

    def __await__(self) -> Generator[Any, None, T]:
        return self._task.__await__()

    def __repr__(self) -> str:
        state = "running" if self.running else "done"
        return f"<RemoteHandle {self._req_id[:8]} ({state})>"


class RemoteSkillProxy:
    """Proxy for remote skill execution."""

    def __init__(self, kernel: RemoteKernelHandle, skill: str, parts: tuple[str, ...] = ()):
        self._kernel = kernel
        self._skill = skill
        self._parts = parts

    def __getattr__(self, name: str) -> RemoteSkillProxy:
        if name.startswith("_"):
            raise AttributeError(f"{type(self).__name__!r} object has no attribute {name!r}")
        return RemoteSkillProxy(self._kernel, self._skill, (*self._parts, name))

    def __call__(self, *args: Any, **kwargs: Any) -> RemoteHandle[Any]:
        return self._kernel._call_skill(self._skill, self._parts, *args, **kwargs)

    def __repr__(self) -> str:
        target = f"{self._skill}.{'.'.join(self._parts)}" if self._parts else self._skill
        return f"<RemoteSkillProxy {target}>"


class RemotePath:
    """Async remote path modeled after AnyIO Path."""

    def __init__(self, path: PathArg, kernel: RemoteKernelHandle):
        self._kernel = kernel
        if isinstance(path, RemotePath):
            if path._kernel is not kernel:
                raise ValueError("RemotePath belongs to a different remote connection")
            self._path = path._path
        elif isinstance(path, PurePosixPath):
            self._path = path
        else:
            self._path = PurePosixPath(_coerce_path(path, kernel))

    @property
    def name(self) -> str:
        return self._path.name

    @property
    def stem(self) -> str:
        return self._path.stem

    @property
    def suffix(self) -> str:
        return self._path.suffix

    @property
    def suffixes(self) -> list[str]:
        return self._path.suffixes

    @property
    def parent(self) -> RemotePath:
        return RemotePath(self._path.parent, self._kernel)

    @property
    def parts(self) -> tuple[str, ...]:
        return self._path.parts

    def with_name(self, name: str) -> RemotePath:
        return RemotePath(self._path.with_name(name), self._kernel)

    def with_suffix(self, suffix: str) -> RemotePath:
        return RemotePath(self._path.with_suffix(suffix), self._kernel)

    def with_stem(self, stem: str) -> RemotePath:
        return RemotePath(self._path.with_stem(stem), self._kernel)

    def joinpath(self, *other: PathArg) -> RemotePath:
        subs = [_coerce_path(o, self._kernel) for o in other]
        return RemotePath(self._path.joinpath(*subs), self._kernel)

    def __truediv__(self, other: PathArg) -> RemotePath:
        sub = _coerce_path(other, self._kernel)
        return RemotePath(self._path / sub, self._kernel)

    def __str__(self) -> str:
        return str(self._path)

    def __repr__(self) -> str:
        return f"RemotePath({str(self._path)!r}, host={self._kernel.host!r})"

    def __eq__(self, other: Any) -> bool:
        if isinstance(other, RemotePath):
            return self._kernel is other._kernel and self._path == other._path
        return NotImplemented

    def __hash__(self) -> int:
        return hash((self._path, id(self._kernel)))

    async def exists(self) -> bool:
        def _remote(p: str) -> bool:
            import os
            return os.path.exists(p)
        return await self._kernel.call(_remote, str(self._path))

    async def is_file(self) -> bool:
        def _remote(p: str) -> bool:
            import os
            return os.path.isfile(p)
        return await self._kernel.call(_remote, str(self._path))

    async def is_dir(self) -> bool:
        def _remote(p: str) -> bool:
            import os
            return os.path.isdir(p)
        return await self._kernel.call(_remote, str(self._path))

    async def read_text(self, encoding: str = "utf-8", errors: str = "strict") -> str:
        def _remote(p: str, enc: str, err: str) -> str:
            with open(p, "r", encoding=enc, errors=err) as f:
                return f.read()
        return await self._kernel.call(_remote, str(self._path), encoding, errors)

    async def read_bytes(self) -> bytes:
        def _remote(p: str) -> bytes:
            with open(p, "rb") as f:
                return f.read()
        return await self._kernel.call(_remote, str(self._path))

    async def write_text(self, data: str, encoding: str = "utf-8", errors: str = "strict") -> int:
        def _remote(p: str, content: str, enc: str, err: str) -> int:
            with open(p, "w", encoding=enc, errors=err) as f:
                return f.write(content)
        return await self._kernel.call(_remote, str(self._path), data, encoding, errors)

    async def write_bytes(self, data: bytes) -> int:
        def _remote(p: str, content: bytes) -> int:
            with open(p, "wb") as f:
                return f.write(content)
        return await self._kernel.call(_remote, str(self._path), data)

    async def stat(self) -> dict[str, Any]:
        def _remote(p: str) -> dict[str, Any]:
            import os
            st = os.stat(p)
            return {
                "st_size": st.st_size,
                "st_mode": st.st_mode,
                "st_mtime": st.st_mtime,
                "st_atime": st.st_atime,
                "st_ctime": st.st_ctime,
                "st_uid": st.st_uid,
                "st_gid": st.st_gid,
            }
        return await self._kernel.call(_remote, str(self._path))

    async def mkdir(self, mode: int = 0o777, parents: bool = False, exist_ok: bool = False) -> None:
        def _remote(p: str, m: int, par: bool, ok: bool) -> None:
            import pathlib
            pathlib.Path(p).mkdir(mode=m, parents=par, exist_ok=ok)
        await self._kernel.call(_remote, str(self._path), mode, parents, exist_ok)

    async def unlink(self, missing_ok: bool = False) -> None:
        def _remote(p: str, ok: bool) -> None:
            import pathlib
            pathlib.Path(p).unlink(missing_ok=ok)
        await self._kernel.call(_remote, str(self._path), missing_ok)

    async def rmdir(self) -> None:
        def _remote(p: str) -> None:
            import os
            os.rmdir(p)
        await self._kernel.call(_remote, str(self._path))

    async def resolve(self) -> RemotePath:
        def _remote(p: str) -> str:
            import pathlib
            return str(pathlib.Path(p).resolve())
        resolved = await self._kernel.call(_remote, str(self._path))
        return RemotePath(resolved, self._kernel)

    async def iterdir(self) -> AsyncIterator[RemotePath]:
        def _remote(p: str) -> list[str]:
            import os
            return os.listdir(p)
        entries = await self._kernel.call(_remote, str(self._path))
        for name in entries:
            yield RemotePath(self._path / name, self._kernel)

    async def glob(self, pattern: str) -> AsyncIterator[RemotePath]:
        def _remote(p: str, pat: str) -> list[str]:
            import pathlib
            return [str(match) for match in pathlib.Path(p).glob(pat)]
        matches = await self._kernel.call(_remote, str(self._path), pattern)
        for path_str in matches:
            yield RemotePath(path_str, self._kernel)


PathArg = str | os.PathLike[str] | RemotePath


class RemotePathlibFacade:
    """Facade for remote pathlib access."""

    def __init__(self, kernel: RemoteKernelHandle):
        self._kernel = kernel
        self.PurePosixPath = PurePosixPath

    def Path(self, *parts: PathArg) -> RemotePath:
        if not parts:
            return RemotePath(".", self._kernel)
        str_parts = [_coerce_path(p, self._kernel) for p in parts]
        return RemotePath(PurePosixPath(*str_parts), self._kernel)

    def __call__(self, *parts: PathArg) -> RemotePath:
        return self.Path(*parts)


class RemoteSysPathFacade:
    def __init__(self, kernel: RemoteKernelHandle):
        self._kernel = kernel

    async def get(self) -> list[str]:
        def _remote() -> list[str]:
            import sys
            return list(sys.path)
        return await self._kernel.call(_remote)

    async def append(self, dir_path: PathArg) -> None:
        p = _coerce_path(dir_path, self._kernel)
        def _remote(target: str) -> None:
            import sys
            if target not in sys.path:
                sys.path.append(target)
        await self._kernel.call(_remote, p)

    async def prepend(self, dir_path: PathArg) -> None:
        p = _coerce_path(dir_path, self._kernel)
        def _remote(target: str) -> None:
            import sys
            if target in sys.path:
                sys.path.remove(target)
            sys.path.insert(0, target)
        await self._kernel.call(_remote, p)


class RemoteSysFacade:
    """Remote process metadata and state facade."""

    def __init__(self, kernel: RemoteKernelHandle):
        self._kernel = kernel
        self.path = RemoteSysPathFacade(kernel)

    @property
    def platform(self) -> str:
        return self._kernel._sys_platform

    @property
    def executable(self) -> str:
        return self._kernel._sys_executable

    @property
    def version(self) -> str:
        return self._kernel._sys_version


class RemoteEnvironFacade:
    def __init__(self, kernel: RemoteKernelHandle):
        self._kernel = kernel

    async def get(self, key: str, default: str | None = None) -> str | None:
        def _remote(k: str, d: str | None) -> str | None:
            import os
            return os.environ.get(k, d)
        return await self._kernel.call(_remote, key, default)

    async def set(self, key: str, val: str) -> None:
        def _remote(k: str, v: str) -> None:
            import os
            os.environ[k] = str(v)
        await self._kernel.call(_remote, key, val)

    async def delete(self, key: str) -> None:
        def _remote(k: str) -> None:
            import os
            os.environ.pop(k, None)
        await self._kernel.call(_remote, key)

    async def update(self, mapping: dict[str, str]) -> None:
        def _remote(m: dict[str, str]) -> None:
            import os
            os.environ.update({str(k): str(v) for k, v in m.items()})
        await self._kernel.call(_remote, mapping)

    async def snapshot(self) -> dict[str, str]:
        def _remote() -> dict[str, str]:
            import os
            return dict(os.environ)
        return await self._kernel.call(_remote)


class RemoteOsFacade:
    """Remote process and filesystem primitives facade."""

    def __init__(self, kernel: RemoteKernelHandle):
        self._kernel = kernel
        self.environ = RemoteEnvironFacade(kernel)

    async def chdir(self, path: PathArg) -> None:
        p = _coerce_path(path, self._kernel)
        def _remote(target: str) -> None:
            import os
            os.chdir(target)
        await self._kernel.call(_remote, p)

    async def getcwd(self) -> str:
        def _remote() -> str:
            import os
            return os.getcwd()
        return await self._kernel.call(_remote)

    async def listdir(self, path: PathArg = ".") -> list[str]:
        p = _coerce_path(path, self._kernel)
        def _remote(target: str) -> list[str]:
            import os
            return os.listdir(target)
        return await self._kernel.call(_remote, p)

    async def mkdir(self, path: PathArg, mode: int = 0o777) -> None:
        p = _coerce_path(path, self._kernel)
        def _remote(target: str, m: int) -> None:
            import os
            os.mkdir(target, mode=m)
        await self._kernel.call(_remote, p, mode)

    async def makedirs(self, path: PathArg, mode: int = 0o777, exist_ok: bool = False) -> None:
        p = _coerce_path(path, self._kernel)
        def _remote(target: str, m: int, ok: bool) -> None:
            import os
            os.makedirs(target, mode=m, exist_ok=ok)
        await self._kernel.call(_remote, p, mode, exist_ok)

    async def remove(self, path: PathArg) -> None:
        p = _coerce_path(path, self._kernel)
        def _remote(target: str) -> None:
            import os
            os.remove(target)
        await self._kernel.call(_remote, p)


class RemoteShutilFacade:
    """Remote high-level filesystem primitives facade."""

    def __init__(self, kernel: RemoteKernelHandle):
        self._kernel = kernel

    async def copy(self, src: PathArg, dst: PathArg) -> str:
        s = _coerce_path(src, self._kernel)
        d = _coerce_path(dst, self._kernel)
        def _remote(src_p: str, dst_p: str) -> str:
            import shutil
            return shutil.copy(src_p, dst_p)
        return await self._kernel.call(_remote, s, d)

    async def copytree(self, src: PathArg, dst: PathArg, **kwargs: Any) -> str:
        s = _coerce_path(src, self._kernel)
        d = _coerce_path(dst, self._kernel)
        def _remote(src_p: str, dst_p: str, kw: dict[str, Any]) -> str:
            import shutil
            return shutil.copytree(src_p, dst_p, **kw)
        return await self._kernel.call(_remote, s, d, kwargs)

    async def rmtree(self, path: PathArg, ignore_errors: bool = False) -> None:
        p = _coerce_path(path, self._kernel)
        def _remote(target: str, ign: bool) -> None:
            import shutil
            shutil.rmtree(target, ignore_errors=ign)
        await self._kernel.call(_remote, p, ignore_errors)

    async def move(self, src: PathArg, dst: PathArg) -> str:
        s = _coerce_path(src, self._kernel)
        d = _coerce_path(dst, self._kernel)
        def _remote(src_p: str, dst_p: str) -> str:
            import shutil
            return shutil.move(src_p, dst_p)
        return await self._kernel.call(_remote, s, d)

    async def which(self, cmd: str) -> str | None:
        def _remote(c: str) -> str | None:
            import shutil
            return shutil.which(c)
        return await self._kernel.call(_remote, cmd)


class RemoteControlPlane:
    """Connection and execution plumbing for a remote kernel."""

    def __init__(self, kernel: RemoteKernelHandle):
        self._kernel = kernel

    @property
    def capabilities(self) -> frozenset[str]:
        return frozenset(self._kernel._capabilities)

    def execute(self, code: str) -> RemoteHandle[Any]:
        return self._kernel._execute_code(code)

    def eval(self, expr: str) -> RemoteHandle[Any]:
        return self._kernel._eval_expr(expr)

    def ship(self, fn: Callable, *args: Any, **kwargs: Any) -> RemoteHandle[Any]:
        return self._kernel._call_fn(fn, *args, **kwargs)

    call = ship

    async def aclose(self) -> None:
        await self._kernel.close()

    close = aclose

    async def refresh_capabilities(self) -> frozenset[str]:
        return await self._kernel._discover_capabilities()

    @property
    def closed(self) -> bool:
        return self._kernel.closed


class RemoteKernelHandle:
    """Active connection to an IPython/RLM kernel running on a remote host."""

    def __init__(
        self,
        host: str,
        proc: asyncio.subprocess.Process,
        remote_python: str,
        remote_dir: str,
        capabilities: Sequence[str] | set[str] | None = None,
    ):
        self.host = host
        self._proc = proc
        self.remote_python = remote_python
        self.remote_dir = remote_dir or ""
        self._closed = False
        self._lock = asyncio.Lock()
        self._imported_modules: set[str] = set()
        self._capabilities: set[str] = set(capabilities) if capabilities is not None else set()
        self._sys_platform: str = ""
        self._sys_executable: str = self.remote_python
        self._sys_version: str = ""
        self._close_task: asyncio.Task[None] | None = None
        self.control = RemoteControlPlane(self)
        self.pathlib = RemotePathlibFacade(self)
        self.sys = RemoteSysFacade(self)
        self.os = RemoteOsFacade(self)
        self.shutil = RemoteShutilFacade(self)

    @property
    def closed(self) -> bool:
        return self._closed or self._proc.returncode is not None

    async def import_(self, *modules: str, reload: bool = False) -> RemoteKernelHandle:
        lines: list[str] = []
        new_modules: list[str] = []
        for mod in modules:
            mod = mod.strip()
            if not mod:
                continue
            if not reload and mod in self._imported_modules:
                continue
            if reload:
                base_name = mod.split()[0]
                lines.append(f"import sys; sys.modules.pop({repr(base_name)}, None)")
            if mod.startswith("import ") or mod.startswith("from "):
                lines.append(mod)
            else:
                lines.append(f"import {mod}")
            new_modules.append(mod)

        if lines:
            import_code = "\n".join(lines)
            await self._execute_raw(import_code)
            self._imported_modules.update(new_modules)
        return self

    imports = import_

    def call(self, fn: Callable, *args: Any, **kwargs: Any) -> RemoteHandle[Any]:
        return self.control.ship(fn, *args, **kwargs)

    ship = call

    def bash(self, command: str, input: Any = None) -> RemoteHandle[BashResult]:
        if hasattr(input, "text") and callable(input.text):
            input_payload = input.text()
        elif hasattr(input, "output"):
            input_payload = input.output() if callable(input.output) else input.output
        else:
            input_payload = input
        code = f"""
import json
from rlm.bash import bash
_h = bash({repr(command)}, input={repr(input_payload)})
_res = await _h
json.dumps({{"exit_code": _res.exit_code, "output": _res.output, "duration": _res.duration, "stdout": _res.stdout, "stderr": _res.stderr}})
"""
        req_id = uuid.uuid4().hex
        stdout_chunks: list[str] = []

        async def _run() -> BashResult:
            raw = await self._execute_request(req_id, code, stdout_chunks)
            raw = _decode_response_text(raw)
            data = json.loads(raw)
            return BashResult(
                exit_code=data["exit_code"],
                output=data["output"],
                duration=data.get("duration", 0.0),
                stdout=data.get("stdout", data["output"]),
                stderr=data.get("stderr", ""),
            )

        task = asyncio.create_task(_run())
        return RemoteHandle(
            req_id=req_id,
            task=task,
            interrupt_fn=self._send_interrupt,
            stdout_chunks=stdout_chunks,
        )

    def _execute_code(self, code: str) -> RemoteHandle[Any]:
        req_id = uuid.uuid4().hex
        stdout_chunks: list[str] = []

        async def _run() -> Any:
            raw = await self._execute_request(req_id, code, stdout_chunks)
            return _decode_response_text(raw)

        task = asyncio.create_task(_run())
        return RemoteHandle(
            req_id=req_id,
            task=task,
            interrupt_fn=self._send_interrupt,
            stdout_chunks=stdout_chunks,
        )

    def _eval_expr(self, expr: str) -> RemoteHandle[Any]:
        code = f"""
import dill, base64, inspect, json
_val = eval({repr(expr)}, globals())
if inspect.isawaitable(_val):
    _val = await _val
try:
    _dumped = json.dumps({{"status": "ok", "result": _val}})
except Exception:
    _dumped = json.dumps({{"status": "dill", "result": base64.b64encode(dill.dumps(_val)).decode()}})
_dumped
"""
        req_id = uuid.uuid4().hex
        stdout_chunks: list[str] = []

        async def _run() -> Any:
            raw = await self._execute_request(req_id, code, stdout_chunks)
            raw = _decode_response_text(raw)
            parsed = json.loads(raw)
            status = parsed.get("status")
            if status == "ok":
                return parsed.get("result")
            if status == "dill":
                import dill
                return dill.loads(base64.b64decode(parsed.get("result")))
            return parsed

        task = asyncio.create_task(_run())
        return RemoteHandle(
            req_id=req_id,
            task=task,
            interrupt_fn=self._send_interrupt,
            stdout_chunks=stdout_chunks,
        )

    def _call_fn(self, fn: Callable, *args: Any, **kwargs: Any) -> RemoteHandle[Any]:
        import dill
        import importlib
        import types

        def _save_mod(pickler, obj):
            pickler.save_reduce(importlib.import_module, (obj.__name__,), obj=obj)

        try:
            dill.register(types.ModuleType)(_save_mod)
            for sub in types.ModuleType.__subclasses__():
                dill.register(sub)(_save_mod)
        except Exception:
            pass

        is_builtin = isinstance(fn, (types.BuiltinFunctionType, types.BuiltinMethodType))
        orig_module = getattr(fn, "__module__", None)
        can_patch_module = not is_builtin and orig_module not in (None, "builtins", "posix", "nt", "os")
        try:
            if can_patch_module:
                fn.__module__ = "__main__"
            fn_b64 = base64.b64encode(dill.dumps(fn, recurse=True)).decode()
        finally:
            if can_patch_module and orig_module is not None:
                try:
                    fn.__module__ = orig_module
                except Exception:
                    pass

        args_b64 = base64.b64encode(dill.dumps(args, recurse=True)).decode()
        kwargs_b64 = base64.b64encode(dill.dumps(kwargs, recurse=True)).decode()

        code = f"""
import dill, base64, inspect, json
_fn = dill.loads(base64.b64decode({repr(fn_b64)}))
_args = dill.loads(base64.b64decode({repr(args_b64)}))
_kwargs = dill.loads(base64.b64decode({repr(kwargs_b64)}))

if hasattr(_fn, "__globals__"):
    _fn.__globals__.update(globals())
    if "bash" not in _fn.__globals__:
        try:
            from rlm.bash import bash as _prime_agent_bash
            _fn.__globals__["bash"] = _prime_agent_bash
        except ImportError:
            pass

_res = _fn(*_args, **_kwargs)
if inspect.isawaitable(_res):
    _res = await _res

try:
    _dumped = json.dumps({{"status": "ok", "result": _res}})
except Exception:
    _dumped = json.dumps({{"status": "dill", "result": base64.b64encode(dill.dumps(_res)).decode()}})
_dumped
"""
        req_id = uuid.uuid4().hex
        stdout_chunks: list[str] = []

        async def _run() -> Any:
            raw = await self._execute_request(req_id, code, stdout_chunks)
            raw = _decode_response_text(raw)
            parsed = json.loads(raw)
            if parsed.get("status") == "ok":
                return parsed.get("result")
            if parsed.get("status") == "dill":
                import dill
                return dill.loads(base64.b64decode(parsed.get("result")))
            return parsed

        task = asyncio.create_task(_run())
        return RemoteHandle(
            req_id=req_id,
            task=task,
            interrupt_fn=self._send_interrupt,
            stdout_chunks=stdout_chunks,
        )

    def _call_skill(self, skill: str, parts: tuple[str, ...], *args: Any, **kwargs: Any) -> RemoteHandle[Any]:
        import dill
        import importlib
        import types

        def _save_mod(pickler, obj):
            pickler.save_reduce(importlib.import_module, (obj.__name__,), obj=obj)

        try:
            dill.register(types.ModuleType)(_save_mod)
            for sub in types.ModuleType.__subclasses__():
                dill.register(sub)(_save_mod)
        except Exception:
            pass

        args_b64 = base64.b64encode(dill.dumps(args, recurse=True)).decode()
        kwargs_b64 = base64.b64encode(dill.dumps(kwargs, recurse=True)).decode()

        code = f"""
import dill, base64, inspect, json, importlib

_mod = importlib.import_module({repr(skill)})
_target = _mod
for _p in {repr(parts)}:
    _target = getattr(_target, _p)

if not callable(_target) and hasattr(_target, "run") and callable(_target.run):
    _fn = _target.run
else:
    _fn = _target

_args = dill.loads(base64.b64decode({repr(args_b64)}))
_kwargs = dill.loads(base64.b64decode({repr(kwargs_b64)}))

_res = _fn(*_args, **_kwargs)
if inspect.isawaitable(_res):
    _res = await _res

try:
    _dumped = json.dumps({{"status": "ok", "result": _res}})
except Exception:
    _dumped = json.dumps({{"status": "dill", "result": base64.b64encode(dill.dumps(_res)).decode()}})
_dumped
"""
        req_id = uuid.uuid4().hex
        stdout_chunks: list[str] = []

        async def _run() -> Any:
            raw = await self._execute_request(req_id, code, stdout_chunks)
            raw = _decode_response_text(raw)
            parsed = json.loads(raw)
            if parsed.get("status") == "ok":
                return parsed.get("result")
            if parsed.get("status") == "dill":
                import dill
                return dill.loads(base64.b64decode(parsed.get("result")))
            return parsed

        task = asyncio.create_task(_run())
        return RemoteHandle(
            req_id=req_id,
            task=task,
            interrupt_fn=self._send_interrupt,
            stdout_chunks=stdout_chunks,
        )

    async def _discover_capabilities(self) -> frozenset[str]:
        if self.closed:
            return frozenset(self._capabilities)
        code = f"""
import importlib.util, json, os, sys

_known = {sorted(list(KNOWN_PRIME_SKILLS))!r}
_skills = set()

for _name in _known:
    try:
        if _name in sys.modules or importlib.util.find_spec(_name) is not None:
            _skills.add(_name)
    except Exception:
        pass

_remote_dir = os.path.expandvars(os.path.expanduser({repr(self.remote_dir)}))
if _remote_dir:
    _skills_path = os.path.join(_remote_dir, "skills")
    if os.path.isdir(_skills_path):
        for _entry in os.listdir(_skills_path):
            if not _entry.startswith((".", "_")) and os.path.isdir(os.path.join(_skills_path, _entry)):
                _norm = _entry.replace("-", "_")
                try:
                    if importlib.util.find_spec(_norm) is not None:
                        _skills.add(_norm)
                except Exception:
                    pass

json.dumps({{
    "skills": sorted(list(_skills)),
    "platform": sys.platform,
    "executable": sys.executable,
    "version": sys.version,
}})
"""
        try:
            raw = await self._execute_raw(code)
            raw = _decode_response_text(raw)
            data = json.loads(raw)
            if isinstance(data, dict):
                skills = data.get("skills", [])
                if isinstance(skills, list):
                    self._capabilities = set(skills)
                self._sys_platform = data.get("platform", "")
                self._sys_executable = data.get("executable", self.remote_python)
                self._sys_version = data.get("version", "")
            elif isinstance(data, list):
                self._capabilities = set(data)
        except Exception:
            pass
        return frozenset(self._capabilities)

    def __getattr__(self, name: str) -> RemoteSkillProxy:
        if name.startswith("_"):
            raise AttributeError(f"{type(self).__name__!r} object has no attribute {name!r}")
        if name not in self._capabilities:
            raise AttributeError(
                f"'{name}' is not an available remote skill. Available: {sorted(self._capabilities)}"
            )
        return RemoteSkillProxy(self, name)

    def __getitem__(self, skill_name: str) -> RemoteSkillProxy:
        norm = skill_name.replace("-", "_")
        if norm in self._capabilities:
            return RemoteSkillProxy(self, norm)
        if skill_name in self._capabilities:
            return RemoteSkillProxy(self, skill_name)
        raise KeyError(
            f"'{skill_name}' is not an available remote skill. Available: {sorted(self._capabilities)}"
        )

    def __contains__(self, key: str) -> bool:
        norm = key.replace("-", "_")
        return norm in self._capabilities or key in self._capabilities

    def __dir__(self) -> list[str]:
        attrs = set(super().__dir__())
        attrs.update(self._capabilities)
        return sorted(attrs)

    def _send_interrupt(self, req_id: str) -> None:
        if self.closed or not self._proc.stdin:
            return
        msg = (json.dumps({"type": "interrupt", "id": req_id}) + "\n").encode("utf-8")
        try:
            self._proc.stdin.write(msg)
        except Exception:
            pass

    async def _execute_raw(self, code: str) -> str:
        req_id = uuid.uuid4().hex
        stdout_chunks: list[str] = []
        return await self._execute_request(req_id, code, stdout_chunks)

    async def _dispatch_host_request(self, event: dict[str, Any]) -> None:
        h_id = event.get("id")
        h_data = event.get("data", {})
        try:
            reply_res = await host_request(h_data.get("type", "unknown"), h_data)
            reply = {"type": "host_reply", "id": h_id, "data": {"status": "ok", "result": reply_res}}
        except Exception as exc:
            reply = {"type": "host_reply", "id": h_id, "data": {"status": "error", "error": str(exc)}}
        if self._proc.stdin and not self._proc.stdin.is_closing():
            self._proc.stdin.write((json.dumps(reply) + "\n").encode("utf-8"))
            await self._proc.stdin.drain()

    async def _drain_until_done(self, req_id: str) -> None:
        assert self._proc.stdout is not None
        while True:
            line = await self._proc.stdout.readline()
            if not line:
                raise RuntimeError(f"Remote kernel on {self.host} closed unexpectedly during drain")
            try:
                event = json.loads(line.decode("utf-8"))
            except Exception:
                continue

            ev_type = event.get("event")
            if ev_type == "host_request":
                await self._dispatch_host_request(event)
                continue

            if event.get("id") == req_id and ev_type == "done":
                return

    async def _execute_request(self, req_id: str, code: str, stdout_chunks: list[str]) -> str:
        if self.closed:
            raise RuntimeError(f"Remote kernel connection to {self.host} is closed")

        async with self._lock:
            req = {"type": "execute", "id": req_id, "code": code}
            data = (json.dumps(req) + "\n").encode("utf-8")

            assert self._proc.stdin is not None
            assert self._proc.stdout is not None

            result_text: str | None = None
            error_data: dict[str, Any] | None = None
            submitted = False

            try:
                self._proc.stdin.write(data)
                submitted = True
                await self._proc.stdin.drain()

                while True:
                    line = await self._proc.stdout.readline()
                    if not line:
                        raise RuntimeError(f"Remote kernel on {self.host} closed unexpectedly during execution")

                    try:
                        event = json.loads(line.decode("utf-8"))
                    except Exception:
                        continue

                    ev_type = event.get("event")
                    ev_id = event.get("id")

                    if ev_type == "host_request":
                        await self._dispatch_host_request(event)
                        continue

                    if ev_id != req_id:
                        continue

                    if ev_type == "stdout":
                        text = event.get("text", "")
                        stdout_chunks.append(text)
                    elif ev_type == "result":
                        result_text = event.get("text", "")
                    elif ev_type == "error":
                        error_data = event
                    elif ev_type == "done":
                        break
            except asyncio.CancelledError:
                if submitted:
                    async def _cancel_cleanup() -> None:
                        self._send_interrupt(req_id)
                        try:
                            await asyncio.wait_for(self._drain_until_done(req_id), timeout=3.0)
                        except (asyncio.TimeoutError, Exception):
                            await self._poison_and_close()

                    cleanup_task = asyncio.create_task(_cancel_cleanup())
                    cancelled = True
                    while not cleanup_task.done():
                        try:
                            await asyncio.shield(cleanup_task)
                        except asyncio.CancelledError:
                            cancelled = True
                    if cancelled:
                        raise asyncio.CancelledError()
                raise

            if error_data is not None:
                ename = error_data.get("ename", "Error")
                evalue = error_data.get("evalue", "")
                tb = error_data.get("traceback", [])
                msg = f"{ename}: {evalue}\n" + "".join(tb)
                raise RemoteExecutionError(msg, ename=ename, evalue=evalue, traceback=tb)

            if result_text is not None:
                return result_text
            return "".join(stdout_chunks)

    async def _do_close(self) -> None:
        try:
            if self._proc.stdin:
                self._proc.stdin.close()
                await self._proc.stdin.wait_closed()
        except Exception:
            pass
        try:
            self._proc.terminate()
            await asyncio.wait_for(self._proc.wait(), timeout=2.0)
        except Exception:
            try:
                self._proc.kill()
                await self._proc.wait()
            except Exception:
                pass

    async def _poison_and_close(self) -> None:
        self._closed = True
        if self._close_task is None:
            self._close_task = asyncio.create_task(self._do_close())
        task = self._close_task
        cancelled = False
        while not task.done():
            try:
                await asyncio.shield(task)
            except asyncio.CancelledError:
                cancelled = True
        if cancelled:
            raise asyncio.CancelledError()

    async def close(self) -> None:
        await self._poison_and_close()

    aclose = close

    async def __aenter__(self) -> RemoteKernelHandle:
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        await self.close()


class _ConnectContext:
    def __init__(self, coro: Awaitable[RemoteKernelHandle]):
        self._coro = coro
        self._handle: RemoteKernelHandle | None = None

    def __await__(self) -> Generator[Any, None, RemoteKernelHandle]:
        return self._coro.__await__()

    async def __aenter__(self) -> RemoteKernelHandle:
        self._handle = await self._coro
        return await self._handle.__aenter__()

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        if self._handle is not None:
            await self._handle.__aexit__(exc_type, exc_val, exc_tb)


def connect(
    host: str,
    *,
    port: int | None = None,
    identity_file: str | None = None,
    ssh_options: Sequence[str] | None = None,
    cache_dir: str | None = None,
    python_cmd: str | None = None,
) -> _ConnectContext:
    """Connect to or spawn an on-demand remote kernel on a remote host via SSH."""

    async def _do_connect() -> RemoteKernelHandle:
        ensure_payload: dict[str, Any] = {"host": host}
        if port is not None:
            ensure_payload["port"] = port
        if identity_file is not None:
            ensure_payload["identityFile"] = identity_file
        if ssh_options:
            ensure_payload["sshOptions"] = list(ssh_options)
        if cache_dir is not None:
            ensure_payload["cacheDir"] = cache_dir
        if python_cmd is not None:
            ensure_payload["pythonCmd"] = python_cmd

        info = await host_request("remote.ensure", ensure_payload)
        remote_python = info.get("remotePython", "python3")
        remote_dir = info.get("remoteDir", "")

        ssh_args = [
            "ssh",
            "-o", "ControlMaster=auto",
            "-o", "ControlPersist=30m",
            "-o", "ControlPath=~/.ssh/cm-%C",
        ]
        if port is not None:
            ssh_args.extend(["-p", str(port)])
        if identity_file is not None:
            ssh_args.extend(["-i", identity_file])
        if ssh_options:
            ssh_args.extend(ssh_options)
        ssh_args.append(host)
        ssh_args.append(f"{remote_python} -m rlm.repl")

        proc = await asyncio.create_subprocess_exec(
            *ssh_args,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )

        try:
            assert proc.stdout is not None
            ready_line = await proc.stdout.readline()
            if not ready_line:
                stderr_out = await proc.stderr.read() if proc.stderr else b""
                raise RuntimeError(f"Failed to start remote kernel on {host}: {stderr_out.decode('utf-8')}")

            try:
                ready_data = json.loads(ready_line.decode("utf-8"))
                if ready_data.get("event") != "ready":
                    raise RuntimeError(f"Unexpected initial message from remote kernel on {host}: {ready_line.decode('utf-8')}")
            except json.JSONDecodeError:
                raise RuntimeError(f"Corrupt startup output from remote kernel on {host}: {ready_line.decode('utf-8')}")

            handle = RemoteKernelHandle(host=host, proc=proc, remote_python=remote_python, remote_dir=remote_dir)
            try:
                await handle.control.refresh_capabilities()
            except Exception:
                pass
            return handle
        except BaseException:
            try:
                proc.kill()
                await proc.wait()
            except BaseException:
                pass
            raise

    return _ConnectContext(_do_connect())


__all__ = [
    "KNOWN_PRIME_SKILLS",
    "PathArg",
    "RemoteControlPlane",
    "RemoteEnvironFacade",
    "RemoteExecutionError",
    "RemoteHandle",
    "RemoteKernelHandle",
    "RemoteOsFacade",
    "RemotePath",
    "RemotePathlibFacade",
    "RemoteShutilFacade",
    "RemoteSkillProxy",
    "RemoteSysFacade",
    "RemoteSysPathFacade",
    "connect",
]