import base64
import asyncio
import io
import json
import uuid
import zipfile
import httpx
from pathlib import Path
from typing import List, Literal, Callable, TYPE_CHECKING
from datetime import timedelta
from loguru import logger
from pydantic import BaseModel, Base64Bytes

from e2b import (
    AsyncSandbox,
    CommandExitException,
    default_build_logger,
    BuildInfo,
    Template,
)

from common.core.utils import TransientStorageMixin
from common.config import common_settings
from common.filestore import get_filestore
from common.utils import attach_protocol_to_url

if TYPE_CHECKING:
    from common.filestore import FileStore


class DataBlob(BaseModel):
    data: Base64Bytes | None = None  # Stored as base64, auto-encodes/decodes for JSON
    description: str | None = None
    url: str | None = None
    name: str | None = None

    def write(self, directory: str) -> str:
        if self.data is None:
            raise ValueError("Data not available.")
        filename = self.name if self.name else f"{uuid.uuid4().hex}"
        path = Path(directory) / filename
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "wb") as f:
            f.write(self.data)
        return str(path)

    @classmethod
    async def from_url(cls, url: str, **kwargs) -> "DataBlob":
        """Create a DataBlob from a URL.

        Handles both regular http(s) URLs and tiptree:// filestore URLs.
        """
        # Check if this is a filestore URL
        protocol_prefix = f"{common_settings.PROTOCOL_NAME}://"
        if url.startswith(protocol_prefix):
            key = url[len(protocol_prefix) :]
            # Extract name from key if not provided
            if "name" not in kwargs:
                kwargs["name"] = Path(key).name
            return await cls.from_filestore(key=key, **kwargs)

        # Regular HTTP(S) URL - download directly

        async with httpx.AsyncClient(follow_redirects=True) as client:
            response = await client.get(url)
            response.raise_for_status()
        data = base64.b64encode(response.content)
        name = kwargs.get("name") or Path(url).name or None
        return cls(
            data=data,
            url=url,
            name=name,
            **{k: v for k, v in kwargs.items() if k != "name"},
        )

    @classmethod
    def from_file(cls, path: str, **kwargs) -> "DataBlob":
        with open(path, "rb") as f:
            data = base64.b64encode(f.read())
        name = kwargs.get("name") or Path(path).name
        return cls(
            data=data, name=name, **{k: v for k, v in kwargs.items() if k != "name"}
        )

    async def upload_to_filestore(
        self,
        key: str,
        filestore: "FileStore | None" = None,
    ) -> str:
        """Upload this DataBlob to a filestore and return the protocol URL.

        Args:
            key: The filestore key (path) where the data will be stored.
            filestore: Optional FileStore instance. If not provided, uses the default filestore.

        Returns:
            The protocol URL (e.g., "tiptree://path/to/file") for the uploaded file.

        Raises:
            ValueError: If the DataBlob has no data to upload.
        """

        if self.data is None:
            raise ValueError("DataBlob has no data to upload")

        if filestore is None:
            filestore = get_filestore()

        await filestore.put_object(key=key, data=self.data)

        return attach_protocol_to_url(key, common_settings.PROTOCOL_NAME)

    @classmethod
    async def from_filestore(
        cls, key: str, filestore: "FileStore | None" = None, **kwargs
    ) -> "DataBlob":
        """Create a DataBlob by downloading from a filestore.

        Args:
            key: The filestore key (path) of the file to download.
            filestore: Optional FileStore instance. If not provided, uses the default filestore.
            **kwargs: Additional arguments passed to the DataBlob constructor.

        Returns:
            A DataBlob containing the downloaded file data.
        """

        if filestore is None:
            filestore = get_filestore()

        # Get a signed URL for the file
        signed_url = await filestore.get_signed_url(
            key=key, expiry=timedelta(minutes=15)
        )

        # Download the file
        async with httpx.AsyncClient(follow_redirects=True) as client:
            response = await client.get(signed_url)
            response.raise_for_status()

        data = base64.b64encode(response.content)
        name = kwargs.get("name") or Path(key).name
        protocol_url = attach_protocol_to_url(key, common_settings.PROTOCOL_NAME)

        return cls(
            data=data,
            url=protocol_url,
            name=name,
            **{k: v for k, v in kwargs.items() if k != "name"},
        )


