"""Provider-neutral contract for external agent-worker sandbox tools."""

from __future__ import annotations

import asyncio
import base64
import json
import os
import shlex
import time
import uuid
from contextlib import contextmanager
from dataclasses import asdict, dataclass
from pathlib import PurePosixPath
from typing import Any, Callable, Iterator, Protocol


WORKSPACE = "/workspace"
STORAGE_CREDENTIAL_ENV = (
    "AIMEM_ACCESS_KEY",
    "AIMEM_SECRET_KEY",
    "AIMEM_SESSION_TOKEN",
    "AISTOR_ROOT_USER",
    "AISTOR_ROOT_PASSWORD",
)


@dataclass(frozen=True)
class ToolEvent:
    call_id: str
    tool: str
    exit_code: int
    stdout: str
    stderr: str
    duration_ms: int


class WorkspaceTools(Protocol):
    """High-level tools supplied by a provider-specific sandbox adapter."""

    sandbox_id: str
    trace: list[ToolEvent]

    def read_text(self, path: str) -> str: ...

    def write_text(self, path: str, content: str) -> None: ...

    def sandbox_env_var_is_set(self, name: str) -> bool: ...

    def create_symlink(self, path: str, target: str) -> None: ...

    def path_exists(self, path: str) -> bool: ...

    def delayed_write(self, path: str, content: str, delay: int, timeout: int) -> None: ...

    def remove_tree(self, path: str) -> None: ...


class AsyncWorkspaceTools(Protocol):
    sandbox_id: str
    trace: list[ToolEvent]

    async def read_text(self, path: str) -> str: ...

    async def write_text(self, path: str, content: str) -> None: ...

    async def sandbox_env_var_is_set(self, name: str) -> bool: ...

    async def create_symlink(self, path: str, target: str) -> None: ...

    async def path_exists(self, path: str) -> bool: ...

    async def delayed_write(self, path: str, content: str, delay: int, timeout: int) -> None: ...

    async def remove_tree(self, path: str) -> None: ...


CommandRunner = Callable[[str, int], tuple[int, str, str]]


class RemoteCommandWorkspaceTools:
    """Shared typed tools for providers whose SDK exposes remote commands."""

    def __init__(
        self,
        sandbox_id: str,
        runner: CommandRunner,
        *,
        workspace: str = WORKSPACE,
        timeout: int = 30,
    ) -> None:
        self.sandbox_id = sandbox_id
        self.runner = runner
        self.workspace = workspace
        self.timeout = timeout
        self.trace: list[ToolEvent] = []

    def _exec(self, tool: str, argv: list[str], *, timeout: int | None = None) -> str:
        call_id = f"tool_{uuid.uuid4().hex}"
        command_timeout = timeout if timeout is not None else self.timeout
        command = shlex.join(["timeout", "--signal=KILL", f"{command_timeout}s", *argv])
        started = time.monotonic()
        try:
            exit_code, stdout, stderr = self.runner(command, command_timeout + 5)
        except Exception as exc:
            stderr = f"runner raised {type(exc).__name__}"
            self.trace.append(
                ToolEvent(
                    call_id=call_id,
                    tool=tool,
                    exit_code=-1,
                    stdout="",
                    stderr=stderr,
                    duration_ms=round((time.monotonic() - started) * 1000),
                )
            )
            raise RuntimeError(redacted_failure(tool, self.sandbox_id, -1, stderr)) from exc
        event = ToolEvent(
            call_id=call_id,
            tool=tool,
            exit_code=exit_code,
            stdout=stdout,
            stderr=stderr,
            duration_ms=round((time.monotonic() - started) * 1000),
        )
        self.trace.append(event)
        if exit_code != 0:
            raise RuntimeError(redacted_failure(tool, self.sandbox_id, exit_code, stderr))
        return stdout

    def read_text(self, path: str) -> str:
        return self._exec("read_text", read_text_argv(self.workspace, path))

    def write_text(self, path: str, content: str) -> None:
        self._exec("write_text", write_text_argv(self.workspace, path, content))

    def sandbox_env_var_is_set(self, name: str) -> bool:
        return self._exec("sandbox_env_probe", env_probe_argv(name)).strip() == "set"

    def create_symlink(self, path: str, target: str) -> None:
        self._exec("create_symlink", create_symlink_argv(self.workspace, path, target))

    def path_exists(self, path: str) -> bool:
        return self._exec("path_exists", path_exists_argv(self.workspace, path)).strip() == "present"

    def delayed_write(self, path: str, content: str, delay: int, timeout: int) -> None:
        self._exec(
            "timeout_probe",
            delayed_write_argv(self.workspace, path, content, delay),
            timeout=timeout,
        )

    def remove_tree(self, path: str) -> None:
        self._exec("security_probe_cleanup", remove_tree_argv(self.workspace, path))


