"""Sandboxed shell execution for Fluxibly.

Uses Anthropic's sandbox-runtime (SRT) for OS-level isolation when available.
Falls back to direct subprocess execution with a warning if SRT is not installed.

uv is used for fast Python environment creation inside the sandbox.

SRT: https://github.com/anthropic-experimental/sandbox-runtime
  - macOS: uses sandbox-exec (Seatbelt)
  - Linux: uses bubblewrap (bwrap)
  - Install: npm install -g @anthropic-ai/sandbox-runtime
"""

from __future__ import annotations

import asyncio
import json
import os
import shutil
from collections.abc import Callable
from pathlib import Path
from typing import Any

from pydantic import BaseModel, Field

from fluxibly.logging import Logger

logger = Logger(component="sandbox")

_srt_available: bool | None = None
_uv_available: bool | None = None


# ══════════════════════════════════════════════════════════════════
# Availability Checks
# ══════════════════════════════════════════════════════════════════


def is_srt_available() -> bool:
    """Check if the ``srt`` CLI is installed and accessible."""
    global _srt_available
    if _srt_available is None:
        _srt_available = shutil.which("srt") is not None
    return _srt_available


def is_uv_available() -> bool:
    """Check if the ``uv`` CLI is installed and accessible."""
    global _uv_available
    if _uv_available is None:
        _uv_available = shutil.which("uv") is not None
    return _uv_available


# ══════════════════════════════════════════════════════════════════
# Configuration
# ══════════════════════════════════════════════════════════════════


class SandboxConfig(BaseModel):
    """Configuration for the sandboxed execution environment."""

    # Filesystem
    allow_write: list[str] = Field(
        default_factory=list,
        description="Extra writable paths beyond sandbox_dir and /tmp.",
    )
    deny_read: list[str] = Field(
        default_factory=lambda: [
            "~/.ssh",
            "~/.aws",
            "~/.config/gcloud",
            "~/.gnupg",
            "**/.env",
            "**/.env.*",
        ],
        description="Paths to deny reads from.",
    )

    # Network
    allowed_domains: list[str] = Field(
        default_factory=lambda: [
            "pypi.org",
            "files.pythonhosted.org",
            "github.com",
            "raw.githubusercontent.com",
        ],
        description="Domains allowed for network access.",
    )
    denied_domains: list[str] = Field(
        default_factory=list,
        description="Domains explicitly denied.",
    )

    # Execution
    timeout_seconds: int = Field(
        default=300,
        description="Maximum seconds for a single command.",
    )

    # uv environment
    use_uv: bool = Field(
        default=True,
        description="Whether to create a uv virtual environment.",
    )
    python_version: str | None = Field(
        default=None,
        description='Python version for uv venv (e.g. "3.11"). None = system default.',
    )
    packages: list[str] = Field(
        default_factory=list,
        description="Packages to pre-install in the sandbox venv.",
    )


# ══════════════════════════════════════════════════════════════════
# Result
# ══════════════════════════════════════════════════════════════════


class SandboxResult(BaseModel):
    """Result from a sandboxed command execution."""

    returncode: int
    stdout: str = ""
    stderr: str = ""
    timed_out: bool = False
    sandboxed: bool = True


# ══════════════════════════════════════════════════════════════════
# SRT Settings
# ══════════════════════════════════════════════════════════════════


def generate_srt_settings(
    sandbox_dir: str,
    config: SandboxConfig | None = None,
) -> dict[str, Any]:
    """Generate SRT settings dict for a session.

    Args:
        sandbox_dir: The session sandbox directory (always writable).
        config: Sandbox configuration overrides.

    Returns:
        Dict suitable for writing as .srt-settings.json.
    """
    config = config or SandboxConfig()

    # Build writable paths: sandbox_dir + /tmp + config
    allow_write = [
        sandbox_dir,
        "/tmp",
        *config.allow_write,
    ]

    # Expand ~ in deny-read paths
    deny_read = [os.path.expanduser(p) for p in config.deny_read]

    settings: dict[str, Any] = {
        "filesystem": {
            "allowWrite": allow_write,
            "denyWrite": [],
            "denyRead": deny_read,
        },
        "network": {
            "allowedDomains": config.allowed_domains,
            "deniedDomains": config.denied_domains,
        },
    }

    return settings


# ══════════════════════════════════════════════════════════════════
# SandboxSession — lazy, reusable per session
# ══════════════════════════════════════════════════════════════════


