"""
NVIDIA OpenShell sandbox launcher.

Implements :class:`~omnigent.onboarding.sandboxes.base.SandboxLauncher`
for `NVIDIA OpenShell <https://github.com/NVIDIA/openshell>`_ sandboxes
on top of the official ``openshell`` Python SDK. Same posture as the
Modal, Daytona, and CoreWeave launchers: the SDK is an optional
dependency (``pip install 'omnigent[openshell]'``) imported lazily, so
the provider can be listed and the module probed without it.

OpenShell is self-hosted: a gateway control plane manages sandbox
lifecycle on a configured compute driver (Docker, Podman, microVM,
Kubernetes), filling the gap for on-prem and air-gapped deployments
where cloud providers are unavailable.

Notes that shape this launcher:

- **gRPC, not REST.** OpenShell's control plane is a gRPC service; the
  ``openshell`` SDK wraps it. The launcher connects through
  :meth:`SandboxClient.from_active_cluster`, which resolves the active
  gateway (its endpoint, TLS material, and OIDC token) from the gateway
  selected by ``openshell gateway select`` — i.e. ``$OPENSHELL_GATEWAY``
  or ``~/.config/openshell/active_gateway``. There is no base-URL knob.
- **Custom host image.** Omnigent boots its prebaked host image, which
  rides in ``SandboxSpec.template.image``. The SDK's public ``create``
  takes only a ``SandboxSpec`` and does not re-export the spec
  protobufs, so the spec is built from the generated ``openshell._proto``
  module — the only path to a non-default image.
- **No file-transfer RPC.** OpenShell exposes command execution but no
  upload primitive, so :meth:`put` streams the file's bytes to ``cat``
  over the exec channel's stdin — the same approach NVIDIA's own
  LangChain backend uses.
- **No local port forwarding.** OpenShell has no local→sandbox port
  forward for the in-sandbox App OAuth callback, so the CLI skips that
  auth step automatically (``supports_local_port_forward = False``).
"""

from __future__ import annotations

import logging
import os
import shlex
from collections.abc import Callable, Sequence
from pathlib import PurePosixPath
from typing import TYPE_CHECKING, ClassVar, TypeVar

import click

from omnigent.onboarding.sandboxes.base import (
    DEFAULT_HOST_IMAGE,
    RemoteCommandResult,
    SandboxLauncher,
    host_image_wheel_install_command,
)

if TYPE_CHECKING:
    from pathlib import Path

    from openshell import ExecResult

_logger = logging.getLogger(__name__)

HOST_IMAGE_ENV_VAR: str = "OMNIGENT_OPENSHELL_HOST_IMAGE"
"""Environment variable overriding :data:`DEFAULT_HOST_IMAGE` for
OpenShell sandboxes."""

SANDBOX_ENV_PASSTHROUGH_ENV_VAR: str = "OMNIGENT_OPENSHELL_SANDBOX_ENV"
"""Comma-separated server-process environment variable names injected
into created OpenShell sandboxes."""

GATEWAY_ENV_VAR: str = "OPENSHELL_GATEWAY"
"""Gateway name read by the SDK's :meth:`SandboxClient.from_active_cluster`;
overrides ``~/.config/openshell/active_gateway``."""

_READY_TIMEOUT_S = 300
_EXEC_TIMEOUT_S = 300
# A foreground host (`omnigent host`) is held open until Ctrl-C, so its
# exec stream must not hit a gRPC deadline mid-session — give it a long
# ceiling. The pidfile records the in-sandbox pid so Ctrl-C can kill the
# remote process (cancelling the local stream doesn't stop it).
_FOREGROUND_TIMEOUT_S = 7 * 24 * 3600
_FOREGROUND_PIDFILE_TEMPLATE = "/tmp/oa-openshell-foreground-{sandbox_id}.pid"