def relative_workspace_path(path: str) -> str:
    parsed = PurePosixPath(path)
    if parsed.is_absolute() or not parsed.parts or ".." in parsed.parts:
        raise ValueError(f"workspace path must be relative and cannot escape: {path!r}")
    return parsed.as_posix()


def read_text_argv(workspace: str, path: str) -> list[str]:
    relative = relative_workspace_path(path)
    script = """\
root=$(realpath "$1") || exit 96
path=$(realpath -m "$root/$2") || exit 96
case "$path" in "$root"/*) ;; *) echo 'path escapes workspace' >&2; exit 97;; esac
cat "$path"
"""
    return ["sh", "-c", script, "sh", workspace, relative]


def read_text_base64_argv(workspace: str, path: str) -> list[str]:
    relative = relative_workspace_path(path)
    script = """\
root=$(realpath "$1") || exit 96
path=$(realpath -m "$root/$2") || exit 96
case "$path" in "$root"/*) ;; *) echo 'path escapes workspace' >&2; exit 97;; esac
base64 -w0 "$path"
"""
    return ["sh", "-c", script, "sh", workspace, relative]


def write_text_argv(workspace: str, path: str, content: str) -> list[str]:
    relative = relative_workspace_path(path)
    encoded = base64.b64encode(content.encode()).decode()
    script = """\
set -e
root=$(realpath "$1") || exit 96
path=$(realpath -m "$root/$2") || exit 96
case "$path" in "$root"/*) ;; *) echo 'path escapes workspace' >&2; exit 97;; esac
mkdir -p "$(dirname "$path")"
printf '%s' "$3" | base64 -d > "$path"
sync "$path" 2>/dev/null || sync
"""
    return ["sh", "-c", script, "sh", workspace, relative, encoded]


def env_probe_argv(name: str) -> list[str]:
    if not name.replace("_", "").isalnum():
        raise ValueError(f"invalid environment variable name: {name!r}")
    script = """\
if printenv "$1" >/dev/null 2>&1; then
    echo set
else
    echo unset
fi
"""
    return ["sh", "-c", script, "sh", name]


def create_symlink_argv(workspace: str, path: str, target: str) -> list[str]:
    relative = relative_workspace_path(path)
    script = """\
root=$(realpath "$1") || exit 96
link=$(realpath -m "$root/$2") || exit 96
case "$link" in "$root"/*) ;; *) echo 'path escapes workspace' >&2; exit 97;; esac
mkdir -p "$(dirname "$link")"
ln -s "$3" "$link"
"""
    return ["sh", "-c", script, "sh", workspace, relative, target]


def path_exists_argv(workspace: str, path: str) -> list[str]:
    relative = relative_workspace_path(path)
    script = """\
root=$(realpath "$1") || exit 96
path=$(realpath -m "$root/$2") || exit 96
case "$path" in "$root"/*) ;; *) echo 'path escapes workspace' >&2; exit 97;; esac
if test -e "$path"; then echo present; else echo absent; fi
"""
    return ["sh", "-c", script, "sh", workspace, relative]


def delayed_write_argv(workspace: str, path: str, content: str, delay: int) -> list[str]:
    relative = relative_workspace_path(path)
    encoded = base64.b64encode(content.encode()).decode()
    script = """\
root=$(realpath "$1") || exit 96
path=$(realpath -m "$root/$2") || exit 96
case "$path" in "$root"/*) ;; *) echo 'path escapes workspace' >&2; exit 97;; esac
mkdir -p "$(dirname "$path")"
sleep "$3"
printf '%s' "$4" | base64 -d > "$path"
"""
    return ["sh", "-c", script, "sh", workspace, relative, str(delay), encoded]


