"""NSJail configuration generation for agent runtime sandbox.

This module generates protobuf-format nsjail configurations specifically
for running the agent runtime in an isolated sandbox.

Security model:
- Network namespace always isolated for private loopback; pasta enables outbound access
- LLM access via internal bridge (localhost:4100) proxied through Unix socket to host LLM gateway
- Namespace isolation (PID, user, mount, IPC, UTS namespaces)
- Fresh read-only /proc inside the jail PID namespace
- All tool execution via MCP socket to trusted server outside sandbox
- Uses same base rootfs as action sandbox (Python 3.12)
- Site-packages or a minimal Claude SDK subtree mounted read-only

Key design:
- Runtime executes a standalone shim script from the per-job /work directory
- Shim mode avoids Tracecat package mounts entirely
- Control socket at /var/run/tracecat/control.sock
- LLM socket at /var/run/tracecat/llm.sock (proxied to LLM gateway)
"""

from __future__ import annotations

import os
import re
from dataclasses import dataclass, field
from pathlib import Path

from tracecat.agent.common.config import (
    CONTROL_SOCKET_NAME,
    JAILED_CONTROL_SOCKET_PATH,
    JAILED_LLM_SOCKET_PATH,
    TRACECAT__AGENT_SANDBOX_MEMORY_MB,
    TRACECAT__AGENT_SANDBOX_TIMEOUT,
    TRUSTED_MCP_SOCKET_PATH,
)
from tracecat.agent.common.exceptions import AgentSandboxValidationError

# Valid environment variable name pattern (POSIX compliant)
_ENV_VAR_KEY_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")


def _contains_dangerous_chars(value: str) -> tuple[bool, str | None]:
    """Check for null bytes and control characters that could affect sandbox behavior.

    Null bytes can truncate strings in C-based tools like nsjail. Control characters
    could affect sandbox behavior in unexpected ways.

    Args:
        value: String value to check.

    Returns:
        Tuple of (is_dangerous, reason). If is_dangerous is False, reason is None.
    """
    if "\x00" in value:
        return True, "contains null byte"

    for char in value:
        code = ord(char)
        # Control chars 0x01-0x1F (except tab 0x09, newline 0x0A, CR 0x0D) and DEL 0x7F
        if code < 0x20 and code not in (0x09, 0x0A, 0x0D):
            return True, f"contains control character 0x{code:02x}"
        if code == 0x7F:
            return True, "contains DEL control character"

    return False, None


@dataclass(frozen=True)
class AgentResourceLimits:
    """Resource limits for agent sandbox execution.

    Defaults are read from environment variables:
    - TRACECAT__AGENT_SANDBOX_MEMORY_MB: memory_mb (default 4096 = 4 GiB)
    - TRACECAT__AGENT_SANDBOX_TIMEOUT: timeout_seconds and cpu_seconds (default 1800s)

    Attributes:
        memory_mb: Maximum memory in megabytes.
        cpu_seconds: Maximum CPU time in seconds.
        max_file_size_mb: Maximum file size in megabytes.
        max_open_files: Maximum number of open file descriptors.
        max_processes: Maximum number of processes/threads.
        timeout_seconds: Maximum wall-clock execution time.
    """

    memory_mb: int = field(default_factory=lambda: TRACECAT__AGENT_SANDBOX_MEMORY_MB)
    cpu_seconds: int = field(default_factory=lambda: TRACECAT__AGENT_SANDBOX_TIMEOUT)
    max_file_size_mb: int = 256
    max_open_files: int = 512
    max_processes: int = 128
    timeout_seconds: int = field(
        default_factory=lambda: TRACECAT__AGENT_SANDBOX_TIMEOUT
    )


@dataclass
class AgentSandboxConfig:
    """Configuration for agent sandbox execution.

    Attributes:
        resources: Resource limits for the sandbox.
        env_vars: Environment variables to inject into the sandbox.
    """

    resources: AgentResourceLimits = field(default_factory=AgentResourceLimits)
    env_vars: dict[str, str] = field(default_factory=dict)


# In-jail path for the standalone Claude shim copied into the per-job workdir.
JAILED_SHIM_ENTRYPOINT_PATH = "/work/shim_entrypoint.py"

# Minimal base environment for sandboxed agent processes
AGENT_SANDBOX_BASE_ENV = {
    "PATH": "/usr/local/bin:/usr/bin:/bin",
    "HOME": "/home/agent",
    "USER": "agent",
    "TRACECAT__DISABLE_NSJAIL": "false",
    "PYTHONDONTWRITEBYTECODE": "1",
    "PYTHONUNBUFFERED": "1",
    "LANG": "C.UTF-8",
    "LC_ALL": "C.UTF-8",
    # Required for Python to find libpython3.12.so in nsjail sandbox
    "LD_LIBRARY_PATH": "/usr/local/lib:/usr/lib:/lib",
    # PYTHONPATH: /site-packages for the mounted Claude SDK package tree.
    "PYTHONPATH": "/site-packages",
}


