# adapters/sprites.py

from __future__ import annotations

import asyncio
import base64
import shlex
from typing import Any

from ..core.types import ExecResult, ExecOpts, SecureConfig
from ..core.shell import shell_escape, fire_and_forget


def sprites(sprite: Any) -> _SpritesAdapter:
    return _SpritesAdapter(sprite)


def sprites_defaults() -> SecureConfig:
    """Return Sprites-optimized defaults for SecureConfig.

    Usage::

        secured = await secure_sandbox(
            sprites(sprite),
            sprites_defaults(),  # or override: sprites_defaults().model_copy(update={...})
        )
    """
    return SecureConfig(
        install_strategy="preinstalled",
        skip_integrity_check=True,
        real_paths=True,
        workspace="/home/sprite",
    )


class _SpritesAdapter:
    """Wraps a ``sprites-py`` SDK ``Sprite`` instance.

    The sprites-py SDK is sync-first with its own internal event loop,
    so all calls are dispatched via ``run_in_executor`` to avoid
    event loop conflicts with the caller's async loop.
    """

    def __init__(self, sprite: Any) -> None:
        self._sprite = sprite

    def _sh_sync(self, cmd: str, **opts: Any) -> tuple[str, str, int]:
        """Run a shell command synchronously. Returns (stdout, stderr, exit_code).

        Manually sets capture flags on the Cmd object and uses the SDK's
        internal ``run_sync`` to execute, ensuring both stdout and stderr
        are captured regardless of exit code.

        Note: The sprites-py SDK may route program output to stderr depending
        on how the underlying binary writes to file descriptors. To ensure
        consistent behavior, callers that need stdout should use ``2>&1``
        redirection in their commands.
        """
        from sprites.loop import run_sync

        c = self._sprite.command("sh", "-c", cmd, **opts)
        c._capture_stdout = True
        c._capture_stderr = True
        c._started = True
        try:
            exit_code: int = run_sync(c._run_async(), timeout=c.timeout)
        except Exception as e:
            # ExitError: has .stdout, .stderr (bytes), .exit_code() (int)
            stdout = getattr(e, "stdout", b"") or b""
            stderr = getattr(e, "stderr", b"") or b""
            code_fn = getattr(e, "exit_code", None)
            exit_code = code_fn() if callable(code_fn) else (code_fn or 1)
            return (
                stdout.decode("utf-8", errors="replace") if isinstance(stdout, bytes) else str(stdout),
                stderr.decode("utf-8", errors="replace") if isinstance(stderr, bytes) else str(stderr),
                exit_code,
            )
        finally:
            c._finished = True

        return (
            c._stdout_data.decode("utf-8", errors="replace"),
            c._stderr_data.decode("utf-8", errors="replace"),
            exit_code,
        )

    async def _sh(self, cmd: str, **opts: Any) -> tuple[str, str, int]:
        """Run a shell command async via executor."""
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(None, lambda: self._sh_sync(cmd, **opts))

    async def exec(
        self,
        cmd: str,
        args: list[str] | None = None,
        opts: ExecOpts | None = None,
    ) -> ExecResult:
        opts = opts or ExecOpts()
        command = shell_escape(cmd, args)
        full_cmd = f"sudo {command}" if opts.sudo else command

        if opts.detached:
            kw: dict[str, Any] = {}
            if opts.cwd:
                kw["cwd"] = opts.cwd

            async def _detached() -> None:
                await self._sh(f"nohup {full_cmd} > /dev/null 2>&1 &", **kw)

            fire_and_forget(_detached())
            return ExecResult(stdout="", stderr="", exit_code=0)

        kw = {}
        if opts.cwd:
            kw["cwd"] = opts.cwd
        stdout, stderr, exit_code = await self._sh(full_cmd, **kw)
        # The sprites-py SDK may route program output to stderr (fd 2)
        # instead of stdout. When stdout is empty but stderr has content,
        # treat stderr as stdout to match other adapter behaviors.
        if not stdout and stderr and exit_code == 0:
            stdout = stderr
            stderr = ""
        return ExecResult(stdout=stdout, stderr=stderr, exit_code=exit_code)

    # Maximum base64 payload per WebSocket command (~48KB raw → ~64KB base64).
    _CHUNK_SIZE = 48 * 1024

    async def write_file(self, path: str, content: str | bytes, *, sudo: bool = False) -> None:
        data = content if isinstance(content, bytes) else content.encode()
        safe_path = path.replace("'", "'\\''")

        if len(data) <= self._CHUNK_SIZE:
            b64 = base64.b64encode(data).decode("ascii")
            inner = f"printf '%s' '{b64}' | base64 -d > '{safe_path}'"
            cmd = f"sudo sh -c {shlex.quote(inner)}" if sudo else inner
            stdout, stderr, exit_code = await self._sh(cmd)
            if exit_code != 0:
                raise RuntimeError(f"writeFile failed (exit {exit_code}): {stderr}")
        else:
            # Chunked write for large files (e.g. agentsh binary upload)
            offset = 0
            first = True
            while offset < len(data):
                chunk = data[offset : offset + self._CHUNK_SIZE]
                b64 = base64.b64encode(chunk).decode("ascii")
                op = ">" if first else ">>"
                inner = f"printf '%s' '{b64}' | base64 -d {op} '{safe_path}'"
                cmd = f"sudo sh -c {shlex.quote(inner)}" if sudo else inner
                stdout, stderr, exit_code = await self._sh(cmd)
                if exit_code != 0:
                    raise RuntimeError(f"writeFile failed (exit {exit_code}): {stderr}")
                first = False
                offset += self._CHUNK_SIZE

    async def read_file(self, path: str) -> str:
        safe_path = path.replace("'", "'\\''")
        stdout, stderr, exit_code = await self._sh(f"cat '{safe_path}'")
        if exit_code != 0:
            raise RuntimeError(f"readFile failed (exit {exit_code}): {stderr}")
        # Apply same stderr→stdout normalization as exec() — sprites-py may
        # route program output to stderr even on success.
        if not stdout and stderr:
            stdout = stderr
        return stdout

    async def stop(self) -> None:
        self._sprite.delete()

    async def file_exists(self, path: str) -> bool:
        result = await self.exec("test", ["-f", path])
        return result.exit_code == 0