class SandboxSession:
    """Lazily-initialized sandbox environment for a Runner session.

    The sandbox is **not** created until the first ``run()`` or
    ``install()`` call.  Subsequent calls reuse the same SRT settings
    and uv virtual environment.

    Lifecycle::

        session = SandboxSession(sandbox_dir)
        # Nothing created yet

        await session.run("echo hello")
        # First call → creates .srt-settings.json, uv venv, installs packages

        await session.run("python script.py")
        # Reuses existing sandbox

        # Runner.finally → shutil.rmtree(sandbox_dir) cleans everything
    """

    def __init__(
        self,
        sandbox_dir: str,
        config: SandboxConfig | None = None,
    ) -> None:
        self.sandbox_dir = sandbox_dir
        self.config = config or SandboxConfig()
        self._ready = False
        self._settings_path: str | None = None
        self._venv_path: str | None = None

    async def ensure_ready(self) -> None:
        """Set up the sandbox on first call; no-op on subsequent calls."""
        if self._ready:
            return

        # Ensure sandbox directory exists
        Path(self.sandbox_dir).mkdir(parents=True, exist_ok=True)

        # 1. Write SRT settings
        settings = generate_srt_settings(self.sandbox_dir, self.config)
        self._settings_path = str(
            Path(self.sandbox_dir) / ".srt-settings.json"
        )
        Path(self._settings_path).write_text(json.dumps(settings, indent=2))
        logger.debug(
            "Wrote SRT settings to {path}",
            path=self._settings_path,
        )

        # 2. Create uv venv (if available and configured)
        if self.config.use_uv and is_uv_available():
            self._venv_path = str(Path(self.sandbox_dir) / ".venv")
            await self._create_venv()

        # Mark ready before pre-installing packages so that
        # install() → ensure_ready() is a no-op (avoids infinite recursion).
        self._ready = True

        # 3. Pre-install packages (after _ready=True to avoid recursion)
        if self.config.use_uv and is_uv_available() and self.config.packages:
            await self.install(self.config.packages)
        logger.info(
            "Sandbox ready: dir={dir}, srt={srt}, uv_venv={venv}",
            dir=self.sandbox_dir,
            srt=is_srt_available(),
            venv=self._venv_path is not None,
        )

    async def _create_venv(self) -> None:
        """Create a uv virtual environment in the sandbox."""
        cmd = ["uv", "venv", self._venv_path]  # type: ignore[list-item]
        if self.config.python_version:
            cmd.extend(["--python", self.config.python_version])

        # Use sandbox-local uv cache to avoid permission issues
        env = os.environ.copy()
        env["UV_CACHE_DIR"] = str(Path(self.sandbox_dir) / ".uv-cache")

        proc = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=env,
        )
        stdout, stderr = await proc.communicate()

        if proc.returncode == 0:
            logger.info("Created uv venv at {path}", path=self._venv_path)
        else:
            logger.warning(
                "Failed to create uv venv (rc={rc}): {err}",
                rc=proc.returncode,
                err=stderr.decode(errors="replace"),
            )
            self._venv_path = None

    def _build_env(
        self, extra_env: dict[str, str] | None = None
    ) -> dict[str, str]:
        """Build environment dict with venv activation."""
        env = os.environ.copy()
        if extra_env:
            env.update(extra_env)

        # Redirect uv cache into sandbox so SRT doesn't block writes to ~/.cache/uv
        env["UV_CACHE_DIR"] = str(Path(self.sandbox_dir) / ".uv-cache")

        # Activate venv by prepending its bin/ to PATH
        if self._venv_path and Path(self._venv_path).exists():
            venv_bin = str(Path(self._venv_path) / "bin")
            env["VIRTUAL_ENV"] = str(self._venv_path)
            env["PATH"] = f"{venv_bin}:{env.get('PATH', '')}"

        return env

    async def run(
        self,
        command: str,
        *,
        env: dict[str, str] | None = None,
        cwd: str | None = None,
    ) -> SandboxResult:
        """Execute a shell command in the sandbox.

        On first call, initializes the sandbox (SRT settings, uv venv).
        If SRT is not installed, falls back to direct subprocess with a warning.

        Args:
            command: Shell command string to execute.
            env: Additional environment variables.
            cwd: Working directory (defaults to sandbox_dir).

        Returns:
            SandboxResult with stdout, stderr, returncode, and metadata.
        """
        await self.ensure_ready()

        run_cwd = cwd or self.sandbox_dir
        run_env = self._build_env(env)
        sandboxed = False

        if is_srt_available():
            full_cmd = [
                "srt",
                "--settings",
                self._settings_path,
                "-c",
                command,  # type: ignore[list-item]
            ]
            sandboxed = True
        else:
            logger.warning(
                "SRT not installed — running unsandboxed: {cmd}",
                cmd=command[:100],
            )
            full_cmd = ["bash", "-c", command]

        try:
            proc = await asyncio.create_subprocess_exec(
                *full_cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=run_env,
                cwd=run_cwd,
            )

            try:
                stdout_bytes, stderr_bytes = await asyncio.wait_for(
                    proc.communicate(),
                    timeout=self.config.timeout_seconds,
                )
                return SandboxResult(
                    returncode=proc.returncode or 0,
                    stdout=stdout_bytes.decode(errors="replace"),
                    stderr=stderr_bytes.decode(errors="replace"),
                    sandboxed=sandboxed,
                )
            except asyncio.TimeoutError:
                proc.kill()
                await proc.communicate()
                return SandboxResult(
                    returncode=-1,
                    stderr=f"Command timed out after {self.config.timeout_seconds}s",
                    timed_out=True,
                    sandboxed=sandboxed,
                )

        except FileNotFoundError as e:
            return SandboxResult(
                returncode=-1,
                stderr=f"Command not found: {e}",
                sandboxed=False,
            )

    # ── Package Management ────────────────────────────────────────

    async def install_requirements(
        self, requirements_path: str
    ) -> SandboxResult:
        """Install packages from a requirements.txt file.

        Uses ``uv pip install -r`` for speed.  Falls back to
        ``pip install -r`` if uv is not available.

        Runs directly (not through SRT) — same as ``install()``.

        Args:
            requirements_path: Path to the requirements.txt file.

        Returns:
            SandboxResult from the install command.
        """
        await self.ensure_ready()

        if is_uv_available():
            cmd = f"uv pip install -r {requirements_path}"
        else:
            cmd = f"pip install -r {requirements_path}"

        logger.info(
            "Installing from requirements: {path}",
            path=requirements_path,
        )

        run_env = self._build_env()
        try:
            proc = await asyncio.create_subprocess_exec(
                "bash",
                "-c",
                cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=run_env,
                cwd=self.sandbox_dir,
            )
            stdout_bytes, stderr_bytes = await asyncio.wait_for(
                proc.communicate(),
                timeout=self.config.timeout_seconds,
            )
            result = SandboxResult(
                returncode=proc.returncode or 0,
                stdout=stdout_bytes.decode(errors="replace"),
                stderr=stderr_bytes.decode(errors="replace"),
                sandboxed=False,
            )
            if result.returncode != 0:
                logger.warning(
                    "Requirements install failed (rc={rc}): {err}",
                    rc=result.returncode,
                    err=result.stderr[:500],
                )
            return result
        except asyncio.TimeoutError:
            proc.kill()
            await proc.communicate()
            return SandboxResult(
                returncode=-1,
                stderr=f"Requirements install timed out after {self.config.timeout_seconds}s",
                timed_out=True,
                sandboxed=False,
            )

    async def install(self, packages: list[str]) -> SandboxResult:
        """Install Python packages into the sandbox venv.

        Uses ``uv pip install`` for speed.  If uv is not available,
        falls back to ``pip install``.

        Runs **directly** (not through SRT) because package installation
        is a trusted setup operation that needs full system access (network,
        cache directories, macOS system APIs used by uv).

        Args:
            packages: List of package specifiers (e.g. ["pandas>=2.0", "matplotlib"]).

        Returns:
            SandboxResult from the install command.
        """
        await self.ensure_ready()

        pkg_str = " ".join(packages)
        if is_uv_available():
            cmd = f"uv pip install {pkg_str}"
        else:
            cmd = f"pip install {pkg_str}"

        logger.info("Installing packages: {pkgs}", pkgs=pkg_str)

        # Run directly (no SRT) — install is a trusted setup operation
        run_env = self._build_env()
        try:
            proc = await asyncio.create_subprocess_exec(
                "bash",
                "-c",
                cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=run_env,
                cwd=self.sandbox_dir,
            )
            stdout_bytes, stderr_bytes = await asyncio.wait_for(
                proc.communicate(),
                timeout=self.config.timeout_seconds,
            )
            return SandboxResult(
                returncode=proc.returncode or 0,
                stdout=stdout_bytes.decode(errors="replace"),
                stderr=stderr_bytes.decode(errors="replace"),
                sandboxed=False,
            )
        except asyncio.TimeoutError:
            proc.kill()
            await proc.communicate()
            return SandboxResult(
                returncode=-1,
                stderr=f"Install timed out after {self.config.timeout_seconds}s",
                timed_out=True,
                sandboxed=False,
            )


# ══════════════════════════════════════════════════════════════════
# Shell Tool Factory
# ══════════════════════════════════════════════════════════════════


def create_shell_tool(
    session: SandboxSession,
) -> tuple[dict[str, Any], Callable[..., Any]]:
    """Create a sandboxed shell execution tool for ToolService.

    Returns a ``(tool_schema, async_handler)`` tuple ready for
    ``ToolService.register_function_from_schema()``.

    Args:
        session: SandboxSession to use for execution.

    Returns:
        Tuple of (OpenAI-format tool dict, async handler callable).
    """

    async def sandboxed_shell(command: str) -> str:
        result = await session.run(command=command)
        output = result.stdout
        if result.stderr:
            output += f"\n[stderr]\n{result.stderr}"
        if result.timed_out:
            output += "\n[TIMED OUT]"
        if not result.sandboxed:
            output = "[WARNING: Ran without sandbox]\n" + output
        return output

    schema: dict[str, Any] = {
        "type": "function",
        "function": {
            "name": "shell",
            "description": (
                "Execute a shell command in a sandboxed environment. "
                "Use this to run scripts, install packages, or process files."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "The shell command to execute.",
                    },
                },
                "required": ["command"],
            },
        },
    }

    return schema, sandboxed_shell