def remove_tree_argv(workspace: str, path: str) -> list[str]:
    relative = relative_workspace_path(path)
    script = """\
root=$(realpath "$1") || exit 96
path=$(realpath -m "$root/$2") || exit 96
case "$path" in "$root"/*) ;; *) echo 'path escapes workspace' >&2; exit 97;; esac
rm -rf "$path"
"""
    return ["sh", "-c", script, "sh", workspace, relative]


def redacted_failure(tool: str, sandbox_id: str, exit_code: int, stderr: str) -> str:
    return (
        f"{tool} failed in sandbox {sandbox_id} (exit {exit_code}); " f"stderr redacted ({len(stderr.encode())} bytes)"
    )


def load_worker_skill() -> str:
    """Best-effort read of the out-of-sandbox worker's host-side aimem skill.

    A real deployment installs it with ``npx skills add minio/skills/aimem``
    (github.com/minio/skills), or points ``AIMEM_SKILL_DIR`` at it. Returns the
    ``SKILL.md`` body if present, else ``""`` — the worker still functions via
    its typed tools; the skill only adds workspace guidance, so it is optional
    for this smoke (which validates the tool boundary, not skill delivery)."""

    skill_dir = os.environ.get("AIMEM_SKILL_DIR") or os.path.join(
        os.path.expanduser("~"), ".agents", "skills", "aimem"
    )
    try:
        with open(os.path.join(skill_dir, "SKILL.md"), encoding="utf-8") as handle:
            return handle.read()
    except OSError:
        return ""


@contextmanager
def without_storage_credentials() -> Iterator[None]:
    """Temporarily remove storage credentials from the worker environment."""

    saved = {name: os.environ.pop(name) for name in STORAGE_CREDENTIAL_ENV if name in os.environ}
    try:
        yield
    finally:
        os.environ.update(saved)


def _assert_local_boundary() -> None:
    if os.path.isdir(f"{WORKSPACE}/.aimem/annot"):
        raise AssertionError("worker unexpectedly sees the sandbox AIMem mount")
    for name in STORAGE_CREDENTIAL_ENV:
        if name in os.environ:
            raise AssertionError(f"{name} leaked into the external agent worker")


def assert_boundary(tools: WorkspaceTools) -> None:
    _assert_local_boundary()
    if tools.sandbox_env_var_is_set("OPENAI_API_KEY"):
        raise AssertionError("OPENAI_API_KEY leaked into the execution sandbox")


async def assert_boundary_async(tools: AsyncWorkspaceTools) -> None:
    _assert_local_boundary()
    if await tools.sandbox_env_var_is_set("OPENAI_API_KEY"):
        raise AssertionError("OPENAI_API_KEY leaked into the execution sandbox")


def run_security_contract(tools: WorkspaceTools) -> None:
    """Verify workspace confinement and remote timeout termination."""

    for path in ("/etc/passwd", "../etc/passwd"):
        try:
            tools.read_text(path)
        except ValueError:
            pass
        else:
            raise AssertionError(f"workspace path escape was accepted: {path!r}")

    probe_dir = f"external-worker/security-{uuid.uuid4().hex}"
    symlink = f"{probe_dir}/escape"
    marker = f"{probe_dir}/late-write.txt"
    try:
        tools.create_symlink(symlink, "/etc")
        try:
            tools.read_text(f"{symlink}/passwd")
        except RuntimeError:
            pass
        else:
            raise AssertionError("workspace symlink escape was accepted")

        try:
            tools.delayed_write(marker, "late", delay=3, timeout=1)
        except RuntimeError:
            pass
        else:
            raise AssertionError("sandbox timeout probe unexpectedly completed")
        time.sleep(3)
        if tools.path_exists(marker):
            raise AssertionError("timed-out sandbox process continued and wrote its marker")
    finally:
        tools.remove_tree(probe_dir)


async def run_security_contract_async(tools: AsyncWorkspaceTools) -> None:
    """Async equivalent of the shared negative security probes."""

    for path in ("/etc/passwd", "../etc/passwd"):
        try:
            await tools.read_text(path)
        except ValueError:
            pass
        else:
            raise AssertionError(f"workspace path escape was accepted: {path!r}")

    probe_dir = f"external-worker/security-{uuid.uuid4().hex}"
    symlink = f"{probe_dir}/escape"
    marker = f"{probe_dir}/late-write.txt"
    try:
        await tools.create_symlink(symlink, "/etc")
        try:
            await tools.read_text(f"{symlink}/passwd")
        except RuntimeError:
            pass
        else:
            raise AssertionError("workspace symlink escape was accepted")

        try:
            await tools.delayed_write(marker, "late", delay=3, timeout=1)
        except RuntimeError:
            pass
        else:
            raise AssertionError("sandbox timeout probe unexpectedly completed")
        await asyncio.sleep(3)
        if await tools.path_exists(marker):
            raise AssertionError("timed-out sandbox process continued and wrote its marker")
    finally:
        await tools.remove_tree(probe_dir)


