"""Subprocess sandbox for secure command execution."""

from __future__ import annotations

import asyncio
import contextlib
import os
import resource
import signal
from dataclasses import dataclass, field

import structlog

log = structlog.get_logger(__name__)


@dataclass
class SandboxConfig:
    """Configuration for subprocess sandboxing."""

    max_memory_mb: int = 512
    max_cpu_time: int = 300  # seconds
    max_file_size_mb: int = 100
    max_processes: int = 50
    max_open_files: int = 256
    nice_level: int = 10  # Lower priority
    env_allowlist: list[str] = field(
        default_factory=lambda: ["PATH", "HOME", "USER", "LANG", "LC_ALL"]
    )
    env_denylist: list[str] = field(
        default_factory=lambda: [
            "AWS_",
            "AZURE_",
            "GCP_",
            "API_KEY",
            "SECRET",
            "TOKEN",
            "PASSWORD",
            "CREDENTIAL",
        ]
    )


def _set_resource_limits(config: SandboxConfig) -> None:
    """Set resource limits for subprocess (Unix only)."""
    try:
        # Memory limit (in bytes)
        mem_bytes = config.max_memory_mb * 1024 * 1024
        resource.setrlimit(resource.RLIMIT_AS, (mem_bytes, mem_bytes))

        # CPU time limit
        resource.setrlimit(resource.RLIMIT_CPU, (config.max_cpu_time, config.max_cpu_time))

        # File size limit
        file_bytes = config.max_file_size_mb * 1024 * 1024
        resource.setrlimit(resource.RLIMIT_FSIZE, (file_bytes, file_bytes))

        # Process limit
        resource.setrlimit(resource.RLIMIT_NPROC, (config.max_processes, config.max_processes))

        # Open files limit
        resource.setrlimit(resource.RLIMIT_NOFILE, (config.max_open_files, config.max_open_files))

        # Set nice level (lower priority)
        os.nice(config.nice_level)

    except (OSError, ValueError) as e:
        log.warning("resource_limit_failed", error=str(e))


def _filter_environment(
    base_env: dict[str, str] | None,
    extra_env: dict[str, str] | None,
    config: SandboxConfig,
) -> dict[str, str]:
    """
    Filter environment variables for subprocess.

    Args:
        base_env: Base environment (or None for current)
        extra_env: Additional environment variables
        config: Sandbox configuration

    Returns:
        Filtered environment dictionary
    """
    # Start with current environment if not specified
    source_env = base_env if base_env is not None else os.environ.copy()

    # Filter based on allowlist
    filtered: dict[str, str] = {}
    for key in config.env_allowlist:
        if key in source_env:
            filtered[key] = source_env[key]

    # Remove any that match denylist patterns
    for key in list(filtered.keys()):
        for pattern in config.env_denylist:
            if pattern.endswith("_"):
                # Prefix match (e.g., "AWS_")
                if key.startswith(pattern):
                    del filtered[key]
                    break
            elif pattern.upper() in key.upper():
                # Contains match
                del filtered[key]
                break

    # Add extra environment variables (also filtered)
    if extra_env:
        for key, value in extra_env.items():
            safe = True
            for pattern in config.env_denylist:
                if pattern.endswith("_"):
                    if key.startswith(pattern):
                        safe = False
                        break
                elif pattern.upper() in key.upper():
                    safe = False
                    break
            if safe:
                filtered[key] = value

    return filtered


@dataclass
class SandboxedProcess:
    """
    Wrapper for a sandboxed subprocess.

    Provides resource limiting and process management.
    """

    command: str
    args: list[str]
    config: SandboxConfig = field(default_factory=SandboxConfig)
    working_dir: str | None = None
    env: dict[str, str] | None = None
    _process: asyncio.subprocess.Process | None = field(default=None, init=False)

    async def start(self) -> asyncio.subprocess.Process:
        """
        Start the sandboxed subprocess.

        Returns:
            The subprocess
        """
        # Build command
        cmd = [self.command, *self.args]

        # Filter environment
        filtered_env = _filter_environment(None, self.env, self.config)

        # Create subprocess
        self._process = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=self.working_dir,
            env=filtered_env,
            preexec_fn=lambda: _set_resource_limits(self.config),
            start_new_session=True,  # Isolate in new process group
        )

        log.debug(
            "sandboxed_process_started",
            pid=self._process.pid,
            command=self.command,
        )

        return self._process

    async def wait(self, timeout: float | None = None) -> tuple[int, bytes, bytes]:
        """
        Wait for process completion.

        Args:
            timeout: Maximum time to wait

        Returns:
            Tuple of (exit_code, stdout, stderr)

        Raises:
            asyncio.TimeoutError: If timeout exceeded
        """
        if self._process is None:
            raise RuntimeError("Process not started")

        try:
            stdout, stderr = await asyncio.wait_for(
                self._process.communicate(),
                timeout=timeout,
            )
            return self._process.returncode or 0, stdout, stderr

        except TimeoutError:
            await self.kill()
            # Collect any partial output
            stdout = b""
            stderr = b""
            if self._process.stdout:
                with contextlib.suppress(TimeoutError):
                    stdout = await asyncio.wait_for(
                        self._process.stdout.read(),
                        timeout=1.0,
                    )
            if self._process.stderr:
                with contextlib.suppress(TimeoutError):
                    stderr = await asyncio.wait_for(
                        self._process.stderr.read(),
                        timeout=1.0,
                    )
            raise

    async def kill(self) -> None:
        """Kill the subprocess and its process group."""
        if self._process is None:
            return

        try:
            # Kill the entire process group
            if self._process.pid:
                with contextlib.suppress(ProcessLookupError, PermissionError):
                    os.killpg(os.getpgid(self._process.pid), signal.SIGKILL)

            # Also try to kill the process directly
            self._process.kill()
            await self._process.wait()

            log.debug("sandboxed_process_killed", pid=self._process.pid)

        except ProcessLookupError:
            pass

    @property
    def pid(self) -> int | None:
        """Get process ID."""
        return self._process.pid if self._process else None

    @property
    def returncode(self) -> int | None:
        """Get process return code."""
        return self._process.returncode if self._process else None


async def run_sandboxed(
    command: str,
    args: list[str] | None = None,
    timeout: float = 300.0,
    working_dir: str | None = None,
    env: dict[str, str] | None = None,
    config: SandboxConfig | None = None,
) -> tuple[int, str, str]:
    """
    Run a command in a sandbox and return results.

    Args:
        command: Command to run
        args: Command arguments
        timeout: Execution timeout
        working_dir: Working directory
        env: Environment variables
        config: Sandbox configuration

    Returns:
        Tuple of (exit_code, stdout, stderr)
    """
    sandbox = SandboxedProcess(
        command=command,
        args=args or [],
        config=config or SandboxConfig(),
        working_dir=working_dir,
        env=env,
    )

    await sandbox.start()
    exit_code, stdout, stderr = await sandbox.wait(timeout=timeout)

    return (
        exit_code,
        stdout.decode("utf-8", errors="replace"),
        stderr.decode("utf-8", errors="replace"),
    )