def _validate_env_key(key: str) -> None:
    """Validate environment variable key is safe for protobuf config.

    Args:
        key: Environment variable name to validate.

    Raises:
        AgentSandboxValidationError: If key contains invalid characters.
    """
    if not _ENV_VAR_KEY_PATTERN.match(key):
        raise AgentSandboxValidationError(
            f"Invalid environment variable key: {key!r}. "
            "Keys must match pattern [A-Za-z_][A-Za-z0-9_]*"
        )


def _validate_env_value(key: str, value: str) -> None:
    """Validate environment variable value is safe for sandbox execution.

    Args:
        key: Environment variable name (for error messages).
        value: Environment variable value to validate.

    Raises:
        AgentSandboxValidationError: If value contains dangerous characters.
    """
    is_dangerous, reason = _contains_dangerous_chars(value)
    if is_dangerous:
        raise AgentSandboxValidationError(
            f"Invalid environment variable value for {key!r}: {reason}"
        )


def _validate_path(path: Path, name: str) -> None:
    """Validate path is safe for protobuf config interpolation.

    Args:
        path: Path to validate.
        name: Human-readable name for error messages.

    Raises:
        AgentSandboxValidationError: If path contains dangerous characters or traversal.
    """
    path_str = str(path)

    # Check for null bytes and control characters
    is_dangerous, reason = _contains_dangerous_chars(path_str)
    if is_dangerous:
        raise AgentSandboxValidationError(f"Invalid {name} path: {reason}")

    # Check for path traversal attempts
    # Resolve to absolute and check it doesn't escape expected boundaries
    try:
        path.resolve()
    except (OSError, ValueError) as e:
        raise AgentSandboxValidationError(
            f"Invalid {name} path: failed to resolve - {e}"
        ) from e

    # Characters that could break protobuf text format parsing
    dangerous_chars = {'"', "'", "\n", "\r", "\\", "{", "}"}
    found_chars = [c for c in dangerous_chars if c in path_str]
    if found_chars:
        raise AgentSandboxValidationError(
            f"Invalid {name} path: contains dangerous characters {found_chars!r}"
        )

    # Reject raw ".." components even before resolution (defense in depth)
    path_parts = path_str.split("/")
    if ".." in path_parts:
        raise AgentSandboxValidationError(
            f"Invalid {name} path: contains '..' path traversal component"
        )