def run_contract(
    tools: WorkspaceTools,
    probe_path: str,
    probe_content: str,
    *,
    verify_only: bool = False,
) -> None:
    """Run the deterministic worker-to-sandbox workspace contract."""

    assert_boundary(tools)
    if not verify_only:
        tools.write_text(probe_path, probe_content)
    actual = tools.read_text(probe_path)
    if actual != probe_content:
        raise AssertionError(f"workspace continuity mismatch for {probe_path!r}")
    if not verify_only:
        run_security_contract(tools)


async def run_async_contract(
    tools: AsyncWorkspaceTools,
    probe_path: str,
    probe_content: str,
    *,
    verify_only: bool = False,
) -> None:
    """Async equivalent for providers whose command API is awaitable."""

    await assert_boundary_async(tools)

    if not verify_only:
        await tools.write_text(probe_path, probe_content)
    if await tools.read_text(probe_path) != probe_content:
        raise AssertionError(f"workspace continuity mismatch for {probe_path!r}")
    if not verify_only:
        await run_security_contract_async(tools)


def model_task(probe_path: str, probe_content: str) -> str:
    """The instruction a model worker gets. The content sits on its own line
    between markers rather than in quotes: a quoted value invites the model to
    write the quotes too, and the file then misses by exactly those bytes."""

    return (
        f"Write the text between the BEGIN and END lines to the workspace file {probe_path}. "
        "Write only those characters: no quotes, no surrounding whitespace, no newline. "
        "Then read the file back and report completion.\n"
        f"BEGIN\n{probe_content}\nEND"
    )


def run_model(tools: WorkspaceTools, probe_path: str, probe_content: str, model: str) -> None:
    """Run an optional real OpenAI Agents SDK worker over the same tools."""

    try:
        from agents import Agent, ModelSettings, Runner, function_tool
    except ImportError as exc:
        raise RuntimeError("model smoke requires the openai-agents package in the worker") from exc

    skill = load_worker_skill()

    @function_tool
    def write_workspace_file(path: str, content: str) -> str:
        """Write UTF-8 text to a relative path in the remote sandbox workspace."""

        tools.write_text(path, content)
        return f"wrote {path}"

    @function_tool
    def read_workspace_file(path: str) -> str:
        """Read UTF-8 text from a relative path in the remote sandbox workspace."""

        return tools.read_text(path)

    agent = Agent(
        name="aimem-external-worker-smoke",
        model=model,
        instructions=(
            "You are an agent worker outside the execution sandbox. Use only the supplied typed "
            "tools to act on the remote sandbox workspace. Write the exact requested content, then "
            "read the file back. Do not claim success until the read matches."
            + (f"\n\nYour aimem skill, loaded host-side:\n\n{skill}" if skill else "")
        ),
        tools=[write_workspace_file, read_workspace_file],
        # Name the tier: left unset, OpenAI picks one per project, and a
        # project that may not use it gets every request rejected.
        model_settings=ModelSettings(extra_args={"service_tier": "default"}),
    )
    result = Runner.run_sync(agent, input=model_task(probe_path, probe_content))
    print(result.final_output)
    actual = tools.read_text(probe_path)
    if actual != probe_content:
        raise AssertionError(
            "external agent did not persist the requested workspace write: "
            f"expected {len(probe_content.encode())} bytes, read {len(actual.encode())}"
        )


def emit_trace(tools: WorkspaceTools) -> None:
    """Emit correlated metadata without copying tool output or secrets to logs."""

    for item in tools.trace:
        payload: dict[str, Any] = asdict(item)
        payload["stdout_bytes"] = len(payload.pop("stdout").encode())
        payload["stderr_bytes"] = len(payload.pop("stderr").encode())
        payload["sandbox_id"] = tools.sandbox_id
        print(json.dumps(payload, sort_keys=True))