"""Modal Sandbox Environment.
Runs all agent tools (bash, read, write, edit, glob, grep) inside a Modal
Sandbox with GPU access. The agent calls `wafer tool eval kernelbench --target local`
since the GPU is attached to the sandbox itself.
Architecture:
- Mirrors RemoteTargetEnvironment but replaces SSH with Modal's async Sandbox API
- Uses Modal's native file ops (sandbox.open) for read/write instead of base64
- All Modal SDK calls use .aio() async variants, bridged via trio_asyncio
Usage:
    from wafer_core.environments.modal_sandbox import (
        ModalSandboxEnvironment,
        ModalSandboxConfig,
    )
    env = ModalSandboxEnvironment(
        config=ModalSandboxConfig(gpu="A100", timeout_seconds=600),
        working_dir=Path("./sample_0"),
    )
    await env.setup()
    # Agent tools now execute inside the Modal sandbox
"""
from __future__ import annotations
import logging
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
    import modal
from wafer_core.rollouts.dtypes import (
    AgentState,
    Message,
    RunConfig,
    Tool,
    ToolCall,
    ToolResult,
)
from wafer_core.tools import (
    BASH_TOOL,
    EDIT_TOOL,
    GLOB_TOOL,
    GREP_TOOL,
    READ_TOOL,
    WRITE_TOOL,
)
logger = logging.getLogger("wafer.environments.modal_sandbox")

# The agent loop calls serialize() + deserialize() after each tool call for checkpointing.
# Without this cache, each cycle makes a Modal API call to re-fetch the sandbox handle.
_sandbox_cache: dict[str, modal.Sandbox] = {}
async def _aio(coro):  # noqa: ANN001, ANN202
    """Bridge an asyncio coroutine into trio via trio_asyncio.
    Modal's .aio() methods return asyncio coroutines. Inside a trio event loop
    (even with trio_asyncio.open_loop()), you can't directly await them — trio
    rejects the asyncio yield. This helper uses aio_as_trio to convert.
    """
    import trio_asyncio
    return await trio_asyncio.aio_as_trio(coro)
class SandboxConnectionError(Exception):
    """Sandbox is unreachable. Harness should reprovision or fail."""
# ── Configuration ────────────────────────────────────────────────────────────
@dataclass
class ModalSandboxConfig:
    """Configuration for Modal sandbox creation."""
    gpu: str = "A100"
    timeout_seconds: int = 600
    app_name: str = "kernelbench-eval"
    # Image config
    cuda_version: str = "12.4.0"
    python_version: str = "3.12"