class CodeAgentOutput(BaseModel):
    output: str
    output_files: List[DataBlob] | None = None

    def write_files(self, directory: str) -> List[str]:
        if not self.output_files:
            return []
        paths = []
        for blob in self.output_files:
            path = blob.write(directory)
            paths.append(path)
        return paths


class CodeAgentSandbox(BaseModel, TransientStorageMixin):
    sandbox_template: Literal["claude-code-python"] = "claude-code-python"
    sandbox_timeout: int = 600

    sandbox_id: str | None = None
    sandbox_state: DataBlob | None = None

    async def get_or_create_sandbox(
        self, restore_state: bool = True, env_vars: dict | None = None
    ) -> AsyncSandbox:
        """Get or create the sandbox instance.

        Args:
            restore_state: If True (default), restore sandbox_state when creating a new sandbox.
                          Set to False to create a fresh sandbox without restoring previous state.
            env_vars: Optional dictionary of environment variables to set in the sandbox.
        """
        storage = self.get_or_create_transient_storage()
        if "sandbox" not in storage.keys():
            await self.create_or_connect_sandbox(
                restore_state=restore_state, env_vars=env_vars
            )
        return storage.get("sandbox")

    async def create_or_connect_sandbox(
        self, restore_state: bool = True, env_vars: dict | None = None
    ) -> "AsyncSandbox":
        storage = self.get_or_create_transient_storage()
        # Check if we have a sandbox_id from a previous session (e.g., paused sandbox)
        if self.sandbox_id is not None:
            # Connect to existing sandbox (automatically resumes if paused)
            sandbox = await AsyncSandbox.connect(
                sandbox_id=self.sandbox_id,
                timeout=self.sandbox_timeout,
            )
        else:
            # Create a new sandbox
            sandbox = await AsyncSandbox.create(
                template=self.sandbox_template,
                timeout=self.sandbox_timeout,
                envs=env_vars,
            )
            # Restore state if we have any and restore_state is True
            if restore_state and self.sandbox_state is not None:
                await self._infiltrate_state(sandbox, self.sandbox_state)
            self.sandbox_id = sandbox.sandbox_id
        storage.set("sandbox", sandbox)
        return sandbox

    async def pause_sandbox(self) -> None:
        """Pause the sandbox, preserving its state (filesystem and memory).

        The sandbox can be resumed later by calling get_sandbox(), which will
        automatically reconnect and resume the paused sandbox using the stored sandbox_id.

        Note: Pausing takes about 4 seconds per 1 GiB of RAM. The sandbox can be
        kept paused for up to 30 days before the data may be deleted.
        """
        storage = self.get_or_create_transient_storage()
        if "sandbox" in storage:
            sandbox = storage.get("sandbox")
            await sandbox.beta_pause()
            # Remove from transient storage but keep sandbox_id for later reconnection
            storage.delete("sandbox")

    async def kill_sandbox(self, save_state: bool = True) -> None:
        """Kill the sandbox and optionally save its state.

        Args:
            save_state: If True (default), exfiltrate and save sandbox state before killing.
                       Set to False to kill without saving (discards any changes).
        """
        storage = self.get_or_create_transient_storage()
        if "sandbox" in storage:
            sandbox = storage.get("sandbox")
            if save_state:
                self.sandbox_state = await self._exfiltrate_state(sandbox)
            await sandbox.kill()
            storage.delete("sandbox")
            self.sandbox_id = None

    def build_sandbox_image(self) -> BuildInfo:
        assets_dir = Path(__file__).parent / "assets"
        dockerfile_content = (
            assets_dir / "templates" / self.sandbox_template / "Dockerfile"
        ).read_text()

        template = Template().from_dockerfile(dockerfile_content)

        return Template.build(
            template=template,
            alias=self.sandbox_template,
            cpu_count=2,
            memory_mb=1024,
            on_build_logs=default_build_logger(),
        )

    @classmethod
    async def _exfiltrate_path(cls, sandbox: AsyncSandbox, path: str) -> DataBlob:
        # The idea is to pull down files from the sandbox, at path.
        exfil_uuid = str(uuid.uuid4().hex)
        sandbox_path = f"/tmp/{exfil_uuid}.zip"
        name = Path(path).name

        # First, verify the source path exists
        check_result = await sandbox.commands.run(
            cmd=f"test -e {path} && echo 'exists' || echo 'missing'",
            timeout=10,
        )
        if "missing" in check_result.stdout:
            raise FileNotFoundError(f"Path does not exist in sandbox: {path}")

        # Zip the content (file or directory)
        zip_result = await sandbox.commands.run(
            cmd=f"zip -r {sandbox_path} {path}",
            timeout=60,
        )
        if zip_result.exit_code != 0:
            raise RuntimeError(f"zip failed: {zip_result.stderr}")

        # Download the zip file
        downloaded_file = await sandbox.files.read(path=sandbox_path, format="bytes")
        # Base64 encode for storage in Base64Bytes field
        return DataBlob(data=base64.b64encode(downloaded_file), name=name)

    async def exfiltrate_path(self, path: str) -> DataBlob:
        sandbox = await self.get_or_create_sandbox()
        return await self._exfiltrate_path(sandbox, path)

    async def exfiltrate_state(self) -> DataBlob:
        sandbox = await self.get_or_create_sandbox()
        return await self._exfiltrate_state(sandbox)

    @classmethod
    async def _exfiltrate_state(cls, sandbox: AsyncSandbox) -> DataBlob:
        """
        Exfiltrate the entire home directory of the sandbox, including a manifest of what's installed.

        Generates manifests for:
        - apt packages (dpkg --get-selections)
        - pip packages (pip freeze)
        - npm global packages (npm list -g)

        These are saved to ~/.state-manifests/ before zipping $HOME.
        """
        manifest_dir = "$HOME/.state-manifests"

        # Create manifest directory and generate package lists
        manifest_result = await sandbox.commands.run(
            cmd=(
                f"mkdir -p {manifest_dir} && "
                f"dpkg --get-selections > {manifest_dir}/apt-packages.txt 2>/dev/null || true && "
                f"pip freeze > {manifest_dir}/pip-requirements.txt 2>/dev/null || true && "
                f"npm list -g --depth=0 > {manifest_dir}/npm-global.txt 2>/dev/null || true"
            ),
            timeout=300,
        )
        if manifest_result.exit_code != 0:
            # Non-fatal: we still want to exfiltrate even if manifest generation fails
            logger.warning(
                f"Warning: manifest generation had issues: {manifest_result.stderr}"
            )

        return await cls._exfiltrate_path(sandbox, "$HOME")

    @classmethod
    async def _infiltrate_state(cls, sandbox: AsyncSandbox, state: DataBlob) -> None:
        """
        Restore a previously exfiltrated state to a sandbox.

        1. Uploads and unzips the home directory
        2. Reinstalls packages from manifests (apt, pip, npm)
        """
        infil_uuid = str(uuid.uuid4().hex)
        sandbox_path = f"/tmp/{infil_uuid}.zip"

        # Upload the zip file
        await sandbox.files.write(path=sandbox_path, data=state.data)

        # Unzip to root (the zip contains full paths like /home/user/...)
        unzip_result = await sandbox.commands.run(
            cmd=f"unzip -o {sandbox_path} -d /",
            timeout=60,
        )
        if unzip_result.exit_code != 0:
            raise RuntimeError(f"unzip failed: {unzip_result.stderr}")

        # Restore packages from manifests (best-effort, non-fatal)
        # Order: apt first (system deps), then pip/npm (which may need system libs)
        manifest_dir = "$HOME/.state-manifests"

        # Restore apt packages first (provides system libs for pip/npm native extensions)
        apt_result = await sandbox.commands.run(
            cmd=(
                f"test -f {manifest_dir}/apt-packages.txt && "
                f"sudo apt-get update -qq && "
                f"sudo dpkg --set-selections < {manifest_dir}/apt-packages.txt && "
                f"sudo apt-get dselect-upgrade -y -qq 2>/dev/null || true"
            ),
            timeout=600,
        )
        if apt_result.exit_code != 0:
            logger.warning(f"Warning: apt restore had issues: {apt_result.stderr}")

        # Restore pip packages
        pip_result = await sandbox.commands.run(
            cmd=(
                f"test -f {manifest_dir}/pip-requirements.txt && "
                f"pip install -r {manifest_dir}/pip-requirements.txt 2>/dev/null || true"
            ),
            timeout=300,
        )
        if pip_result.exit_code != 0:
            logger.warning(f"Warning: pip restore had issues: {pip_result.stderr}")

        # Restore npm global packages (parse the list and install)
        npm_result = await sandbox.commands.run(
            cmd=(
                f"test -f {manifest_dir}/npm-global.txt && "
                # Extract package names from npm list output, skip first line (path) and empty lines
                f"tail -n +2 {manifest_dir}/npm-global.txt | "
                f"grep -oP '(?<=── )[^@]+' | "
                f"xargs -r npm install -g 2>/dev/null || true"
            ),
            timeout=300,
        )
        if npm_result.exit_code != 0:
            logger.warning(f"Warning: npm restore had issues: {npm_result.stderr}")

        # Cleanup
        await sandbox.commands.run(cmd=f"rm -f {sandbox_path}", timeout=10)

    async def infiltrate_state(self, state: DataBlob) -> None:
        sandbox = await self.get_or_create_sandbox()
        await self._infiltrate_state(sandbox, state)

    @classmethod
    async def _infiltrate_file(
        cls, sandbox: AsyncSandbox, data: DataBlob, path: str
    ) -> AsyncSandbox:
        if data.data is None:
            raise ValueError("DataBlob has no data to infiltrate")
        # Ensure parent directory exists
        parent_dir = str(Path(path).parent)
        await sandbox.commands.run(cmd=f"mkdir -p {parent_dir}", timeout=10)
        # Write the file
        await sandbox.files.write(path=path, data=data.data)
        return sandbox

    @classmethod
    async def _infiltrate_files(
        cls, sandbox: AsyncSandbox, data: DataBlob | List[DataBlob], path: str
    ) -> AsyncSandbox:
        """Infiltrate one or more files to the sandbox.

        Args:
            sandbox: The sandbox to infiltrate files into.
            data: A single DataBlob or list of DataBlobs to infiltrate.
            path: For single blob, the full file path. For multiple blobs, the directory path.
                  Filenames for multiple blobs are taken from each DataBlob's name attribute.

        Returns:
            The sandbox instance.
        """
        # Handle single blob - delegate to existing method
        if isinstance(data, DataBlob):
            return await cls._infiltrate_file(sandbox, data, path)

        # Handle empty list
        if not data:
            return sandbox

        # Handle multiple blobs - zip locally, upload, unzip at path
        zip_buffer = io.BytesIO()
        with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
            for blob in data:
                if blob.data is None:
                    raise ValueError(
                        f"DataBlob '{blob.name}' has no data to infiltrate"
                    )
                filename = blob.name or f"file_{uuid.uuid4().hex[:8]}"
                zf.writestr(filename, blob.data)

        zip_data = zip_buffer.getvalue()

        # Create unique temp path for the zip on the sandbox
        infil_uuid = uuid.uuid4().hex
        sandbox_zip_path = f"/tmp/infil_{infil_uuid}.zip"

        # Ensure target directory exists
        await sandbox.commands.run(cmd=f"mkdir -p {path}", timeout=10)

        # Upload the zip file
        await sandbox.files.write(path=sandbox_zip_path, data=zip_data)

        # Unzip to target directory
        unzip_result = await sandbox.commands.run(
            cmd=f"unzip -o {sandbox_zip_path} -d {path}",
            timeout=60,
        )
        if unzip_result.exit_code != 0:
            raise RuntimeError(f"unzip failed: {unzip_result.stderr}")

        # Cleanup temp zip
        await sandbox.commands.run(cmd=f"rm -f {sandbox_zip_path}", timeout=10)

        return sandbox

    @classmethod
    async def _infiltrate_directory(
        cls, sandbox: AsyncSandbox, local_path: str | Path, sandbox_path: str
    ) -> AsyncSandbox:
        """Infiltrate a local directory to the sandbox, preserving structure.

        Args:
            sandbox: The sandbox to infiltrate the directory into.
            local_path: Path to the local directory to copy.
            sandbox_path: Destination path in the sandbox.

        Returns:
            The sandbox instance.
        """
        local_path = Path(local_path)
        if not local_path.exists():
            raise FileNotFoundError(f"Local path does not exist: {local_path}")
        if not local_path.is_dir():
            raise ValueError(f"Local path is not a directory: {local_path}")

        # Create an in-memory zip preserving directory structure
        zip_buffer = io.BytesIO()
        with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
            for file_path in local_path.rglob("*"):
                if file_path.is_file():
                    # Preserve relative path from local_path
                    arcname = file_path.relative_to(local_path)
                    zf.write(file_path, arcname)

        zip_data = zip_buffer.getvalue()

        # If zip is empty (no files), skip
        if len(zip_data) <= 22:  # Empty zip is ~22 bytes
            return sandbox

        # Create unique temp path for the zip on the sandbox
        infil_uuid = uuid.uuid4().hex
        sandbox_zip_path = f"/tmp/infil_dir_{infil_uuid}.zip"

        # Ensure target directory exists
        await sandbox.commands.run(cmd=f"mkdir -p {sandbox_path}", timeout=10)

        # Upload the zip file
        await sandbox.files.write(path=sandbox_zip_path, data=zip_data)

        # Unzip to target directory
        unzip_result = await sandbox.commands.run(
            cmd=f"unzip -o {sandbox_zip_path} -d {sandbox_path}",
            timeout=60,
        )
        if unzip_result.exit_code != 0:
            raise RuntimeError(f"unzip failed: {unzip_result.stderr}")

        # Cleanup temp zip
        await sandbox.commands.run(cmd=f"rm -f {sandbox_zip_path}", timeout=10)

        return sandbox

    @staticmethod
    def _unpack_zip_blob(zip_blob: DataBlob) -> List[DataBlob]:
        """Unpack a zip blob into individual file blobs."""
        if zip_blob.data is None:
            return []

        blobs: List[DataBlob] = []
        with zipfile.ZipFile(io.BytesIO(zip_blob.data), "r") as zf:
            for info in zf.infolist():
                # Skip directories
                if info.is_dir():
                    continue
                file_data = zf.read(info.filename)
                # Extract just the filename from the full path
                name = Path(info.filename).name
                blobs.append(DataBlob(data=base64.b64encode(file_data), name=name))
        return blobs

    async def _exfil_new_output_files(
        self, sandbox: AsyncSandbox, force: bool = False
    ) -> List[DataBlob] | None:
        """Exfiltrate only new output files (created after last exfil).

        Uses a timestamp file at /home/user/work/.meta/last_exfil_time to track
        when files were last exfiltrated. Only files newer than this timestamp
        are downloaded.

        Args:
            sandbox: The sandbox instance to exfiltrate from.
            force: If True, exfiltrate all files regardless of timestamp.

        Returns:
            List of DataBlobs for new files, or None if no new files.
        """
        output_dir = "/home/user/work/output_files"
        meta_dir = "/home/user/work/.meta"
        timestamp_file = f"{meta_dir}/last_exfil_time"

        # Check if output directory has any files
        check_output = await sandbox.commands.run(
            cmd=f"ls -A {output_dir} 2>/dev/null | head -1 || true",
            timeout=10,
        )
        if not check_output.stdout.strip():
            return None

        # Find files newer than last exfil (or all files if force or no timestamp exists)
        if force:
            find_cmd = f"find {output_dir} -type f"
        else:
            # Find files newer than timestamp file, or all if timestamp doesn't exist
            find_cmd = (
                f"if [ -f {timestamp_file} ]; then "
                f"find {output_dir} -type f -newer {timestamp_file}; "
                f"else find {output_dir} -type f; fi"
            )

        find_result = await sandbox.commands.run(cmd=find_cmd, timeout=30)
        new_files = [
            f.strip() for f in find_result.stdout.strip().split("\n") if f.strip()
        ]

        if not new_files:
            return None

        # Create a temp directory with just the new files, preserving structure
        exfil_uuid = uuid.uuid4().hex
        temp_dir = f"/tmp/exfil_{exfil_uuid}"
        await sandbox.commands.run(cmd=f"mkdir -p {temp_dir}", timeout=10)

        # Copy new files to temp dir
        for filepath in new_files:
            # Get relative path from output_dir
            rel_path = filepath.replace(f"{output_dir}/", "")
            dest_path = f"{temp_dir}/{rel_path}"
            dest_dir = str(Path(dest_path).parent)
            await sandbox.commands.run(
                cmd=f"mkdir -p {dest_dir} && cp {filepath} {dest_path}",
                timeout=10,
            )

        # Exfiltrate the temp directory
        output_blob = await self._exfiltrate_path(sandbox, temp_dir)
        output_files = self._unpack_zip_blob(output_blob)

        # Update the timestamp file
        await sandbox.commands.run(
            cmd=f"mkdir -p {meta_dir} && touch {timestamp_file}",
            timeout=10,
        )

        # Cleanup temp dir
        await sandbox.commands.run(cmd=f"rm -rf {temp_dir}", timeout=10)

        return output_files if output_files else None

    async def run(
        self,
        task: str,
        input_files: List[DataBlob] | None = None,
        stateful: bool = False,
        pause: bool = False,
        terminate: bool = False,
        env_vars: dict | None = None,
        on_transcript_event: Callable[[dict], None] | None = None,
        force_output_file_exfil: bool = False,
        model_name: str | None = None,
    ) -> CodeAgentOutput:
        # Validate
        if pause:
            # Cannot terminate if pausing
            terminate = False

        # Get the sandbox
        sandbox = await self.get_or_create_sandbox(restore_state=stateful)
        logger.debug(f"Built sandbox with ID {sandbox.sandbox_id}")

        # Infil the sandbox_claude.md file to /home/user/work/CLAUDE.md
        assets_dir = Path(__file__).parent / "assets"
        sandbox_claude_md_path = (
            assets_dir / "templates" / self.sandbox_template / "sandbox_claude.md"
        )
        if sandbox_claude_md_path.exists():
            claude_md = DataBlob.from_file(str(sandbox_claude_md_path))
            await self._infiltrate_file(sandbox, claude_md, "/home/user/work/CLAUDE.md")
            logger.debug("Infiltrated CLAUDE.md into sandbox.")

        # Infil skills from sandbox_skills/ to /home/user/work/.claude/skills/
        sandbox_skills_path = (
            assets_dir / "templates" / self.sandbox_template / "sandbox_skills"
        )
        if sandbox_skills_path.exists() and sandbox_skills_path.is_dir():
            await self._infiltrate_directory(
                sandbox, sandbox_skills_path, "/home/user/work/.claude/skills"
            )
            logger.debug("Infiltrated skills into sandbox.")

        # Infil any input_files to /home/user/work/input_files/...
        if input_files:
            await self._infiltrate_files(
                sandbox, input_files, "/home/user/work/input_files"
            )
            logger.debug(f"Infiltrated {len(input_files)} input files into sandbox.")

        # Run the claude command at /home/user/work
        # Use heredoc to safely pass the task without shell escaping issues
        output_format = "stream-json" if on_transcript_event else "text"
        buffer = ""
        agent_output = None

        def _on_stdout(data):
            nonlocal buffer
            nonlocal agent_output
            if on_transcript_event:
                if isinstance(data, dict):
                    buffer += data.get("line", "") + "\n"
                else:
                    buffer += str(data)
                while "\n" in buffer:
                    line, buffer = buffer.split("\n", 1)
                    line = line.strip()
                    if line:
                        try:
                            line_as_json = json.loads(line)
                            on_transcript_event(line_as_json)
                            # If this line has a final output, capture it
                            if line_as_json.get("type") == "result":
                                agent_output = line_as_json.get("result")
                        except json.JSONDecodeError:
                            pass

        # Build the command piecemeal
        cmd_parts = [
            "cd /home/user/work &&",
            "claude -p",
            "--dangerously-skip-permissions",
            "--continue",
            "--verbose",
            f"--output-format {output_format}",
        ]
        if model_name:
            cmd_parts.append(f"--model {model_name}")
        cmd_parts.append(f"<< 'TASK_EOF'\n{task}\nTASK_EOF")
        cmd = " ".join(cmd_parts)

        try:
            result = await sandbox.commands.run(
                cmd=cmd,
                timeout=0,
                envs=env_vars,
                on_stdout=_on_stdout if on_transcript_event else None,
            )
            logger.debug("Executed command in sandbox.")
        except CommandExitException as e:
            logger.exception(
                f"Command execution failed in sandbox. Details:\n"
                f"stderr: {e.stderr}\n"
                f"stdout: {e.stdout}\n"
                f"exit_code: {e.exit_code}\n"
                f"error: {e.error}"
            )
            raise

        # Set the agent output
        if agent_output is None:
            agent_output = result.stdout

        # Exfil any output files from /home/user/work/output_files/
        output_files = await self._exfil_new_output_files(
            sandbox, force=force_output_file_exfil
        )
        if output_files:
            logger.debug(
                f"Exfiltrated {len(output_files)} new output file(s) from sandbox."
            )
        else:
            logger.debug("No new output files to exfiltrate.")

        # Optionally terminate the sandbox
        if pause:
            await self.pause_sandbox()
            logger.debug("Paused sandbox.")
        if terminate:
            await self.kill_sandbox(save_state=stateful)
            logger.debug("Terminated sandbox.")

        return CodeAgentOutput(output=agent_output, output_files=output_files)