# OpenShell runs the agent as the non-root ``sandbox`` user (its image
# contract; see deploy/docker/Dockerfile), whose home is ``/sandbox``.
# The host image keeps ``WORKDIR /root`` for the root-based providers, so we
# pin every exec's cwd + ``$HOME`` to the sandbox user's writable home here
# rather than changing the shared image -- otherwise ``omnigent host`` resolves
# its config under ``/root`` (unreadable to the sandbox user) and crashes, and
# the managed flow's ``$HOME/workspace`` lands somewhere unwritable.
# ``/home/sandbox`` is denied by the k8s Landlock LSM policy; ``/sandbox`` is
# the permitted path.
_SANDBOX_HOME = "/sandbox"

_T = TypeVar("_T")


def _ensure_sdk() -> None:
    """Verify the openshell SDK is importable, with an install hint when not."""
    try:
        import openshell  # noqa: F401
    except ImportError as exc:
        raise click.ClickException(
            "The openshell SDK is required for the 'openshell' sandbox provider. "
            "Install it with `pip install 'omnigent[openshell]'`, then select a "
            "gateway with `openshell gateway select <name>` (or set OPENSHELL_GATEWAY)."
        ) from exc


class _OpenShellClient:
    """Thin wrapper over the ``openshell`` gRPC SandboxClient.

    Owns the launcher's contact with the SDK: it builds the create
    spec, maps the public sandbox name back to the id ``exec`` needs,
    and translates SDK / gRPC errors into ``click.ClickException`` so
    the launcher surface stays clean.
    """

    def __init__(self, *, cluster: str | None = None) -> None:
        _ensure_sdk()
        from openshell import SandboxClient, SandboxError

        try:
            self._client: SandboxClient = SandboxClient.from_active_cluster(cluster=cluster)
        except SandboxError as exc:
            raise click.ClickException(
                f"Could not connect to an OpenShell gateway: {exc}. Select one with "
                "`openshell gateway select <name>` (or set OPENSHELL_GATEWAY)."
            ) from exc
        # Petname (public handle) -> opaque sandbox id, which exec needs.
        self._ids: dict[str, str] = {}
        # Daemon threads holding long-lived exec streams open (see
        # exec_background): OpenShell kills an exec's processes when the
        # ExecSandbox RPC returns, so a backgrounded host must be kept on
        # an open stream for its lifetime.
        self._bg_threads: list[object] = []

    def close(self) -> None:
        """Release the gRPC channel and any bearer-auth resources."""
        self._client.close()

    def create_sandbox(self, *, image: str, env: dict[str, str]) -> str:
        """Create a sandbox from *image*, wait until ready, return its name."""
        from openshell._proto import openshell_pb2

        spec = openshell_pb2.SandboxSpec(
            template=openshell_pb2.SandboxTemplate(image=image),
            environment=env or {},
        )
        ref = self._guard(
            "OpenShell sandbox creation failed",
            lambda: self._client.create(spec=spec),
        )
        ready = self._guard(
            "OpenShell sandbox did not become ready",
            lambda: self._client.wait_ready(ref.name, timeout_seconds=_READY_TIMEOUT_S),
        )
        sandbox_name: str = ready.name
        self._ids[sandbox_name] = ready.id
        return sandbox_name

    def execute(
        self,
        name: str,
        command: Sequence[str],
        *,
        stdin: bytes | None = None,
        timeout: int = _EXEC_TIMEOUT_S,
    ) -> ExecResult:
        """Run *command* (argv) in the sandbox and return its ExecResult."""
        sandbox_id = self._id_for(name)
        return self._guard(
            f"Remote command failed on OpenShell sandbox '{name}'",
            lambda: self._client.exec(
                sandbox_id,
                command,
                stdin=stdin,
                timeout_seconds=timeout,
                workdir=_SANDBOX_HOME,
                env={"HOME": _SANDBOX_HOME},
            ),
        )

    def run_foreground(self, name: str, command: Sequence[str], *, timeout: int) -> int:
        """Stream a command's output to the terminal; return its exit code."""
        import grpc
        from openshell import ExecChunk, ExecResult, SandboxError

        sandbox_id = self._id_for(name)
        exit_code = 0
        try:
            for item in self._client.exec_stream(
                sandbox_id,
                command,
                timeout_seconds=timeout,
                workdir=_SANDBOX_HOME,
                env={"HOME": _SANDBOX_HOME},
            ):
                if isinstance(item, ExecResult):
                    exit_code = int(item.exit_code)
                elif isinstance(item, ExecChunk):
                    click.echo(
                        item.data.decode("utf-8", errors="replace"),
                        nl=False,
                        err=item.stream != "stdout",
                    )
        except (grpc.RpcError, SandboxError) as exc:
            raise click.ClickException(
                f"Foreground command failed on OpenShell sandbox '{name}': {exc}"
            ) from exc
        return exit_code

    def exec_background(self, name: str, command: Sequence[str], *, timeout: int) -> None:
        """
        Start a long-lived foreground command, holding its exec stream open.

        OpenShell terminates an exec's process tree when the ``ExecSandbox``
        RPC returns, so the usual ``setsid nohup … &`` detach (which works on
        Modal/CoreWeave) is reaped immediately. Instead we run the command in
        the FOREGROUND of an ``exec_stream`` drained on a daemon thread: the
        stream stays open for the process's lifetime, so OpenShell keeps it
        alive. The thread ends when the process exits or the sandbox is
        deleted (the stream then errors out).
        """
        import threading

        sandbox_id = self._id_for(name)

        def _pump() -> None:
            try:
                for _ in self._client.exec_stream(
                    sandbox_id,
                    command,
                    timeout_seconds=timeout,
                    workdir=_SANDBOX_HOME,
                    env={"HOME": _SANDBOX_HOME},
                ):
                    pass
            except Exception:
                _logger.debug(
                    "exec_background stream ended for sandbox %s",
                    name,
                    exc_info=True,
                )

        self._bg_threads = [t for t in self._bg_threads if t.is_alive()]
        thread = threading.Thread(target=_pump, name=f"openshell-host-{name}", daemon=True)
        thread.start()
        self._bg_threads.append(thread)

    def get_status(self, name: str) -> None:
        """Resolve a sandbox by name (validates access) and cache its id."""
        ref = self._guard(
            f"Could not resolve OpenShell sandbox '{name}'",
            lambda: self._client.get(name),
        )
        self._ids[name] = ref.id

    def delete_sandbox(self, name: str) -> None:
        """Delete a sandbox; a missing sandbox is treated as already gone."""
        import grpc
        from openshell import SandboxError

        try:
            self._client.delete(name)
        except grpc.RpcError as exc:
            if isinstance(exc, grpc.Call) and exc.code() == grpc.StatusCode.NOT_FOUND:
                self._ids.pop(name, None)
                return
            raise click.ClickException(
                f"Failed to delete OpenShell sandbox '{name}': {exc}"
            ) from exc
        except SandboxError as exc:
            if "not found" in str(exc).lower():
                self._ids.pop(name, None)
                return
            raise click.ClickException(
                f"Failed to delete OpenShell sandbox '{name}': {exc}"
            ) from exc
        self._ids.pop(name, None)

    def _id_for(self, name: str) -> str:
        cached = self._ids.get(name)
        if cached is None:
            ref = self._guard(
                f"Could not resolve OpenShell sandbox '{name}'",
                lambda: self._client.get(name),
            )
            cached = ref.id
            self._ids[name] = cached
        return cached

    def _guard(self, message: str, call: Callable[[], _T]) -> _T:
        """Run an SDK *call*, surfacing gRPC / SandboxError as ClickException."""
        import grpc
        from openshell import SandboxError

        try:
            return call()
        except (grpc.RpcError, SandboxError) as exc:
            raise click.ClickException(f"{message}: {exc}") from exc