# ── Tool definitions ─────────────────────────────────────────────────────────
ALL_TOOLS = {
    "read": READ_TOOL,
    "write": WRITE_TOOL,
    "edit": EDIT_TOOL,
    "glob": GLOB_TOOL,
    "grep": GREP_TOOL,
    "bash": BASH_TOOL,
}
# ── Environment ──────────────────────────────────────────────────────────────
@dataclass
class ModalSandboxEnvironment:
    """Environment that runs all tools inside a Modal Sandbox with GPU.
    The agent sees standard tools (bash, read, write, etc.) but every operation
    executes inside the sandbox. For kernelbench, the agent runs:
        wafer tool eval kernelbench --target local --backend cuda ...
    since the GPU is attached directly to the sandbox.
    """
    config: ModalSandboxConfig
    working_dir: Path
    remote_working_dir: str = "/workspace"
    enabled_tools: list[str] | None = None
    bash_timeout: int = 300
    # Internal state
    _sandbox: modal.Sandbox | None = field(init=False, repr=False, default=None)
    _sandbox_id: str | None = field(init=False, repr=False, default=None)
    _setup_done: bool = field(init=False, repr=False, default=False)
    def __post_init__(self) -> None:
        if self.enabled_tools is not None:
            unknown = set(self.enabled_tools) - set(ALL_TOOLS)
            if unknown:
                raise ValueError(f"Unknown tools: {sorted(unknown)}")
    # ── Lifecycle ────────────────────────────────────────────────────────────
    async def setup(self) -> None:
        """Idempotent setup. Creates sandbox and syncs working_dir."""
        if self._setup_done:
            return
        import modal
        logger.info(
            "Building Modal image and creating sandbox",
            extra={
                "gpu": self.config.gpu,
                "timeout": self.config.timeout_seconds,
            },
        )
        setup_start = time.time()
        image = self._build_image()
        app = await _aio(
            modal.App.lookup.aio(
                self.config.app_name,
                create_if_missing=True,
            )
        )
        self._sandbox = await _aio(
            modal.Sandbox.create.aio(
                image=image,
                gpu=self.config.gpu,
                timeout=self.config.timeout_seconds,
                workdir=self.remote_working_dir,
                app=app,
            )
        )
        self._sandbox_id = self._sandbox.object_id
        _sandbox_cache[self._sandbox_id] = self._sandbox
        logger.info(
            "Sandbox created",
            extra={
                "sandbox_id": self._sandbox_id,
                "duration_ms": int((time.time() - setup_start) * 1000),
            },
        )
        # Ensure working dir exists in sandbox
        proc = await _aio(self._sandbox.exec.aio("mkdir", "-p", self.remote_working_dir))
        await _aio(proc.wait.aio())
        # Sync local working_dir contents to sandbox
        await self._sync_working_dir()
        self._setup_done = True
    async def teardown(self) -> None:
        """Terminate sandbox and clean up."""
        if self._sandbox is not None:
            logger.info("Terminating sandbox", extra={"sandbox_id": self._sandbox_id})
            if self._sandbox_id and self._sandbox_id in _sandbox_cache:
                del _sandbox_cache[self._sandbox_id]
            await _aio(self._sandbox.terminate.aio())
            self._sandbox = None
            self._sandbox_id = None
            self._setup_done = False
    async def __aenter__(self) -> ModalSandboxEnvironment:
        return self
    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:  # noqa: ANN401
        await self.teardown()
    # ── Image Building ───────────────────────────────────────────────────────
    def _build_image(self) -> modal.Image:
        """Build Modal image with CUDA, PyTorch, and wafer CLI."""
        import modal
        cuda_ver = self.config.cuda_version
        python_ver = self.config.python_version
        image = (
            modal.Image
            .from_registry(
                f"nvidia/cuda:{cuda_ver}-devel-ubuntu22.04",
                add_python=python_ver,
            )
            .apt_install("git", "build-essential", "cmake")
            .pip_install(
                "torch",
                index_url="https://download.pytorch.org/whl/cu124",
                extra_index_url="https://pypi.org/simple",
            )
            .pip_install(
                "numpy",
                "scipy",
                "tqdm",
                "packaging",
                "pytest",
                "ninja",
                "einops",
                "triton",
                "wafer-cli",
            )
            .env({
                "CUDA_HOME": "/usr/local/cuda",
            })
        )
        return image
    # ── File Sync ────────────────────────────────────────────────────────────
    async def _sync_working_dir(self) -> None:
        """Upload local working_dir contents to sandbox.
        Skips sync if working_dir has more than 100 files (likely a repo root,
        not a sample directory). The agent creates files inside the sandbox.
        """
        assert self._sandbox is not None
        if not self.working_dir.exists():
            return
        # Count files first — don't upload entire repos
        file_count = sum(1 for _ in self.working_dir.rglob("*") if _.is_file())
        if file_count > 100:
            logger.info(
                "Skipping working_dir sync (too many files)",
                extra={"file_count": file_count, "working_dir": str(self.working_dir)},
            )
            return
        sync_start = time.time()
        count = 0
        for local_path in self.working_dir.rglob("*"):
            if not local_path.is_file():
                continue
            rel = local_path.relative_to(self.working_dir)
            remote_path = f"{self.remote_working_dir}/{rel}"
            # Ensure parent dir exists
            remote_parent = str(Path(remote_path).parent)
            proc = await _aio(self._sandbox.exec.aio("mkdir", "-p", remote_parent))
            await _aio(proc.wait.aio())
            # Upload via sandbox.open
            content = local_path.read_bytes()
            f = await _aio(self._sandbox.open.aio(remote_path, "wb"))
            await _aio(f.write.aio(content))
            await _aio(f.close.aio())
            count += 1
        logger.info(
            "Synced working_dir to sandbox",
            extra={
                "file_count": count,
                "duration_ms": int((time.time() - sync_start) * 1000),
            },
        )
    # ── Sandbox Exec Helpers ─────────────────────────────────────────────────
    async def _ensure_sandbox(self) -> modal.Sandbox:
        """Return sandbox, raising SandboxConnectionError if dead."""
        if self._sandbox is None:
            raise SandboxConnectionError("No sandbox available")
        return self._sandbox
    async def _exec_in_sandbox(
        self,
        *args: str,
        timeout: int | None = None,
    ) -> tuple[int, str, str]:
        """Execute command in sandbox. Returns (exit_code, stdout, stderr)."""
        import modal.exception
        sandbox = await self._ensure_sandbox()
        exec_start = time.time()
        try:
            proc = await _aio(
                sandbox.exec.aio(
                    *args,
                    timeout=timeout or self.bash_timeout,
                    workdir=self.remote_working_dir,
                )
            )
            exit_code = await _aio(proc.wait.aio())
            stdout = await _aio(proc.stdout.read.aio())
            stderr = await _aio(proc.stderr.read.aio())
        except modal.exception.NotFoundError as e:
            if self._sandbox_id and self._sandbox_id in _sandbox_cache:
                del _sandbox_cache[self._sandbox_id]
            self._sandbox = None
            self._sandbox_id = None
            self._setup_done = False
            raise SandboxConnectionError(f"Sandbox died: {e}") from e
        duration_ms = int((time.time() - exec_start) * 1000)
        logger.debug(
            "sandbox_exec",
            extra={
                "command": " ".join(args)[:200],
                "exit_code": exit_code,
                "duration_ms": duration_ms,
                "stdout_len": len(stdout),
                "stderr_len": len(stderr),
            },
        )
        return exit_code, stdout, stderr
    # ── Environment Protocol ─────────────────────────────────────────────────
    def get_tools(self) -> list[Tool]:
        if self.enabled_tools is None:
            return list(ALL_TOOLS.values())
        return [ALL_TOOLS[name] for name in self.enabled_tools if name in ALL_TOOLS]
    def requires_confirmation(self, tool_call: ToolCall) -> bool:
        return False
    def get_tool_formatter(self, tool_name: str) -> None:
        return None
    def get_status_info(self) -> dict[str, str] | None:
        if self._sandbox_id:
            return {"sandbox": self._sandbox_id[:12]}
        return None
    def get_system_prompt(self) -> str | None:
        return None
    async def on_session_start(self, session_id: str) -> None:
        await self.setup()
    async def on_assistant_message(self, message: Message, state: AgentState) -> AgentState:
        return state
    async def exec_tool(
        self,
        tool_call: ToolCall,
        current_state: AgentState,
        run_config: RunConfig,
        cancel_scope: Any = None,  # noqa: ANN401
    ) -> ToolResult:
        """Execute tool call inside the Modal sandbox."""
        if self.enabled_tools is not None and tool_call.name not in self.enabled_tools:
            return ToolResult(
                tool_call_id=tool_call.id,
                is_error=True,
                content="",
                error=f"Tool '{tool_call.name}' is not enabled.",
            )
        try:
            if tool_call.name == "bash":
                return await self._exec_bash(tool_call)
            elif tool_call.name == "read":
                return await self._exec_read(tool_call)
            elif tool_call.name == "write":
                return await self._exec_write(tool_call)
            elif tool_call.name == "edit":
                return await self._exec_edit(tool_call)
            elif tool_call.name == "glob":
                return await self._exec_glob(tool_call)
            elif tool_call.name == "grep":
                return await self._exec_grep(tool_call)
            else:
                return ToolResult(
                    tool_call_id=tool_call.id,
                    is_error=True,
                    content="",
                    error=f"Unknown tool: {tool_call.name}",
                )
        except SandboxConnectionError:
            raise
        except Exception as e:
            return ToolResult(
                tool_call_id=tool_call.id,
                is_error=True,
                content="",
                error=str(e),
            )
    # ── Tool Implementations ─────────────────────────────────────────────────
    async def _exec_bash(self, tool_call: ToolCall) -> ToolResult:
        """Execute bash command in sandbox."""
        args = tool_call.args
        command = args.get("command", "")
        timeout_ms = args.get("timeout", self.bash_timeout * 1000)
        timeout_s = timeout_ms // 1000
        exit_code, stdout, stderr = await self._exec_in_sandbox(
            "bash",
            "-c",
            command,
            timeout=timeout_s,
        )
        output = stdout
        if stderr:
            output = f"{stdout}\n{stderr}" if stdout else stderr
        if exit_code != 0:
            return ToolResult(
                tool_call_id=tool_call.id,
                is_error=True,
                content=output,
                error=f"Exit code {exit_code}",
            )
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=False,
            content=output,
        )
    async def _exec_read(self, tool_call: ToolCall) -> ToolResult:
        """Read file from sandbox using sandbox.open."""
        args = tool_call.args
        file_path = args.get("path", "")
        offset = args.get("offset", 0)
        limit = args.get("limit")
        remote_path = self._to_remote_path(file_path)
        try:
            f = await _aio(self._sandbox.open.aio(remote_path, "r"))
            content = await _aio(f.read.aio())
            await _aio(f.close.aio())
        except Exception as e:
            return ToolResult(
                tool_call_id=tool_call.id,
                is_error=True,
                content="",
                error=f"Failed to read {file_path}: {e}",
            )
        # Format with line numbers (cat -n style)
        lines = content.splitlines()
        if offset:
            lines = lines[offset:]
        if limit:
            lines = lines[:limit]
        numbered = "\n".join(f"{i + offset + 1:>6}\t{line}" for i, line in enumerate(lines))
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=False,
            content=numbered,
        )
    async def _exec_write(self, tool_call: ToolCall) -> ToolResult:
        """Write file to sandbox using sandbox.open."""
        args = tool_call.args
        file_path = args.get("path", "")
        content = args.get("content", "")
        remote_path = self._to_remote_path(file_path)
        # Ensure parent directory exists
        remote_parent = str(Path(remote_path).parent)
        proc = await _aio(self._sandbox.exec.aio("mkdir", "-p", remote_parent))
        await _aio(proc.wait.aio())
        f = await _aio(self._sandbox.open.aio(remote_path, "w"))
        await _aio(f.write.aio(content))
        await _aio(f.close.aio())
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=False,
            content=f"Wrote {len(content)} bytes to {file_path}",
        )
    async def _exec_edit(self, tool_call: ToolCall) -> ToolResult:
        """Edit file in sandbox: read, find/replace, write back."""
        args = tool_call.args
        file_path = args.get("path", "")
        old_string = args.get("old_string", "")
        new_string = args.get("new_string", "")
        replace_all = args.get("replace_all", False)
        remote_path = self._to_remote_path(file_path)
        # Read current content
        try:
            f = await _aio(self._sandbox.open.aio(remote_path, "r"))
            content = await _aio(f.read.aio())
            await _aio(f.close.aio())
        except Exception:
            return ToolResult(
                tool_call_id=tool_call.id,
                is_error=True,
                content="",
                error=f"File not found: {file_path}",
            )
        if old_string not in content:
            return ToolResult(
                tool_call_id=tool_call.id,
                is_error=True,
                content="",
                error=f"old_string not found in {file_path}",
            )
        if not replace_all and content.count(old_string) > 1:
            return ToolResult(
                tool_call_id=tool_call.id,
                is_error=True,
                content="",
                error=(
                    f"old_string appears {content.count(old_string)} times. "
                    "Use replace_all=true or provide more context."
                ),
            )
        if replace_all:
            new_content = content.replace(old_string, new_string)
        else:
            new_content = content.replace(old_string, new_string, 1)
        f = await _aio(self._sandbox.open.aio(remote_path, "w"))
        await _aio(f.write.aio(new_content))
        await _aio(f.close.aio())
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=False,
            content=f"Edited {file_path}",
        )
    async def _exec_glob(self, tool_call: ToolCall) -> ToolResult:
        """Execute glob via Python's glob module on sandbox."""
        args = tool_call.args
        pattern = args.get("pattern", "")
        path = args.get("path", self.remote_working_dir)

        cmd = (
            f'python3 -c "'
            f"import glob, os; "
            f"os.chdir('{path}'); "
            f"files = sorted(glob.glob('{pattern}', recursive=True)); "
            f"print('\\n'.join(f for f in files if os.path.isfile(f)))\""
        )
        exit_code, stdout, stderr = await self._exec_in_sandbox("bash", "-c", cmd)
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=exit_code != 0,
            content=stdout.strip() if stdout.strip() else "No files found",
            error=stderr if exit_code != 0 else None,
        )
    async def _exec_grep(self, tool_call: ToolCall) -> ToolResult:
        """Execute grep on sandbox."""
        args = tool_call.args
        pattern = args.get("pattern", "")
        path = args.get("path", ".")
        output_mode = args.get("output_mode", "files_with_matches")
        if output_mode == "files_with_matches":
            cmd = f"grep -rl '{pattern}' {path} 2>/dev/null | head -50"
        elif output_mode == "count":
            cmd = f"grep -rc '{pattern}' {path} 2>/dev/null | grep -v ':0$' | head -50"
        else:  # content
            context = args.get("-C", 0)
            if context:
                cmd = f"grep -rn -C {context} '{pattern}' {path} 2>/dev/null | head -200"
            else:
                cmd = f"grep -rn '{pattern}' {path} 2>/dev/null | head -200"
        exit_code, stdout, stderr = await self._exec_in_sandbox("bash", "-c", cmd)
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=False,
            content=stdout if stdout else "No matches found",
        )
    # ── Path Helpers ─────────────────────────────────────────────────────────
    def _to_remote_path(self, file_path: str) -> str:
        """Convert agent-provided path to sandbox path."""
        p = Path(file_path).expanduser()
        if p.is_absolute():
            # Try to make relative to working_dir
            try:
                rel = p.resolve().relative_to(self.working_dir.resolve())
                return f"{self.remote_working_dir}/{rel}"
            except ValueError:
                return str(p)
        else:
            return f"{self.remote_working_dir}/{file_path}"
    # ── Serialization ────────────────────────────────────────────────────────
    async def serialize(self) -> dict:
        from dataclasses import asdict
        return {
            "env_kind": "modal_sandbox",
            "version": "1.0",
            "config": asdict(self.config),
            "working_dir": str(self.working_dir),
            "remote_working_dir": self.remote_working_dir,
            "enabled_tools": self.enabled_tools,
            "bash_timeout": self.bash_timeout,
            "sandbox_id": self._sandbox_id,
            "setup_done": self._setup_done,
        }
    @staticmethod
    async def deserialize(data: dict) -> ModalSandboxEnvironment:
        assert data.get("env_kind") == "modal_sandbox"
        assert data.get("version", "").startswith("1.")
        import modal
        env = ModalSandboxEnvironment(
            config=ModalSandboxConfig(**data.get("config", {})),
            working_dir=Path(data["working_dir"]),
            remote_working_dir=data.get("remote_working_dir", "/workspace"),
            enabled_tools=data.get("enabled_tools"),
            bash_timeout=data.get("bash_timeout", 300),
        )
        # Reconnect to existing sandbox if available
        sandbox_id = data.get("sandbox_id")
        if sandbox_id:

            # serialize/deserialize cycle in the agent loop)
            cached = _sandbox_cache.get(sandbox_id)
            if cached is not None:
                env._sandbox = cached
                env._sandbox_id = sandbox_id
                env._setup_done = data.get("setup_done", False)
            else:
                try:
                    env._sandbox = await _aio(modal.Sandbox.from_id.aio(sandbox_id))
                    env._sandbox_id = sandbox_id
                    env._setup_done = data.get("setup_done", False)
                    _sandbox_cache[sandbox_id] = env._sandbox
                    logger.info("Reconnected to sandbox", extra={"sandbox_id": sandbox_id})
                except Exception as e:
                    logger.warning(
                        "Could not reconnect to sandbox",
                        extra={"sandbox_id": sandbox_id, "error": str(e)},
                    )
        return env