def build_agent_nsjail_config(
    rootfs: Path,
    job_dir: Path,
    socket_dir: Path,
    config: AgentSandboxConfig,
    site_packages_dir: Path,
    llm_socket_path: Path | None,
    mcp_socket_path: Path | None = None,
    *,
    mount_control_socket: bool = True,
    control_socket_path: Path | None = None,
    session_home_dir: Path | None = None,
    session_project_dir: Path | None = None,
    enable_internet_access: bool = False,
    skills_dir: Path | None = None,
) -> str:
    """Build nsjail protobuf config for agent runtime execution.

    Args:
        rootfs: Path to the sandbox rootfs (same as action sandbox).
        job_dir: Directory for job-specific data.
        socket_dir: Directory containing the per-job control socket
            (control.sock) created by the orchestrator.
        config: Agent sandbox configuration.
        site_packages_dir: Path to site-packages with Claude SDK deps.
        llm_socket_path: Optional path to the LLM socket for proxied LLM gateway
            access.
        mcp_socket_path: Optional path to the trusted MCP socket for tool calls.
        mount_control_socket: Whether to mount the per-job control socket into the
            jail. The current shim path does not require this.
        control_socket_path: Optional explicit control socket path. When omitted
            and mount_control_socket is True, defaults to socket_dir/control.sock.
        session_home_dir: Optional host directory mounted as the jailed Claude home.
        session_project_dir: Optional host directory mounted as the jailed project.
        enable_internet_access: If True, enables pasta userspace networking for
            outbound internet access. Default is False (network isolated with
            private loopback only).
        skills_dir: Optional host path containing staged workspace skills.

    Returns:
        nsjail protobuf configuration as a string.

    Raises:
        AgentSandboxValidationError: If any input contains dangerous characters.
    """
    # Import lazily so sandboxed runtime imports of this module do not require
    # the full tracecat.sandbox package to be mounted inside the jail.
    from tracecat.sandbox.networking import (
        pasta_dns_mount_config_lines,
        pasta_user_net_config_lines,
        write_pasta_network_files,
    )
    from tracecat.sandbox.seccomp import build_untrusted_seccomp_policy

    # Validate inputs to prevent injection into protobuf config
    _validate_path(rootfs, "rootfs")
    _validate_path(job_dir, "job_dir")
    _validate_path(socket_dir, "socket_dir")
    _validate_path(site_packages_dir, "site_packages_dir")
    if llm_socket_path is not None:
        _validate_path(llm_socket_path, "llm_socket_path")
    if mcp_socket_path is not None:
        _validate_path(mcp_socket_path, "mcp_socket_path")
    if skills_dir is not None:
        _validate_path(skills_dir, "skills_dir")

    # Derive control socket path from socket_dir using well-known name when enabled.
    resolved_control_socket_path: Path | None
    if mount_control_socket:
        resolved_control_socket_path = control_socket_path or (
            socket_dir / CONTROL_SOCKET_NAME
        )
    else:
        resolved_control_socket_path = None
    if resolved_control_socket_path is not None:
        _validate_path(resolved_control_socket_path, "control_socket_path")
    if session_home_dir is not None:
        _validate_path(session_home_dir, "session_home_dir")
    if session_project_dir is not None:
        _validate_path(session_project_dir, "session_project_dir")
    claude_sdk_package_dir = site_packages_dir / "claude_agent_sdk"
    _validate_path(claude_sdk_package_dir, "claude_sdk_package_dir")
    # JAILED_LLM_SOCKET_PATH is a constant, no validation needed.

    # Network behavior:
    # - always isolate network namespace for private loopback
    # - internet enabled: add pasta userspace networking for outbound access
    # - internet disabled: no route out of the isolated namespace
    lines = [
        'name: "agent_sandbox"',
        "mode: ONCE",
        'hostname: "agent"',
        "keep_env: false",
        "",
        "# Namespace isolation",
        "clone_newnet: true",
        "clone_newuser: true",
        "clone_newns: true",
        "clone_newpid: true",
        "clone_newipc: true",
        "clone_newuts: true",
        "",
        "# UID/GID mapping - map container user to sandbox user",
        f'uidmap {{ inside_id: "1000" outside_id: "{os.getuid()}" count: 1 }}',
        f'gidmap {{ inside_id: "1000" outside_id: "{os.getgid()}" count: 1 }}',
        "",
        "# Syscall filtering",
        f'seccomp_string: "{build_untrusted_seccomp_policy()}"',
        "",
        "# Rootfs mounts - read-only base system (same rootfs as action sandbox)",
        f'mount {{ src: "{rootfs}/usr" dst: "/usr" is_bind: true rw: false }}',
        f'mount {{ src: "{rootfs}/lib" dst: "/lib" is_bind: true rw: false }}',
        f'mount {{ src: "{rootfs}/bin" dst: "/bin" is_bind: true rw: false }}',
        f'mount {{ src: "{rootfs}/etc" dst: "/etc" is_bind: true rw: false }}',
    ]

    if enable_internet_access:
        lines.extend(pasta_user_net_config_lines())

    # Optional mounts - only include if the directories exist in rootfs
    lib64_path = rootfs / "lib64"
    if lib64_path.exists():
        lines.append(
            f'mount {{ src: "{lib64_path}" dst: "/lib64" is_bind: true rw: false }}'
        )

    sbin_path = rootfs / "sbin"
    if sbin_path.exists():
        lines.append(
            f'mount {{ src: "{sbin_path}" dst: "/sbin" is_bind: true rw: false }}'
        )

    if enable_internet_access:
        network_files = write_pasta_network_files(socket_dir)
        lines.extend(pasta_dns_mount_config_lines(network_files))

    # Fresh procfs avoids leaking executor-container process metadata. Docker
    # runtimes must run these containers with systempaths=unconfined; otherwise
    # masked proc entries like /proc/kcore prevent nested procfs mounts.
    lines.extend(
        [
            "",
            "# /proc - fresh read-only procfs scoped to the jail PID namespace",
            'mount { dst: "/proc" fstype: "proc" rw: false }',
            "",
            "# /dev essentials",
            'mount { src: "/dev/null" dst: "/dev/null" is_bind: true rw: true }',
            'mount { src: "/dev/urandom" dst: "/dev/urandom" is_bind: true rw: false }',
            'mount { src: "/dev/random" dst: "/dev/random" is_bind: true rw: false }',
            'mount { src: "/dev/zero" dst: "/dev/zero" is_bind: true rw: false }',
            "",
            "# Temporary filesystems",
            'mount { dst: "/tmp" fstype: "tmpfs" rw: true options: "size=256M" }',
            "",
            "# Job directory - contains copied runtime code",
            f'mount {{ src: "{job_dir}" dst: "/work" is_bind: true rw: true }}',
        ]
    )
    lines.extend(
        [
            "",
            "# Standalone shim: mount only the Claude SDK package tree",
            "# so the jailed Claude binary can execute without exposing the",
            "# entire host site-packages directory.",
            'mount { dst: "/site-packages" fstype: "tmpfs" rw: false options: "size=1M" }',
            f'mount {{ src: "{claude_sdk_package_dir}" dst: "/site-packages/claude_agent_sdk" is_bind: true rw: false }}',
        ]
    )

    if resolved_control_socket_path is not None:
        lines.extend(
            [
                "",
                "# Per-job control socket",
                f'mount {{ src: "{resolved_control_socket_path}" dst: "{JAILED_CONTROL_SOCKET_PATH}" is_bind: true rw: false }}',
            ]
        )

    if llm_socket_path is not None:
        lines.extend(
            [
                "",
                "# Per-job LLM socket (proxied to LLM gateway on host)",
                f'mount {{ src: "{llm_socket_path}" dst: "{JAILED_LLM_SOCKET_PATH}" is_bind: true rw: false }}',
            ]
        )

    if mcp_socket_path is not None:
        lines.extend(
            [
                "",
                "# Trusted MCP socket (proxied to the agent runtime over localhost)",
                f'mount {{ src: "{mcp_socket_path}" dst: "{TRUSTED_MCP_SOCKET_PATH}" is_bind: true rw: false }}',
            ]
        )

    if session_home_dir is not None and session_project_dir is not None:
        lines.extend(
            [
                "",
                "# Stable Claude session directories shared across jailed turns",
                f'mount {{ src: "{session_home_dir}" dst: "/work/claude-home" is_bind: true rw: true }}',
                f'mount {{ src: "{session_project_dir}" dst: "/work/claude-project" is_bind: true rw: true }}',
            ]
        )

    if skills_dir is not None:
        if session_home_dir is not None:
            lines.extend(
                [
                    "",
                    "# Workspace skills for the Claude Code Skill tool",
                    f'mount {{ src: "{skills_dir}" dst: "/work/claude-home/.claude/skills" is_bind: true rw: false }}',
                ]
            )
        else:
            lines.extend(
                [
                    "",
                    "# Workspace skills for the Claude Code Skill tool",
                    'mount { dst: "/home/agent/.claude" fstype: "tmpfs" rw: true options: "size=64M" }',
                    f'mount {{ src: "{skills_dir}" dst: "/home/agent/.claude/skills" is_bind: true rw: false }}',
                ]
            )

    # Resource limits
    lines.extend(
        [
            "",
            "# Resource limits",
            f"rlimit_as: {config.resources.memory_mb * 1024 * 1024}",
            f"rlimit_cpu: {config.resources.cpu_seconds}",
            f"rlimit_fsize: {config.resources.max_file_size_mb * 1024 * 1024}",
            f"rlimit_nofile: {config.resources.max_open_files}",
            f"rlimit_nproc: {config.resources.max_processes}",
            f"time_limit: {config.resources.timeout_seconds}",
        ]
    )

    # Execution settings.
    lines.extend(
        [
            "",
            'cwd: "/work"',
            "# Execution - standalone shim script",
            f'exec_bin {{ path: "/usr/local/bin/python3" arg: "{JAILED_SHIM_ENTRYPOINT_PATH}" }}',
        ]
    )

    return "\n".join(lines)


def build_agent_env_map(config: AgentSandboxConfig) -> dict[str, str]:
    """Construct a sanitized environment for the nsjail agent process.

    Args:
        config: Agent sandbox configuration.

    Returns:
        Dictionary of environment variables to pass to nsjail.

    Raises:
        AgentSandboxValidationError: If any env var key or value is invalid.
    """
    env_map: dict[str, str] = {**AGENT_SANDBOX_BASE_ENV}
    if value := os.environ.get("TRACECAT__LITELLM_BASE_URL"):
        env_map["TRACECAT__LITELLM_BASE_URL"] = value

    for key, value in config.env_vars.items():
        _validate_env_key(key)
        _validate_env_value(key, value)
        if key in AGENT_SANDBOX_BASE_ENV:
            raise AgentSandboxValidationError(
                f"Cannot override protected env var: {key}"
            )
        env_map[key] = value

    return env_map