class OpenShellSandboxLauncher(SandboxLauncher):
    """
    :class:`SandboxLauncher` for NVIDIA OpenShell, over the ``openshell`` SDK.

    Lifecycle and transport map onto the SDK's gRPC client: create /
    wait-ready / delete for the sandbox, command execution for running
    code, and an exec-over-stdin stream for shipping wheels.
    """

    provider: ClassVar[str] = "openshell"
    supports_local_port_forward: ClassVar[bool] = False

    def __init__(
        self,
        *,
        image: str | None = None,
        env: Sequence[str] | None = None,
        cluster: str | None = None,
    ) -> None:
        """
        :param image: Registry image to provision from
            (``sandbox.openshell.image``); ``None`` resolves
            :data:`HOST_IMAGE_ENV_VAR` then the official host image.
        :param env: Server-process env var NAMES to inject into every
            sandbox (``sandbox.openshell.env``); ``None`` resolves
            :data:`SANDBOX_ENV_PASSTHROUGH_ENV_VAR`.
        :param cluster: Gateway name to connect to
            (``sandbox.openshell.cluster``); ``None`` lets the SDK
            resolve the active gateway (``$OPENSHELL_GATEWAY`` or
            ``~/.config/openshell/active_gateway``).
        """
        self._image_ref = image
        self._env_names = tuple(env) if env is not None else None
        self._cluster = cluster
        self._client: _OpenShellClient | None = None

    def prepare(self) -> None:
        """Preflight: the SDK must be installed and a gateway resolvable."""
        _ensure_sdk()
        # Constructing the client resolves the active gateway from on-disk
        # state, failing fast with a remediation hint before remote work.
        self._openshell()

    def provision(self, name: str) -> str:
        """Create a sandbox from the host image and wait until it is ready."""
        image = self._image_ref or os.environ.get(HOST_IMAGE_ENV_VAR) or DEFAULT_HOST_IMAGE
        env_vars = self._resolve_sandbox_env()
        click.echo(f"▸ Creating OpenShell sandbox from {image}")
        # OpenShell assigns its own petname; the requested `name` is advisory.
        sandbox_name = self._openshell().create_sandbox(image=image, env=env_vars)
        click.echo(f"  → created {sandbox_name}")
        return sandbox_name

    def attach(self, sandbox_id: str) -> None:
        """Validate access to an existing OpenShell sandbox by name."""
        click.echo(f"▸ Reusing existing OpenShell sandbox '{sandbox_id}'")
        self._openshell().get_status(sandbox_id)

    def keep_alive(self, sandbox_id: str) -> None:
        """No idle auto-stop management is exposed by the OpenShell API."""
        click.echo(f"  → OpenShell sandbox '{sandbox_id}' remains active until destroyed")

    def run(self, sandbox_id: str, command: str, *, check: bool = True) -> RemoteCommandResult:
        """Run ``bash -lc <command>`` in the sandbox and capture its output."""
        result = self._openshell().execute(sandbox_id, ["bash", "-lc", command])
        if result.stdout:
            click.echo(result.stdout, nl=False)
        if result.stderr:
            click.echo(result.stderr, nl=False, err=True)
        if check and result.exit_code != 0:
            raise click.ClickException(
                f"Remote command failed on OpenShell sandbox '{sandbox_id}' "
                f"(exit {result.exit_code}): {command}"
            )
        return RemoteCommandResult(
            returncode=result.exit_code, stdout=result.stdout, stderr=result.stderr
        )

    def run_background(
        self, sandbox_id: str, command: str, *, log_path: str = "/tmp/omnigent-host.log"
    ) -> RemoteCommandResult:
        """Hold *command* on a long-lived exec stream instead of detaching.

        OpenShell kills an exec's processes when the RPC returns, so the
        base class's ``setsid nohup`` detach pattern doesn't work. Instead
        the command runs in the foreground of an ``exec_stream`` drained on
        a daemon thread — the stream stays open for the process's lifetime.
        """
        bg_command = f"{command} > {log_path} 2>&1 < /dev/null"
        self._openshell().exec_background(
            sandbox_id, ["bash", "-lc", bg_command], timeout=_FOREGROUND_TIMEOUT_S
        )
        return RemoteCommandResult(returncode=0, stdout="launched\n", stderr="")

    def put(self, sandbox_id: str, local_path: Path, remote_path: str) -> None:
        """Copy a local file into the sandbox by piping its bytes to ``cat``."""
        content = local_path.read_bytes()
        parent = shlex.quote(str(PurePosixPath(remote_path).parent))
        dest = shlex.quote(remote_path)
        result = self._openshell().execute(
            sandbox_id,
            ["bash", "-c", f"mkdir -p {parent} && cat > {dest}"],
            stdin=content,
        )
        if result.exit_code != 0:
            raise click.ClickException(
                f"File upload to OpenShell sandbox '{sandbox_id}' failed "
                f"(exit {result.exit_code}): {result.stderr.strip()}"
            )

    def exec_foreground(self, sandbox_id: str, command: str) -> int:
        """
        Run *command* in the sandbox, streaming its output, until it exits.

        Holds ``omnigent host`` open for ``omnigent sandbox connect``;
        Ctrl-C kills the remote process and re-raises ``KeyboardInterrupt``.
        """
        client = self._openshell()
        pidfile = _FOREGROUND_PIDFILE_TEMPLATE.format(sandbox_id=sandbox_id)
        # `exec` keeps the recorded pid across the shell swap, so a Ctrl-C
        # can kill the remote process — cancelling the local gRPC stream
        # stops our reads but doesn't stop the remote command.
        remote = f"echo $$ > {shlex.quote(pidfile)} && exec {command}"
        try:
            return client.run_foreground(
                sandbox_id, ["bash", "-lc", remote], timeout=_FOREGROUND_TIMEOUT_S
            )
        except KeyboardInterrupt:
            click.echo("\n  → detaching; stopping the remote process")
            client.execute(
                sandbox_id,
                [
                    "bash",
                    "-lc",
                    f"kill $(cat {shlex.quote(pidfile)}) 2>/dev/null || true",
                ],
            )
            raise

    def wheel_install_command(self, remote_tgz_path: str) -> str:
        """Overlay shipped wheels onto the prebaked host image."""
        return host_image_wheel_install_command(remote_tgz_path)

    def terminate(self, sandbox_id: str) -> None:
        """Delete a sandbox, releasing its compute."""
        try:
            self._openshell().delete_sandbox(sandbox_id)
        finally:
            if self._client is not None:
                self._client.close()
                self._client = None

    def close(self) -> None:
        """Release the underlying gRPC channel, if one was opened."""
        if self._client is not None:
            self._client.close()
            self._client = None

    def __del__(self) -> None:
        self.close()

    def _openshell(self) -> _OpenShellClient:
        if self._client is None:
            self._client = _OpenShellClient(cluster=self._cluster)
        return self._client

    def _resolve_sandbox_env(self) -> dict[str, str]:
        if self._env_names is not None:
            names: Sequence[str] = self._env_names
        else:
            names = [
                name.strip()
                for name in os.environ.get(SANDBOX_ENV_PASSTHROUGH_ENV_VAR, "").split(",")
                if name.strip()
            ]
        resolved: dict[str, str] = {}
        for name in names:
            value = os.environ.get(name)
            if value is None:
                raise click.ClickException(
                    f"sandbox env passthrough lists '{name}' but it is not set "
                    "in the server's environment — set it (or remove it from "
                    f"sandbox.openshell.env / {SANDBOX_ENV_PASSTHROUGH_ENV_VAR})."
                )
            resolved[name] = value
        return resolved