"""Runtime — standard and passthrough SecuredSandbox implementations.

See spec Section 9.1.1 for the full design.
"""

from __future__ import annotations

import asyncio
import base64
import json
import shlex

from .errors import TransportError
from .types import (
    ExecOpts,
    ExecResult,
    ReadFileFailure,
    ReadFileResult,
    ReadFileSuccess,
    SandboxAdapter,
    SecuredSandbox,
    SecurityMode,
    WriteFileFailure,
    WriteFileResult,
    WriteFileSuccess,
)


def create_secured_sandbox(
    adapter: SandboxAdapter,
    session_id: str,
    security_mode: SecurityMode,
    *,
    passthrough: bool = False,
) -> SecuredSandbox:
    """Create the appropriate SecuredSandbox for the runtime mode."""
    if passthrough:
        return _PassthroughSecuredSandbox(adapter, session_id, security_mode)
    return _StandardSecuredSandbox(adapter, session_id, security_mode)


class _StandardSecuredSandbox:
    """Wraps every operation in `agentsh exec --output json $SID --`."""

    def __init__(
        self,
        adapter: SandboxAdapter,
        sid: str,
        mode: SecurityMode,
    ) -> None:
        self._adapter = adapter
        self._session_id = sid
        self._security_mode = mode

    @property
    def session_id(self) -> str:
        return self._session_id

    @property
    def security_mode(self) -> SecurityMode:
        return self._security_mode

    async def exec(
        self,
        command: str,
        *,
        cwd: str | None = None,
        timeout: float | None = None,
    ) -> ExecResult:
        args = ["exec", "--output", "json", self._session_id, "--", "bash", "-c"]
        full_cmd = f"cd {shlex.quote(cwd)} && {command}" if cwd else command
        args.append(full_cmd)

        if timeout is not None:
            raw = await asyncio.wait_for(
                self._adapter.exec("agentsh", args),
                timeout=timeout,
            )
        else:
            raw = await self._adapter.exec("agentsh", args)

        # Transport failure detection
        if raw.exit_code == 127 and "agentsh" in raw.stderr:
            raise TransportError(
                session_id=self._session_id,
                command=command,
                stderr=raw.stderr,
            )

        # Parse JSON envelope
        return _parse_exec_envelope(raw)

    async def write_file(
        self,
        path: str,
        content: str | bytes,
    ) -> WriteFileResult:
        data = content if isinstance(content, bytes) else content.encode()
        b64 = base64.b64encode(data).decode("ascii")
        safe_path = shlex.quote(path)
        result = await self.exec(
            f'printf "%s" {shlex.quote(b64)} | base64 -d > {safe_path}',
        )
        if result.exit_code != 0:
            return WriteFileFailure(
                success=False,
                path=path,
                error=result.stderr,
            )
        return WriteFileSuccess(success=True, path=path)

    async def read_file(self, path: str) -> ReadFileResult:
        result = await self.exec(f"cat {shlex.quote(path)}")
        if result.exit_code != 0:
            return ReadFileFailure(
                success=False,
                path=path,
                error=result.stderr,
            )
        return ReadFileSuccess(success=True, path=path, content=result.stdout)

    async def stop(self) -> None:
        # Stop agentsh session, then stop provider sandbox
        try:
            await self._adapter.exec(
                "agentsh",
                ["session", "stop", self._session_id],
            )
        except Exception:
            pass
        await self._adapter.stop()


class _PassthroughSecuredSandbox:
    """
    Used with install_strategy="running". Commands go directly through
    the adapter; the shell shim enforces policy on every command.
    """

    def __init__(
        self,
        adapter: SandboxAdapter,
        sid: str,
        mode: SecurityMode,
    ) -> None:
        self._adapter = adapter
        self._session_id = sid
        self._security_mode = mode

    @property
    def session_id(self) -> str:
        return self._session_id

    @property
    def security_mode(self) -> SecurityMode:
        return self._security_mode

    async def exec(
        self,
        command: str,
        *,
        cwd: str | None = None,
        timeout: float | None = None,
    ) -> ExecResult:
        opts = ExecOpts(cwd=cwd)
        if timeout is not None:
            return await asyncio.wait_for(
                self._adapter.exec("bash", ["-c", command], opts),
                timeout=timeout,
            )
        return await self._adapter.exec("bash", ["-c", command], opts)

    async def write_file(
        self,
        path: str,
        content: str | bytes,
    ) -> WriteFileResult:
        # Route through adapter.exec (NOT adapter.write_file) so the
        # shell shim can enforce policy.
        data = content if isinstance(content, bytes) else content.encode()
        b64 = base64.b64encode(data).decode("ascii")
        safe_path = shlex.quote(path)
        result = await self.exec(
            f'printf "%s" {shlex.quote(b64)} | base64 -d > {safe_path}',
        )
        if result.exit_code != 0:
            return WriteFileFailure(
                success=False,
                path=path,
                error=result.stderr,
            )
        return WriteFileSuccess(success=True, path=path)

    async def read_file(self, path: str) -> ReadFileResult:
        result = await self.exec(f"cat {shlex.quote(path)}")
        if result.exit_code != 0:
            return ReadFileFailure(
                success=False,
                path=path,
                error=result.stderr,
            )
        return ReadFileSuccess(success=True, path=path, content=result.stdout)

    async def stop(self) -> None:
        await self._adapter.stop()


def _parse_exec_envelope(raw: ExecResult) -> ExecResult:
    """
    Parse the agentsh exec JSON envelope. Falls back to raw result
    if parsing fails (e.g., mock adapters or non-JSON output).
    """
    try:
        envelope = json.loads(raw.stdout)
        result = envelope["result"]
        return ExecResult(
            stdout=result.get("stdout", ""),
            stderr=result.get("stderr", ""),
            exit_code=result.get("exit_code", 0),
        )
    except (json.JSONDecodeError, KeyError, TypeError):
        return raw
